File size: 2,995 Bytes
0052e99 | 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 | """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 # noqa: E402
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 = {} # (prompt id, Beam effort, reasoning text) -> level it was trained as
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()
|