| """Write the public dataset files from data/final into release/dataset/data with a stable schema. |
| |
| .venv/bin/python -I scripts/export_release_data.py [--run-a data/final_random/sft_verified.jsonl] |
| |
| sft.parquet one verified-correct Beam trace per prompt (shortest correct), messages form |
| traces_all.parquet every generation, correct or not (held-out eval prompts are never included) |
| |
| Everything in the dataset repo is Parquet: the Hub picks one file reader per dataset, so formats can't be mixed. |
| run_a_ids.txt prompt ids in the 52,660-example snapshot Run A was trained on |
| """ |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| TEACHER = {"beam-api": "api", "beam-nvfp4-selfhost-sglang": "selfhost-nvfp4", "beam-fp8-selfhost-sglang": "selfhost-fp8"} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--final", default=str(ROOT / "data" / "final")) |
| ap.add_argument("--run-a", default=str(ROOT / "data" / "final_random" / "sft_verified.jsonl"), |
| help="any SFT build from the Run A snapshot (same prompt ids)") |
| ap.add_argument("--out", default=str(ROOT / "release" / "dataset" / "data")) |
| args = ap.parse_args() |
| out = Path(args.out) |
| out.mkdir(parents=True, exist_ok=True) |
| eval_ids = {l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()} |
|
|
| rows = [] |
| for line in open(Path(args.final) / "sft_verified.jsonl"): |
| r = json.loads(line) |
| rows.append({"id": r["id"], "domain": r["domain"], "source": r["source"], "difficulty": r["difficulty"], |
| "teacher": TEACHER[r["teacher"]], "reasoning_effort": r["effort"], |
| "messages": r["messages"], "gold": r["gold"], |
| "n_samples": r["n_samples"], "n_correct": r["n_correct"]}) |
| pq.write_table(pa.Table.from_pylist(rows), out / "sft.parquet", compression="zstd") |
| print(f"sft.parquet: {len(rows)}") |
|
|
| rows = [] |
| for line in open(Path(args.final) / "traces_all.jsonl"): |
| r = json.loads(line) |
| assert r["id"] not in eval_ids, f"held-out eval prompt {r['id']} in traces" |
| rows.append({ |
| "id": r["id"], "domain": r["domain"], "source": r["source"], "topic": r.get("topic"), |
| "difficulty": r["difficulty"], |
| "teacher": TEACHER[r.get("teacher", "beam-api")], "reasoning_effort": r["effort"], |
| "prompt": r["prompt"], "reasoning": r["reasoning"], "content": r["content"], |
| "gold": r["gold"], "answer_type": r["answer_type"], "pred": r["pred"], "correct": r["correct"], |
| "finish_reason": r["finish_reason"], |
| |
| "completion_tokens": r["usage"].get("completion_tokens"), |
| }) |
| pq.write_table(pa.Table.from_pylist(rows), out / "traces_all.parquet", compression="zstd") |
| print(f"traces_all.parquet: {len(rows)}") |
|
|
| ids = sorted({json.loads(l)["id"] for l in open(args.run_a)}) |
| (out.parent / "run_a_ids.txt").write_text("\n".join(ids) + "\n") |
| print(f"run_a_ids.txt: {len(ids)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
