| """Training mix for the coding fix (run rC): Beam's test-verified TACO solutions plus a replay sample of the effort |
| data that Dory-Beam-4B v2 was trained on, so the effort levels and the math are not forgotten. |
| |
| .venv/bin/python -I scripts/build_code_sft.py --beam 'data/code/beam/run-code-*.jsonl' \ |
| --verdicts data/code/verdicts.jsonl [--max-code 8000] [--replay-ratio 1.0] |
| # -> data/code/sft_code_mix.jsonl, data/code/sft_code_mix_stats.json |
| |
| Coding examples are Beam at medium effort, tagged `{reasoning effort: medium}` half the time and untagged otherwise, |
| as in scripts/build_effort_dataset.py. Replay examples are drawn uniformly from data/final_effort/sft_effort.jsonl. |
| """ |
|
|
| import argparse |
| import glob |
| import json |
| import random |
| from collections import Counter |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--beam", required=True) |
| ap.add_argument("--verdicts", required=True) |
| ap.add_argument("--max-code", type=int, default=0) |
| ap.add_argument("--replay-ratio", type=float, default=1.0) |
| ap.add_argument("--max-chars", type=int, default=40000, help="drop very long traces (~12k tokens)") |
| ap.add_argument("--prefer-fn", action="store_true", |
| help="fill --max-code with function-style problems first (the HumanEval/MBPP format)") |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
| rng = random.Random(args.seed) |
|
|
| passed = {v["key"] for v in map(json.loads, open(args.verdicts)) if v["passed"]} |
| code = [] |
| for path in sorted(glob.glob(args.beam)): |
| for line in open(path): |
| r = json.loads(line) |
| key = f"{r.get('id')}#{r.get('sample_idx', 0)}" |
| if key not in passed or not r["reasoning"].strip(): |
| continue |
| if len(r["reasoning"]) + len(r["content"]) > args.max_chars: |
| continue |
| code.append(r) |
| rng.shuffle(code) |
| if args.prefer_fn: |
| fn = {t["id"] for t in map(json.loads, open(ROOT / "data" / "code" / "tests.jsonl")) if t["fn_name"]} |
| code.sort(key=lambda r: r["id"] not in fn) |
| if args.max_code: |
| code = code[:args.max_code] |
|
|
| rows = [] |
| for r in code: |
| tagged = rng.random() < 0.5 |
| user = r["prompt"] + ("\n\n{reasoning effort: medium}" if tagged else "") |
| rows.append({"id": f"{r['id']}#code", "prompt_id": r["id"], "effort": "medium", "beam_effort": "medium", |
| "tagged": tagged, "domain": "code", "teacher": r["teacher"], |
| "messages": [{"role": "user", "content": user}, |
| {"role": "assistant", "reasoning_content": r["reasoning"], "content": r["content"]}], |
| "gold": "", "chars": len(r["reasoning"]) + len(r["content"])}) |
| effort = [json.loads(l) for l in open(ROOT / "data" / "final_effort" / "sft_effort.jsonl")] |
| replay = rng.sample(effort, min(len(effort), round(len(rows) * args.replay_ratio))) |
| mix = rows + replay |
| rng.shuffle(mix) |
|
|
| out = ROOT / "data" / "code" / "sft_code_mix.jsonl" |
| with open(out, "w") as f: |
| for r in mix: |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") |
| stats = {"code_examples": len(rows), "replay_examples": len(replay), |
| "code_chars_mean": round(sum(r["chars"] for r in rows) / max(1, len(rows))), |
| "replay_chars_mean": round(sum(r.get("chars", 0) for r in replay) / max(1, len(replay))), |
| "replay_mix": Counter(f"{r['effort']}{'' if r['tagged'] else '-untagged'}" for r in replay), |
| "code_sources": Counter(r["source"] for r in code), |
| "code_function_style": sum("starter code" in r["prompt"] for r in code)} |
| stats["approx_tokens_M"] = round((stats["code_chars_mean"] * len(rows) + stats["replay_chars_mean"] * len(replay)) |
| / 3.3 / 1e6, 1) |
| (ROOT / "data" / "code" / "sft_code_mix_stats.json").write_text(json.dumps(stats, indent=1)) |
| print(json.dumps(stats, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
