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

"""Hard benchmark suite aggregator for TinyMind."""

from __future__ import annotations

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

from evaluation.logic_eval import run_logic_eval
from evaluation.official_hard_eval import run_official_hard_eval


HARD_BENCHMARK_TARGETS = (
    "MMLU-Pro",
    "GPQA Diamond",
    "IFEval",
    "LiveCodeBench",
    "SWE-bench Verified/Pro",
    "AIME/MATH-500",
    "Long-context passkey/chunk recall",
)


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 _score_percent(value: float) -> float:
    if value <= 1.0:
        return max(0.0, min(100.0, value * 100.0))
    return max(0.0, min(100.0, value))


def _metric(axis: str, score: float, source: str, samples: int | None = None, notes: str = "") -> dict:
    return {
        "axis": axis,
        "score": _score_percent(score),
        "source": source,
        "samples": samples,
        "notes": notes,
    }


def _import_memory(memory_report: str | Path | None) -> dict:
    memory = _load(memory_report)
    passed = bool(memory.get("passkey_recall", {}).get("passed"))
    measured = int(memory.get("measured_tokens", 0) or 0)
    score = 100.0 if passed and measured >= 10_000_000 else 0.0
    return _metric(
        "long_context_exact",
        score,
        str(memory_report) if memory_report else "missing",
        samples=measured or None,
        notes="10M exact archive passkey gate" if score else "missing or below 10M exact recall gate",
    )


def run_hard_benchmark_suite(
    checkpoint_path: str | Path,
    out_dir: str | Path,
    mmlu_limit: int = 20,
    memory_report: str | Path | None = "reports/extreme_memory_10m/extreme_memory_report.json",
    safetensors_path: str | Path | None = None,
    int4_artifact_path: str | Path | None = None,
    skip_mmlu: bool = False,
    skip_logic: bool = False,
) -> dict:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    metrics: list[dict] = []
    artifacts: dict[str, Any] = {}
    official = None
    if not skip_mmlu:
        official = run_official_hard_eval(
            checkpoint_path=checkpoint_path,
            out_dir=out / "official_hard_eval",
            mmlu_limit=mmlu_limit,
            safetensors_path=safetensors_path,
            int4_artifact_path=int4_artifact_path,
        )
        artifacts["official_hard_eval"] = official.get("json_path")
        mmlu = official.get("results", {}).get("mmlu_pro", {})
        metrics.append(
            _metric(
                "mmlu_pro",
                float(mmlu.get("accuracy", 0.0) or 0.0),
                str(artifacts["official_hard_eval"]),
                samples=int(mmlu.get("samples", 0) or 0),
                notes="TIGER-Lab/MMLU-Pro local public harness",
            )
        )
        size = official.get("size", {})
    else:
        size = {"total_params": None, "million_params": None}
        metrics.append(_metric("mmlu_pro", 0.0, "skipped", notes="skipped by flag"))

    if not skip_logic:
        logic = run_logic_eval(checkpoint_path, out / "logic_ifeval_style")
        artifacts["logic_ifeval_style"] = logic.get("report_path")
        metrics.append(
            _metric(
                "logic_ifeval_style",
                float(logic.get("accuracy", 0.0) or 0.0),
                str(logic.get("report_path")),
                samples=int(logic.get("samples", 0) or 0),
                notes="local strict option-following logic/IFEval-style smoke",
            )
        )
    else:
        metrics.append(_metric("logic_ifeval_style", 0.0, "skipped", notes="skipped by flag"))

    metrics.extend(
        [
            _metric(
                "livecodebench_style",
                0.0,
                "not_run",
                notes="requires code-generation model endpoint and official/live dataset harness",
            ),
            _metric(
                "swe_bench_style",
                0.0,
                "not_run",
                notes="requires full agentic repo-edit harness with Docker test execution",
            ),
            _metric(
                "aime_math_style",
                0.0,
                "not_run",
                notes="requires math answer extraction/evaluator and calibrated prompts",
            ),
            _import_memory(memory_report),
        ]
    )

    measured = [row for row in metrics if row["source"] not in {"missing", "not_run", "skipped"}]
    average = sum(row["score"] for row in metrics) / max(len(metrics), 1)
    hard_blockers = [row["axis"] for row in metrics if row["score"] <= 0.0]
    report = {
        "schema_version": "tinymind-hard-benchmark-suite-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "checkpoint_path": str(checkpoint_path),
        "targets": list(HARD_BENCHMARK_TARGETS),
        "size": size,
        "metrics": metrics,
        "artifacts": artifacts,
        "summary": {
            "average_score": average,
            "measured_axes": len(measured),
            "axis_count": len(metrics),
            "hard_blockers": hard_blockers,
        },
        "suite_gate": {
            "passed": bool(metrics),
            "meaning": "Report generation succeeded; this does not mean the model is strong on every benchmark.",
        },
        "claim_gate": {
            "world_best_claim_allowed": False,
            "official_rank_claim_allowed": False,
            "reason": "Hard local/public benchmark evidence is useful, but world-best claims require official external rank evidence.",
        },
    }
    json_path = out / "hard_benchmark_suite_report.json"
    md_path = out / "hard_benchmark_suite_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(_markdown(report), encoding="utf-8")
    return report


def _markdown(report: dict) -> str:
    lines = [
        "# TinyMind Hard Benchmark Suite",
        "",
        f"- Average score: {report['summary']['average_score']:.2f}",
        f"- Measured axes: {report['summary']['measured_axes']}/{report['summary']['axis_count']}",
        f"- World-best claim allowed: {report['claim_gate']['world_best_claim_allowed']}",
        "",
        "| Axis | Score | Samples | Source | Notes |",
        "|---|---:|---:|---|---|",
    ]
    for row in report["metrics"]:
        samples = "" if row.get("samples") is None else row["samples"]
        lines.append(f"| {row['axis']} | {row['score']:.2f} | {samples} | {row['source']} | {row['notes']} |")
    lines.extend(["", "## Hard Blockers", ""])
    for blocker in report["summary"]["hard_blockers"]:
        lines.append(f"- {blocker}")
    return "\n".join(lines) + "\n"