scripts/teacher_bakeoff.py

"""Pick a teacher empirically: keep-rate + trace-length-vs-ceiling per model.

    # after setting TFY_API_KEY in .env:
    python scripts/teacher_bakeoff.py --models claude-fable-5 gpt-5.6-sol --n 40

For each model it distills N solutions (same problems), grades the boxed answer,
and reports: keep-rate (higher = better/cheaper data), median/p90 trace length in
STUDENT tokens, and the fraction that would overflow the generator's max_length
(dropped). Highest keep-rate that mostly fits the ceiling wins.

The V-critique job is easier (label is known from PRM800K), so a lighter/cheaper
model usually suffices there regardless of who wins for G.
"""
import argparse
import statistics as stats
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from mathcompose.common.chat import load_tokenizer  # noqa: E402
from mathcompose.common.math_grade import extract_last_boxed, grade_answer  # noqa: E402
from mathcompose.datagen.gen_generator_data import Problem, synthesize_solution  # noqa: E402
from mathcompose.teachers import get_teacher  # noqa: E402


def load_problems(source, split, n, problem_field, solution_field):
    """Stream the source (no full download) and collect the first n problems that
    have a parseable boxed gold answer."""
    from datasets import load_dataset

    ds = load_dataset(source, split=split, streaming=True)
    out = []
    for row in ds:
        prob = row.get(problem_field)
        ans = extract_last_boxed(row.get(solution_field) or "")
        if prob and ans:
            out.append(Problem(problem=prob, answer=ans))
        if len(out) >= n:
            break
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", nargs="+", required=True, help="gateway model slugs to compare")
    ap.add_argument("--teacher", default="promptlens")
    ap.add_argument("--n", type=int, default=40)
    ap.add_argument("--source", default="AI-MO/NuminaMath-CoT")
    ap.add_argument("--split", default="train")
    ap.add_argument("--problem-field", default="problem")
    ap.add_argument("--solution-field", default="solution")
    ap.add_argument("--max-tokens", type=int, default=2048, help="teacher output cap")
    ap.add_argument("--ceiling", type=int, default=2048, help="generator max_length; traces over this are dropped")
    ap.add_argument("--base-id", default="Qwen/Qwen2.5-Math-1.5B-Instruct")
    args = ap.parse_args()

    problems = load_problems(args.source, args.split, args.n, args.problem_field, args.solution_field)
    print(f"loaded {len(problems)} problems from {args.source}\n")
    tok = load_tokenizer(args.base_id)

    rows = []
    for slug in args.models:
        teacher = get_teacher(args.teacher, model=slug)
        kept, lens, t0 = 0, [], time.time()
        for p in problems:
            try:
                sol = synthesize_solution(teacher, p.problem, temperature=0.7, max_tokens=args.max_tokens)
            except Exception as e:
                print(f"  [{slug}] call failed: {type(e).__name__}: {e}")
                continue
            n_tok = len(tok(sol)["input_ids"])
            lens.append(n_tok)
            if grade_answer(extract_last_boxed(sol), p.answer):
                kept += 1
        dt = time.time() - t0
        n = len(lens) or 1
        rows.append({
            "model": slug,
            "keep_rate": kept / len(problems),
            "median_tok": int(stats.median(lens)) if lens else 0,
            "p90_tok": int(sorted(lens)[int(0.9 * (len(lens) - 1))]) if lens else 0,
            "over_ceiling": sum(1 for x in lens if x > args.ceiling) / n,
            "sec_per_ex": dt / len(problems),
        })

    print(f"\n{'model':<24}{'keep':>8}{'med_tok':>9}{'p90_tok':>9}{'>ceil':>8}{'s/ex':>8}")
    print("-" * 66)
    for r in sorted(rows, key=lambda r: -r["keep_rate"]):
        print(f"{r['model']:<24}{r['keep_rate']:>7.0%}{r['median_tok']:>9}{r['p90_tok']:>9}"
              f"{r['over_ceiling']:>7.0%}{r['sec_per_ex']:>7.1f}")
    print("\nPick the highest keep-rate whose p90_tok comfortably fits the ceiling "
          f"({args.ceiling}). Set it via TFY_MODEL or --model at build time.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())