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

"""Official hard benchmark runner for TinyMind.

Runs what is publicly accessible immediately and records blockers for gated or
provider-only suites. Scores are real local measurements, not official ranks.
"""

from __future__ import annotations

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

import torch

from evaluation.local_evidence import _encode
from model.architecture import OmegaModel


LETTERS = "ABCDEFGHIJ"
REFERENCE_MODEL_SIZES = {
    "sshleifer/tiny-gpt2": {
        "params": 102714,
        "notes": "HF tiny GPT-2 smoke baseline; size reference only.",
    },
    "distilgpt2": {
        "params": 81912576,
        "notes": "Distilled GPT-2 baseline; size reference only.",
    },
    "gpt2": {
        "params": 124439808,
        "notes": "GPT-2 small baseline; size reference only.",
    },
}


def load_dataset(*args, **kwargs):
    """Lazy datasets import to keep CLI/test collection free of pyarrow startup cost."""
    from datasets import load_dataset as _load_dataset

    return _load_dataset(*args, **kwargs)


def _load_model(checkpoint_path: str | Path) -> OmegaModel:
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    model = OmegaModel(ckpt["model_cfg"])
    model.load_state_dict(ckpt["model_state"])
    model.eval()
    return model


def _mb(path: str | Path | None) -> float | None:
    if not path:
        return None
    p = Path(path)
    if not p.exists():
        return None
    return p.stat().st_size / (1024 * 1024)


def profile_model_size(
    model: OmegaModel,
    checkpoint_path: str | Path,
    safetensors_path: str | Path | None = None,
    int4_artifact_path: str | Path | None = None,
) -> dict:
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    dense_bf16_mb_estimate = total_params * 2 / (1024 * 1024)
    dense_fp32_mb_estimate = total_params * 4 / (1024 * 1024)
    int4_raw_mb_estimate = total_params * 0.5 / (1024 * 1024)
    return {
        "model": "TinyMind-PureField-ReGenesis",
        "architecture_mode": getattr(model.cfg, "architecture_mode", "unknown"),
        "dim": model.cfg.dim,
        "layers": model.cfg.n_layers,
        "vocab_size": model.cfg.vocab_size,
        "max_seq_len": model.cfg.max_seq_len,
        "total_params": total_params,
        "trainable_params": trainable_params,
        "million_params": total_params / 1_000_000,
        "checkpoint_mb": _mb(checkpoint_path),
        "safetensors_mb": _mb(safetensors_path),
        "int4_artifact_mb": _mb(int4_artifact_path),
        "dense_bf16_mb_estimate": dense_bf16_mb_estimate,
        "dense_fp32_mb_estimate": dense_fp32_mb_estimate,
        "int4_raw_mb_estimate": int4_raw_mb_estimate,
    }


def build_size_comparison(tinymind_size: dict) -> list[dict]:
    tiny_params = max(int(tinymind_size["total_params"]), 1)
    rows = [
        {
            "model": tinymind_size["model"],
            "params": tinymind_size["total_params"],
            "million_params": tinymind_size["million_params"],
            "relative_to_tinymind_params": 1.0,
            "checkpoint_mb": tinymind_size["checkpoint_mb"],
            "safetensors_mb": tinymind_size["safetensors_mb"],
            "int4_artifact_mb": tinymind_size["int4_artifact_mb"],
            "comparison_scope": "measured local artifact",
        }
    ]
    for name, info in REFERENCE_MODEL_SIZES.items():
        params = int(info["params"])
        rows.append(
            {
                "model": name,
                "params": params,
                "million_params": params / 1_000_000,
                "relative_to_tinymind_params": params / tiny_params,
                "checkpoint_mb": None,
                "safetensors_mb": None,
                "int4_artifact_mb": None,
                "comparison_scope": info["notes"],
            }
        )
    return rows


def _mmlu_pro_prompt(row: dict) -> str:
    options = "\n".join(f"{LETTERS[i]}. {opt}" for i, opt in enumerate(row["options"]))
    return (
        "Answer the following MMLU-Pro question. Return only the option letter.\n\n"
        f"Question: {row['question']}\n"
        f"{options}\n"
        "Answer:"
    )


@torch.no_grad()
def _choose_letter(model: OmegaModel, prompt: str) -> tuple[str, dict[str, float]]:
    ids = _encode(prompt, model.cfg.max_seq_len, model.cfg.vocab_size).unsqueeze(0)
    logits = model(ids)["logits"][0, -1].float()
    scores: dict[str, float] = {}
    for letter in LETTERS:
        token_id = 4 + (ord(letter) % max(model.cfg.vocab_size - 4, 1))
        scores[letter] = float(logits[token_id].item())
    return max(scores, key=scores.get), scores


def run_mmlu_pro(
    model: OmegaModel,
    split: str = "validation",
    limit: int = 20,
) -> dict:
    ds = load_dataset("TIGER-Lab/MMLU-Pro", split=split, streaming=True)
    rows = []
    correct = 0
    total = 0
    by_category: dict[str, list[int]] = {}
    for row in ds:
        if total >= limit:
            break
        pred, scores = _choose_letter(model, _mmlu_pro_prompt(row))
        gold = str(row["answer"]).strip().upper()
        ok = pred == gold
        correct += int(ok)
        total += 1
        cat = str(row.get("category", "unknown"))
        by_category.setdefault(cat, [0, 0])
        by_category[cat][0] += int(ok)
        by_category[cat][1] += 1
        rows.append(
            {
                "question_id": row.get("question_id"),
                "category": cat,
                "gold": gold,
                "prediction": pred,
                "correct": ok,
                "scores": scores,
            }
        )
    return {
        "benchmark": "MMLU-Pro",
        "dataset": "TIGER-Lab/MMLU-Pro",
        "split": split,
        "limit": limit,
        "samples": total,
        "accuracy": correct / max(total, 1),
        "correct": correct,
        "by_category": {key: wins / max(count, 1) for key, (wins, count) in by_category.items()},
        "rows": rows,
        "official_rank": None,
        "official_rank_claimed": False,
    }


def _probe_dataset(name: str, path: str, config: str | None = None) -> dict:
    try:
        ds = load_dataset(path, config, split="train", streaming=True) if config else load_dataset(path, split="train", streaming=True)
        first = next(iter(ds))
        return {"name": name, "path": path, "status": "accessible", "keys": list(first.keys())}
    except Exception as exc:
        return {"name": name, "path": path, "status": "blocked", "error_type": type(exc).__name__, "error": str(exc)[:500]}


def run_official_hard_eval(
    checkpoint_path: str | Path,
    out_dir: str | Path,
    mmlu_limit: int = 20,
    safetensors_path: str | Path | None = None,
    int4_artifact_path: str | Path | None = None,
) -> dict:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    model = _load_model(checkpoint_path)
    size = profile_model_size(model, checkpoint_path, safetensors_path, int4_artifact_path)
    mmlu = run_mmlu_pro(model, split="validation", limit=mmlu_limit)
    blockers = [
        _probe_dataset("GPQA Diamond", "Idavidrein/gpqa", "gpqa_diamond"),
        _probe_dataset("Humanity's Last Exam", "cais/hle", None),
        {
            "name": "FrontierMath",
            "status": "requires_official_access",
            "path": "https://epoch.ai/benchmarks/frontiermath/",
        },
        {
            "name": "SWE-bench Pro",
            "status": "requires_linux_docker_harness_or_provider",
            "path": "official SWE-bench Pro harness",
        },
    ]
    report = {
        "schema_version": "tinymind-official-hard-eval-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "checkpoint_path": str(checkpoint_path),
        "size": size,
        "size_comparison": build_size_comparison(size),
        "results": {"mmlu_pro": mmlu},
        "blockers": blockers,
        "world_best_claim_allowed": False,
        "claim_note": "This is a real local run on accessible official/public data. It is not an official leaderboard rank.",
    }
    json_path = out / "official_hard_eval_report.json"
    report["json_path"] = str(json_path)
    md_path = out / "official_hard_eval_report.md"
    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(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict) -> str:
    mmlu = report["results"]["mmlu_pro"]
    size = report["size"]
    lines = [
        "# TinyMind Official Hard Eval",
        "",
        "## Current Model Size",
        "",
        f"- Parameters: {size['total_params']:,} ({size['million_params']:.6f}M)",
        f"- Checkpoint: {_fmt_mb(size['checkpoint_mb'])}",
        f"- Safetensors: {_fmt_mb(size['safetensors_mb'])}",
        f"- INT4 sparse artifact: {_fmt_mb(size['int4_artifact_mb'])}",
        f"- Raw BF16 estimate: {size['dense_bf16_mb_estimate']:.4f} MB",
        f"- Raw INT4 estimate: {size['int4_raw_mb_estimate']:.4f} MB",
        "",
        "## Size Comparison",
        "",
        "| Model | Params | x TinyMind | Artifact scope |",
        "|---|---:|---:|---|",
    ]
    for row in report["size_comparison"]:
        lines.append(
            f"| {row['model']} | {row['params']:,} | {row['relative_to_tinymind_params']:.1f}x | {row['comparison_scope']} |"
        )
    lines += [
        "",
        "## Accessible Official/Public Harness",
        "",
        f"- MMLU-Pro samples: {mmlu['samples']}",
        f"- MMLU-Pro accuracy: {mmlu['accuracy']:.4f}",
        "- Official rank claimed: false",
        "- World-best claim allowed: false",
        "",
        "## Blockers",
    ]
    for row in report["blockers"]:
        lines.append(f"- {row['name']}: {row['status']} ({row.get('error_type', row.get('path', ''))})")
    lines.append("")
    return "\n".join(lines)


def _fmt_mb(value: float | None) -> str:
    return "missing" if value is None else f"{value:.4f} MB"