| """Export the reasoning-effort data into release/dataset/data/ (Parquet). |
| |
| .venv/bin/python -I scripts/export_effort_data.py |
| |
| effort_sft.parquet the training set behind Dory-Beam-4B's reasoning_effort (scripts/build_effort_dataset.py) |
| effort_traces.parquet every Beam generation at low, high and xhigh effort, correct or not, with its grade |
| """ |
|
|
| import glob |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT / "scripts")) |
| from verify import check |
|
|
| TEACHER = {"beam-api": "api", "beam-nvfp4-selfhost-sglang": "selfhost-nvfp4", "beam-fp8-selfhost-sglang": "selfhost-fp8"} |
| OUT = ROOT / "release" / "dataset" / "data" |
|
|
|
|
| def main(): |
| eval_ids = {l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()} |
| sft = [json.loads(l) for l in open(ROOT / "data" / "final_effort" / "sft_effort.jsonl")] |
| used = {} |
| rows = [] |
| for r in sft: |
| assert r["prompt_id"] not in eval_ids |
| used[(r["prompt_id"], r["beam_effort"], r["messages"][1]["reasoning_content"])] = r["effort"] |
| rows.append({"id": r["id"], "prompt_id": r["prompt_id"], "reasoning_effort": r["effort"], |
| "beam_effort": r["beam_effort"], "tagged": r["tagged"], "domain": r["domain"], |
| "teacher": TEACHER[r["teacher"]], "messages": r["messages"], "gold": r["gold"]}) |
| pq.write_table(pa.Table.from_pylist(rows), OUT / "effort_sft.parquet", compression="zstd") |
| print(f"effort_sft.parquet: {len(rows)}") |
|
|
| rows = [] |
| for path in sorted(glob.glob(str(ROOT / "data" / "selfhost_effort" / "*.jsonl"))): |
| for line in open(path): |
| r = json.loads(line) |
| if "error" in r: |
| continue |
| assert r["id"] not in eval_ids |
| ok, pred = check(r["content"], r["gold"], r["answer_type"]) |
| rows.append({ |
| "id": r["id"], "domain": r["domain"], "source": r["source"], "topic": r.get("topic"), |
| "difficulty": r["difficulty"], "teacher": TEACHER[r["teacher"]], "reasoning_effort": r["effort"], |
| "prompt": r["prompt"], "reasoning": r["reasoning"], "content": r["content"], "gold": r["gold"], |
| "answer_type": r["answer_type"], "pred": pred, |
| "correct": bool(ok and r["finish_reason"] == "stop" and r["reasoning"].strip()), |
| "finish_reason": r["finish_reason"], "completion_tokens": r["usage"].get("completion_tokens"), |
| "trained_as": used.get((r["id"], r["effort"], r["reasoning"])), |
| }) |
| pq.write_table(pa.Table.from_pylist(rows), OUT / "effort_traces.parquet", compression="zstd") |
| n_used = sum(r["trained_as"] is not None for r in rows) |
| print(f"effort_traces.parquet: {len(rows)} ({n_used} used in effort_sft)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
