| """Assemble generated Beam traces into release/SFT files with summary stats. |
| |
| Outputs (in --out-dir): |
| traces_all.jsonl every successful generation, correct or not, with metadata |
| (held-out eval prompts included, marked by eval_ids.txt) |
| sft_verified.jsonl verified-correct, untruncated, non-eval traces in "messages" form |
| (assistant turn carries reasoning_content, which Dory's chat |
| template renders as <think>...</think>) |
| stats.json counts, accuracy and token statistics per domain/band/effort |
| """ |
|
|
| 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 band(r): |
| if r["domain"] != "math": |
| return "science" |
| return f"math_d{min(int(r['difficulty'] or 0), 9)}" |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--inputs", nargs="+", |
| default=[str(ROOT / "data" / "gen" / "*.jsonl"), str(ROOT / "data" / "selfhost" / "*.jsonl")]) |
| ap.add_argument("--out-dir", default=str(ROOT / "data" / "final")) |
| ap.add_argument("--max-per-prompt", type=int, default=1) |
| |
| |
| ap.add_argument("--select", choices=["shortest", "random"], default="shortest") |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--eval-ids", default=str(ROOT / "data" / "eval_ids.txt")) |
| args = ap.parse_args() |
| eval_ids = {line.strip() for line in open(args.eval_ids) if line.strip()} |
|
|
| recs = [] |
| for path in sorted(p for pattern in args.inputs for p in glob.glob(pattern)): |
| for line in open(path): |
| r = json.loads(line) |
| if "error" not in r: |
| |
| ok, pred = check(r["content"], r["gold"], r["answer_type"]) |
| r["pred"], r["correct"] = pred, ok and r.get("finish_reason") == "stop" |
| recs.append(r) |
|
|
| out = Path(args.out_dir) |
| out.mkdir(parents=True, exist_ok=True) |
|
|
| by_id = defaultdict(list) |
| for r in recs: |
| by_id[r["id"]].append(r) |
|
|
| |
| |
| rng = random.Random(args.seed) |
| n_sft = 0 |
| with open(out / "traces_all.jsonl", "w") as fa, open(out / "sft_verified.jsonl", "w") as fs: |
| for r in recs: |
| fa.write(json.dumps(r, ensure_ascii=False) + "\n") |
| for pid, rs in by_id.items(): |
| if pid in eval_ids: |
| continue |
| good = [r for r in rs if r.get("correct") and r.get("reasoning", "").strip()] |
| good.sort(key=lambda r: len(r["reasoning"]) + len(r["content"])) |
| if args.select == "random": |
| rng.shuffle(good) |
| for r in good[: args.max_per_prompt]: |
| n_sft += 1 |
| fs.write( |
| json.dumps( |
| { |
| "id": r["id"], |
| "domain": r["domain"], |
| "source": r["source"], |
| "effort": r["effort"], |
| "teacher": r.get("teacher", "beam-api"), |
| "difficulty": r.get("difficulty"), |
| "messages": [ |
| {"role": "user", "content": r["prompt"]}, |
| { |
| "role": "assistant", |
| "reasoning_content": r["reasoning"], |
| "content": r["content"], |
| }, |
| ], |
| "gold": r["gold"], |
| "n_samples": len(rs), |
| "n_correct": len(good), |
| "chars": len(r["reasoning"]) + len(r["content"]), |
| "completion_tokens": r["usage"].get("completion_tokens"), |
| }, |
| ensure_ascii=False, |
| ) |
| + "\n" |
| ) |
|
|
| groups = defaultdict(list) |
| for r in recs: |
| split = "eval" if r["id"] in eval_ids else "train" |
| groups[(r["effort"], split, band(r))].append(r) |
| groups[(r["effort"], split, "ALL")].append(r) |
| stats = {"generations": len(recs), "sft_examples": n_sft, "groups": {}} |
| for (effort, split, b), rs in sorted(groups.items()): |
| toks = [r["usage"].get("completion_tokens", 0) for r in rs] |
| ok_toks = [r["usage"].get("completion_tokens", 0) for r in rs if r.get("correct")] |
| stats["groups"][f"{effort}/{split}/{b}"] = { |
| "n": len(rs), |
| "accuracy": round(sum(r.get("correct", False) for r in rs) / len(rs), 3), |
| "truncated": sum(r.get("finish_reason") == "length" for r in rs), |
| "out_tokens_mean": round(statistics.mean(toks)), |
| "out_tokens_median": round(statistics.median(toks)), |
| "out_tokens_correct_mean": round(statistics.mean(ok_toks)) if ok_toks else None, |
| "out_tokens_total": sum(toks), |
| } |
| (out / "stats.json").write_text(json.dumps(stats, indent=1)) |
| print(json.dumps(stats, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
