| """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 |
| from mathcompose.common.math_grade import extract_last_boxed, grade_answer |
| from mathcompose.datagen.gen_generator_data import Problem, synthesize_solution |
| from mathcompose.teachers import get_teacher |
|
|
|
|
| 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()) |
|
|
