File size: 4,293 Bytes
014b14c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """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())
|