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

"""Evidence-first system auto tuner for TinyMind adapters.

The tuner does not modify weights. It selects the best available adapter from
saved train/eval evidence, quarantines regressions, and writes the next safe
training plan. This prevents bad micro-distill runs from becoming the default
model just because they ran later.
"""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import json
import math
from pathlib import Path
from typing import Any


DEFAULT_ADAPTER_ROOT = Path("model/tinymind-12b/adapters")


@dataclass(frozen=True)
class AdapterScore:
    adapter: str
    manifest_path: str
    eval_loss: float | None
    perplexity: float | None
    train_loss: float | None
    eval_records: int
    train_records: int
    max_steps: int
    dataset: str
    status: str
    quarantine_reason: str | None

    @property
    def usable(self) -> bool:
        return self.eval_loss is not None and self.quarantine_reason is None


def _read_json(path: Path) -> dict[str, Any]:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {}


def _float_or_none(value: Any) -> float | None:
    try:
        out = float(value)
    except (TypeError, ValueError):
        return None
    if not math.isfinite(out):
        return None
    return out


def _score_manifest(path: Path) -> AdapterScore:
    data = _read_json(path)
    adapter = str(data.get("adapter") or path.parent)
    eval_loss = _float_or_none(data.get("eval_loss") or (data.get("eval_metrics") or {}).get("eval_loss"))
    perplexity = _float_or_none(data.get("perplexity"))
    train_loss = _float_or_none((data.get("train_metrics") or {}).get("train_loss"))
    eval_records = int(data.get("eval_records") or 0)
    train_records = int(data.get("train_records") or 0)
    max_steps = int(data.get("max_steps") or 0)
    dataset = str(data.get("dataset") or "")
    status = "measured" if eval_loss is not None else "missing_eval"
    quarantine_reason: str | None = None
    eval_credible = data.get("eval_credible")
    if eval_loss is None:
        quarantine_reason = "missing_eval_loss"
    elif eval_credible is False:
        quarantine_reason = "eval_not_credible"
    elif eval_records < 2:
        quarantine_reason = "eval_split_too_small"
    elif "gateway_teacher_distill" in dataset.replace("\\", "/").lower() and train_records < 32:
        quarantine_reason = "micro_teacher_transfer_not_enough_rows"
    return AdapterScore(
        adapter=adapter,
        manifest_path=str(path),
        eval_loss=eval_loss,
        perplexity=perplexity,
        train_loss=train_loss,
        eval_records=eval_records,
        train_records=train_records,
        max_steps=max_steps,
        dataset=dataset,
        status=status,
        quarantine_reason=quarantine_reason,
    )


def discover_adapter_scores(adapter_root: str | Path = DEFAULT_ADAPTER_ROOT) -> list[AdapterScore]:
    root = Path(adapter_root)
    manifests = sorted(root.rglob("tinymind_12b_manifest.json"), key=lambda p: p.stat().st_mtime, reverse=True)
    return [_score_manifest(path) for path in manifests]


def choose_champion(scores: list[AdapterScore]) -> AdapterScore | None:
    usable = [score for score in scores if score.usable]
    if not usable:
        return None
    max_eval_records = max(score.eval_records for score in usable)
    if max_eval_records >= 128:
        usable = [score for score in usable if score.eval_records >= 128]
    return min(usable, key=lambda row: (row.eval_loss if row.eval_loss is not None else float("inf"), -row.eval_records))


def build_system_auto_tuner_report(
    out_dir: str | Path,
    *,
    adapter_root: str | Path = DEFAULT_ADAPTER_ROOT,
    min_teacher_rows: int = 32,
) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    scores = discover_adapter_scores(adapter_root)
    champion = choose_champion(scores)
    rows = [score.__dict__ for score in scores]
    quarantined = [row for row in rows if row["quarantine_reason"]]
    next_plan = {
        "base_adapter": champion.adapter if champion else "",
        "avoid_adapters": [row["adapter"] for row in quarantined],
        "teacher_transfer_min_rows": int(min_teacher_rows),
        "recommended_next_dataset": "reports/dataset_quality_governor/tinymind_12b_quality_governed_mix.jsonl",
        "recommended_eval_records": 128,
        "recommended_max_seq_length": 1024,
        "recommended_steps": 32,
        "rules": [
            "Never promote a checkpoint with eval_records < 2.",
            "Never promote micro teacher-transfer rows as champion without a held-out eval gain.",
            "Prefer lower eval_loss only when the eval split is credible and claim gates remain false.",
        ],
    }
    report = {
        "schema_version": "tinymind-system-auto-tuner-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "adapter_root": str(adapter_root),
        "adapter_count": len(rows),
        "champion": champion.__dict__ if champion else None,
        "adapters": rows,
        "quarantined": quarantined,
        "next_plan": next_plan,
        "claim_gate": {
            "champion_selected": champion is not None,
            "bad_adapter_quarantined": bool(quarantined),
            "auto_promote_allowed": champion is not None,
            "world_best_claim_allowed": False,
            "reason": "Champion selection is based on local eval manifests only, not official external rank.",
        },
    }
    json_path = out / "system_auto_tuner_report.json"
    champion_path = out / "champion_adapter.json"
    md_path = out / "system_auto_tuner_report.md"
    report["json_path"] = str(json_path)
    report["champion_path"] = str(champion_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    champion_path.write_text(
        json.dumps({"champion": report["champion"], "next_plan": next_plan}, ensure_ascii=False, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def promote_candidate_if_better(
    out_dir: str | Path,
    *,
    candidate_manifest: str | Path,
    champion_path: str | Path = "reports/system_auto_tuner/champion_adapter.json",
    baseline_manifest: str | Path | None = None,
    min_eval_records: int = 16,
    min_delta: float = 0.0,
) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    candidate = _score_manifest(Path(candidate_manifest))
    champion_file = Path(champion_path)
    champion_data = _read_json(champion_file).get("champion") if champion_file.exists() else None
    champion_score: AdapterScore | None = None
    if baseline_manifest:
        champion_score = _score_manifest(Path(baseline_manifest))
    elif champion_data:
        champion_manifest = champion_data.get("manifest_path")
        if not champion_manifest and champion_data.get("adapter"):
            champion_manifest = str(Path(champion_data["adapter"]) / "tinymind_12b_manifest.json")
        if champion_manifest:
            champion_score = _score_manifest(Path(champion_manifest))

    reject_reason: str | None = None
    if candidate.quarantine_reason:
        reject_reason = candidate.quarantine_reason
    elif candidate.eval_records < min_eval_records:
        reject_reason = f"candidate_eval_records_below_minimum_{min_eval_records}"
    elif champion_score and champion_score.eval_loss is not None and candidate.eval_loss is not None:
        if candidate.eval_loss >= champion_score.eval_loss - min_delta:
            reject_reason = "candidate_not_better_than_champion"
    elif candidate.eval_loss is None:
        reject_reason = "candidate_missing_eval_loss"

    promoted = reject_reason is None
    report = {
        "schema_version": "tinymind-promotion-gate-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "candidate": candidate.__dict__,
        "previous_champion": champion_score.__dict__ if champion_score else None,
        "baseline_manifest": str(baseline_manifest) if baseline_manifest else None,
        "min_eval_records": int(min_eval_records),
        "min_delta": float(min_delta),
        "promoted": promoted,
        "reject_reason": reject_reason,
        "claim_gate": {
            "promotion_allowed": promoted,
            "world_best_claim_allowed": False,
            "reason": "Promotion is local eval-loss selection only, not official external rank.",
        },
    }
    if promoted:
        next_plan = _read_json(champion_file).get("next_plan", {}) if champion_file.exists() else {}
        champion_file.parent.mkdir(parents=True, exist_ok=True)
        champion_file.write_text(
            json.dumps({"champion": candidate.__dict__, "next_plan": next_plan}, ensure_ascii=False, indent=2, sort_keys=True),
            encoding="utf-8",
        )
        report["champion_path"] = str(champion_file)
    path = out / "promotion_gate_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


def _markdown(report: dict[str, Any]) -> str:
    champion = report.get("champion") or {}
    lines = [
        "# TinyMind System Auto Tuner",
        "",
        f"- Adapter count: {report['adapter_count']}",
        f"- Champion selected: {report['claim_gate']['champion_selected']}",
        f"- Bad adapter quarantined: {report['claim_gate']['bad_adapter_quarantined']}",
        f"- World-best claim allowed: {report['claim_gate']['world_best_claim_allowed']}",
    ]
    if champion:
        lines.extend(
            [
                f"- Champion adapter: {champion['adapter']}",
                f"- Champion eval_loss: {champion['eval_loss']}",
                f"- Champion perplexity: {champion['perplexity']}",
                f"- Champion eval_records: {champion['eval_records']}",
            ]
        )
    lines.append("")
    lines.append("## Quarantined")
    for row in report["quarantined"]:
        lines.append(f"- {row['adapter']}: {row['quarantine_reason']}")
    return "\n".join(lines) + "\n"