Dory-Beam-4B / code /scripts /build_space_data.py
ProCreations's picture
v2: Beam reasoning-effort levels (low/medium/high), Accurate variant; v1 is tagged
0052e99 verified
Raw History Blame Contribute Delete
3.23 kB
"""Build space/data.json for the side-by-side Space: base Dory vs Dory-Beam-4B on the same eval problems.
python3 scripts/build_space_data.py [--loops 2] [--seed 0]
Problems are a seeded random sample per benchmark (not hand-picked); both answers are each model's first
sample for that problem from the eval runs. Also writes per-benchmark summary stats over the full eval.
"""
import argparse
import json
import random
from collections import defaultdict
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
RES = ROOT / "data" / "eval" / "results"
# No GPQA examples: its authors ask that questions not appear online in plain text. GPQA stays in the stats table.
QUOTA = {"math500": 10, "aime24": 5, "aime25": 5, "amc23": 5, "heldout": 9}
EFFORTS = ("low", "default", "high")
NAMES = {"gpqa": "GPQA-Diamond", "math500": "MATH-500", "aime24": "AIME 2024", "aime25": "AIME 2025",
"amc23": "AMC 2023", "heldout": "Held-out"}
def load(tag):
by = defaultdict(dict)
for line in open(RES / f"{tag}.jsonl"):
r = json.loads(line)
by[r["id"]].setdefault(r["sample"], r)
return by
def side(r):
reasoning, _, answer = r["text"].partition("</think>")
return {"reasoning": reasoning.strip(), "answer": answer.strip(), "tokens": r["tokens"],
"correct": r["correct"], "pred": r["pred"], "truncated": r["finish"] == "length"}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--loops", type=int, default=2)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
base = load(f"base_loops{args.loops}-64k")
# Dory-Beam-4B v2 at each reasoning effort ("default" = no effort tag)
beam = {e: load(f"rE-t050_loops{args.loops}-{e}-64k") for e in EFFORTS}
probs = [json.loads(l) for l in open(ROOT / "data" / "evalsets.jsonl")]
rng = random.Random(args.seed)
items = []
for bench, k in QUOTA.items():
pool = [p for p in probs if p["bench"] == bench]
for p in sorted(rng.sample(pool, k), key=lambda p: p["id"]):
items.append({"id": p["id"], "bench": NAMES[bench], "prompt": p["prompt"], "gold": p["answer"],
"base": side(base[p["id"]][0]), "beam": {e: side(beam[e][p["id"]][0]) for e in EFFORTS}})
stats = []
for bench, name in NAMES.items():
row = {"bench": name}
def agg(runs):
rs = [r for pid, samples in runs.items() for r in samples.values() if r["bench"] == bench]
return {"acc": round(100 * sum(r["correct"] for r in rs) / len(rs), 1),
"tokens": round(sum(r["tokens"] for r in rs) / len(rs))}
row["base"] = agg(base)
row["beam"] = {e: agg(beam[e]) for e in EFFORTS}
stats.append(row)
out = ROOT / "space" / "data.json"
out.parent.mkdir(exist_ok=True)
out.write_text(json.dumps({"loops": args.loops, "seed": args.seed, "stats": stats, "items": items},
ensure_ascii=False))
both = sum(i["base"]["correct"] and i["beam"]["default"]["correct"] for i in items)
print(f"{len(items)} problems ({both} both correct) -> {out} ({out.stat().st_size / 1e6:.1f} MB)")
if __name__ == "__main__":
main()