tinymind-native-8b-remote-handoff/bundle/evaluation/gateway_teacher_distill.py

from __future__ import annotations

from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
from typing import Any

from evaluation.llm_stats_gateway import GatewayFetcher, LLMStatsGatewayClient


def _read_prompts(path: str | Path) -> list[dict[str, str]]:
    rows: list[dict[str, str]] = []
    for i, line in enumerate(Path(path).read_text(encoding="utf-8", errors="replace").splitlines(), start=1):
        if not line.strip():
            continue
        try:
            row = json.loads(line)
        except json.JSONDecodeError:
            row = {"prompt": line}
        prompt = str(row.get("prompt") or row.get("question") or row.get("text") or "").strip()
        if prompt:
            rows.append({"id": str(row.get("id") or f"prompt-{i}"), "prompt": prompt})
    return rows


def _response_text(response: dict[str, Any]) -> str:
    try:
        return str(response["choices"][0]["message"]["content"])
    except (KeyError, IndexError, TypeError):
        return ""


def _messages(prompt: str, *, concise_retry: bool = False) -> list[dict[str, str]]:
    if concise_retry:
        system = (
            "You are a concise expert teacher. Return a direct training answer only. "
            "Do not return null, None, tool calls, or hidden reasoning."
        )
        user = f"Teach this as a clean SFT answer in 250-450 words:\n{prompt}"
    else:
        system = "You are a careful teacher. Answer with grounded, concise, high-signal reasoning. Do not reveal secrets."
        user = prompt
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def _reward(prompt: str, answer: str) -> dict[str, float]:
    prompt_terms = {t.lower() for t in prompt.split() if len(t) > 2}
    answer_terms = {t.lower().strip(".,:;!?()[]{}") for t in answer.split() if len(t) > 2}
    coverage = len(prompt_terms & answer_terms) / max(1, len(prompt_terms))
    length_quality = min(1.0, len(answer) / 400.0) if answer else 0.0
    total = 0.55 * coverage + 0.45 * length_quality
    return {"coverage": round(coverage, 6), "length_quality": round(length_quality, 6), "total": round(total, 6)}


def build_gateway_teacher_distill(
    out_dir: str | Path,
    *,
    prompts_path: str | Path,
    model: str = "glm-5.1",
    min_reward: float = 0.15,
    api_key: str | None = None,
    fetcher: GatewayFetcher | None = None,
) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    client = LLMStatsGatewayClient(api_key=api_key)
    prompts = _read_prompts(prompts_path)
    rows: list[dict[str, Any]] = []
    rejected: list[dict[str, Any]] = []
    for item in prompts:
        response = client.chat_completion(
            model=model,
            messages=_messages(item["prompt"]),
            fetcher=fetcher,
        )
        answer = _response_text(response)
        retried = False
        if not answer.strip() or answer.strip().lower() in {"none", "null", "n/a"}:
            retried = True
            response = client.chat_completion(
                model=model,
                messages=_messages(item["prompt"], concise_retry=True),
                temperature=0.1,
                fetcher=fetcher,
            )
            answer = _response_text(response)
        reward = _reward(item["prompt"], answer)
        reject_reason = ""
        if not answer.strip() or answer.strip().lower() in {"none", "null", "n/a"}:
            reject_reason = "empty_or_none_answer"
        elif reward["total"] < min_reward:
            reject_reason = "low_reward"
        row = {
            "messages": [
                {"role": "system", "content": "Answer using distilled teacher reasoning with evidence discipline."},
                {"role": "user", "content": item["prompt"]},
                {"role": "assistant", "content": answer},
            ],
            "source": "llm_stats_gateway_teacher",
            "prompt_id": item["id"],
            "prompt_sha256": hashlib.sha256(item["prompt"].encode("utf-8")).hexdigest(),
            "answer_sha256": hashlib.sha256(answer.encode("utf-8")).hexdigest(),
            "reward": reward,
            "teacher": {
                "model": model,
                "provider": "llm-stats-gateway",
                "api_key_stored": False,
                "raw_response_stored": False,
                "retried_empty_answer": retried,
            },
        }
        if reject_reason:
            rejected.append({"prompt_id": item["id"], "reason": reject_reason, "reward": reward})
        else:
            rows.append(row)
    sft_path = out / "gateway_teacher_distill_sft.jsonl"
    with sft_path.open("w", encoding="utf-8", newline="\n") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
    avg_reward = sum(row["reward"]["total"] for row in rows) / max(1, len(rows))
    report = {
        "schema_version": "tinymind-gateway-teacher-distill-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "model": model,
        "prompts_path": str(prompts_path),
        "sft_path": str(sft_path),
        "records_written": len(rows),
        "records_rejected": len(rejected),
        "rejected": rejected,
        "min_reward": min_reward,
        "avg_reward": round(avg_reward, 6),
        "api_key_redacted": client.redacted_key(),
        "claim_gate": {
            "teacher_distill_ready": len(rows) > 0,
            "quality_gate_passed": len(rows) > 0 and not any(r["reason"] == "empty_or_none_answer" for r in rejected),
            "hidden_weight_extraction_claim_allowed": False,
            "all_knowledge_extracted_claim_allowed": False,
            "reason": "This stores teacher-generated answers with provenance. It does not extract model weights or complete internal knowledge.",
        },
    }
    path = out / "gateway_teacher_distill_report.json"
    report["json_path"] = str(path)
    path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    return report