| """Build the reasoning-effort SFT set: Beam traces at low, medium and high effort for the same prompts, each |
| tagged the way Dory-Beam's chat template renders `reasoning_effort` (Dory's own `{reasoning effort: ...}` |
| convention, appended to the user turn). |
| |
| .venv/bin/python -I scripts/build_effort_dataset.py --efforts low,medium,high --pool xhigh:high --longest high |
| |
| Inputs: data/effort_prompts.jsonl (the 20k prompts), low/high/xhigh generations in data/selfhost_effort/*.jsonl, |
| and the existing medium generations (data/gen, data/selfhost). Per prompt and level, one randomly drawn |
| verified-correct trace, or the longest with --longest (used for high, pooled from Beam's high and xhigh). Medium examples are tagged half the time and left untagged otherwise, so the untagged |
| default keeps behaving like medium. Writes data/final_effort/sft_effort.jsonl and stats.json. |
| """ |
|
|
| import argparse |
| import glob |
| import json |
| import random |
| import statistics |
| import sys |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT / "scripts")) |
| from verify import check |
|
|
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--prompts", default=str(ROOT / "data" / "effort_prompts.jsonl")) |
| ap.add_argument("--effort-inputs", default=str(ROOT / "data" / "selfhost_effort" / "*.jsonl")) |
| ap.add_argument("--medium-inputs", nargs="+", |
| default=[str(ROOT / "data" / "gen" / "*.jsonl"), str(ROOT / "data" / "selfhost" / "*.jsonl")]) |
| ap.add_argument("--untagged-medium", type=float, default=0.5) |
| ap.add_argument("--out-dir", default=str(ROOT / "data" / "final_effort")) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--efforts", default="low,medium,high", help="Beam effort levels to include") |
| ap.add_argument("--pool", default="", help="use one Beam level's traces as candidates for another, e.g. xhigh:high") |
| ap.add_argument("--longest", default="", help="levels that take the longest correct candidate instead of a random one") |
| args = ap.parse_args() |
| EFFORTS = tuple(args.efforts.split(",")) |
| pool = dict(x.split(":") for x in args.pool.split(",") if x) |
| longest = set(filter(None, args.longest.split(","))) |
|
|
| prompts = {json.loads(l)["id"]: json.loads(l) for l in open(args.prompts)} |
| eval_ids = {l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()} |
| assert not set(prompts) & eval_ids |
|
|
| gens = defaultdict(list) |
| teachers = defaultdict(set) |
| files = sorted(glob.glob(args.effort_inputs)) |
| files += sorted(p for pat in args.medium_inputs for p in glob.glob(pat)) |
| for path in files: |
| for line in open(path): |
| r = json.loads(line) |
| if "error" in r or r["id"] not in prompts: |
| continue |
| beam_eff = r.get("effort", "medium") |
| if (beam_eff == "medium") != ("selfhost_effort" not in path): |
| continue |
| eff = pool.get(beam_eff, beam_eff) |
| if eff not in EFFORTS: |
| continue |
| r["beam_effort"] = beam_eff |
| ok, _ = check(r["content"], r["gold"], r["answer_type"]) |
| r["correct"] = ok and r.get("finish_reason") == "stop" and bool(r.get("reasoning", "").strip()) |
| gens[(r["id"], eff)].append(r) |
| teachers[eff].add(r.get("teacher", "beam-api")) |
|
|
| rng = random.Random(args.seed) |
| out = Path(args.out_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| stats = {e: {"generations": 0, "correct": 0, "examples": 0, "chars": [], "tokens": []} for e in EFFORTS} |
| n_untagged = 0 |
| with open(out / "sft_effort.jsonl", "w") as f: |
| for pid in sorted(prompts): |
| for eff in EFFORTS: |
| rs = gens.get((pid, eff), []) |
| good = sorted((r for r in rs if r["correct"]), key=lambda r: len(r["reasoning"])) |
| st = stats[eff] |
| st["generations"] += len(rs) |
| st["correct"] += len(good) |
| for r in rs: |
| st["tokens"].append(r["usage"].get("completion_tokens", 0)) |
| if not good: |
| continue |
| r = good[-1] if eff in longest else rng.choice(good) |
| tagged = eff != "medium" or rng.random() >= args.untagged_medium |
| n_untagged += not tagged |
| user = r["prompt"] + (f"\n\n{{reasoning effort: {eff}}}" if tagged else "") |
| st["examples"] += 1 |
| st["chars"].append(len(r["reasoning"]) + len(r["content"])) |
| f.write(json.dumps({ |
| "id": f"{pid}#{eff}", "prompt_id": pid, "effort": eff, "beam_effort": r["beam_effort"], |
| "tagged": tagged, |
| "domain": r["domain"], "teacher": r.get("teacher", "beam-api"), |
| "messages": [{"role": "user", "content": user}, |
| {"role": "assistant", "reasoning_content": r["reasoning"], "content": r["content"]}], |
| "gold": r["gold"], "chars": st["chars"][-1], |
| }, ensure_ascii=False) + "\n") |
|
|
| summary = {"prompts": len(prompts), "untagged_medium": n_untagged, "pool": pool, "longest": sorted(longest), |
| "teachers": {e: sorted(t) for e, t in teachers.items()}} |
| for e, st in stats.items(): |
| summary[e] = {"generations": st["generations"], "correct": st["correct"], |
| "accuracy": round(st["correct"] / max(1, st["generations"]), 3), "examples": st["examples"], |
| "chars_mean": round(statistics.mean(st["chars"])) if st["chars"] else None, |
| "tokens_mean": round(statistics.mean(st["tokens"])) if st["tokens"] else None} |
| (out / "stats.json").write_text(json.dumps(summary, indent=1)) |
| print(json.dumps(summary, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
