File size: 4,191 Bytes
079e13c | 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 | """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) # stable: function-style first, each part still shuffled
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()
|