tinymind-native-colab-handoff/bundle/evaluation/ten_million_step_readiness.py

"""Evidence gate for real long-step TinyMind training runs."""

from __future__ import annotations

from datetime import datetime, timezone
import json
import math
from pathlib import Path


DEFAULT_TARGET_STEPS = 10_000_000


def _load(path: str | Path | None) -> dict:
    if not path:
        return {}
    p = Path(path)
    if not p.exists():
        return {}
    return json.loads(p.read_text(encoding="utf-8"))


def _estimate_seconds(target_steps: int, calibration_steps: int, calibration_seconds: float) -> dict:
    sec_per_step = calibration_seconds / max(calibration_steps, 1)
    total_seconds = sec_per_step * target_steps
    return {
        "target_steps": int(target_steps),
        "calibration_steps": calibration_steps,
        "calibration_seconds": calibration_seconds,
        "seconds_per_step": sec_per_step,
        "estimated_total_seconds": total_seconds,
        "estimated_total_hours": total_seconds / 3600.0,
        "estimated_total_days": total_seconds / 86400.0,
    }


def build_ten_million_step_readiness(
    out_dir: str | Path,
    baseline_report: str | Path | None = None,
    target_steps: int = DEFAULT_TARGET_STEPS,
    calibration_steps: int = 256,
    calibration_seconds: float = 181.6,
    checkpoint_every: int = 10_000,
    eval_every: int = 5_000,
) -> dict:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    baseline = _load(baseline_report)
    target_steps = max(1, int(target_steps))
    estimate = _estimate_seconds(target_steps, calibration_steps, calibration_seconds)
    latest_train = baseline.get("train_eval", {})
    artifacts = baseline.get("artifacts", {})

    command = (
        "python -m train.cli hf-pure-auto-refine-train "
        f"--preset thai-code --out-dir reports\\hf_pure_thai_code_{target_steps}_steps "
        f"--rows-per-source 100 --train-steps {target_steps}"
    )
    report = {
        "schema_version": "tinymind-long-step-readiness-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "target_steps": target_steps,
        "status": "ready_not_started",
        "baseline_report": str(baseline_report) if baseline_report else "",
        "baseline_train_eval": latest_train,
        "baseline_artifacts": artifacts,
        "estimate": estimate,
        "recommended_runtime": {
            "checkpoint_every_steps": int(checkpoint_every),
            "eval_every_steps": int(eval_every),
            "save_best": True,
            "resume_required": True,
            "power_guard_required": True,
            "nan_inf_guard_required": True,
            "stop_on_loss_divergence": True,
        },
        "launch_command": command,
        "claim_gate": {
            "target_steps_completed": False,
            "ten_million_steps_completed": False,
            "world_best_claim_allowed": False,
            "top1_claim_allowed": False,
            "reason": "Long-run claim remains blocked until a checkpoint/report with step >= target_steps exists and external benchmarks pass.",
        },
    }
    json_path = out / "long_step_readiness_report.json"
    md_path = out / "long_step_readiness_report.md"
    report["json_path"] = str(json_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")
    md_path.write_text(
        "\n".join(
            [
                "# TinyMind Long-Step Readiness",
                "",
                f"- Target steps: {target_steps:,}",
                f"- Status: {report['status']}",
                f"- Calibration: {calibration_steps:,} steps / {calibration_seconds:.2f}s",
                f"- Estimated days: {estimate['estimated_total_days']:.2f}",
                f"- Baseline eval loss: {latest_train.get('eval_loss')}",
                f"- Baseline perplexity: {latest_train.get('perplexity')}",
                f"- Target completed: {report['claim_gate']['target_steps_completed']}",
                f"- World-best claim allowed: {report['claim_gate']['world_best_claim_allowed']}",
                "",
                "Launch command:",
                "",
                f"```powershell\n{command}\n```",
                "",
            ]
        ),
        encoding="utf-8",
    )
    return report