File size: 3,256 Bytes
f698686 b1cf75e f698686 b1cf75e f698686 b1cf75e f698686 b1cf75e f698686 | 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 | """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"],
# SGLang n>1 responses report only total usage, so self-hosted counts are per-request estimates
"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()
|