File size: 2,886 Bytes
84bcfb1 | 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 | #!/usr/bin/env python
"""Write the problem files every lab uses (run once, on a machine with internet).
math500.jsonl MATH-500 test (500), gold = its 'answer' field (the paper's benchmark)
gsm8k.jsonl GSM8K main test (1,319), gold = the number after #### (the paper's benchmark)
ifeval.jsonl IFEval (541), graded by the official checkers (the paper's benchmark)
dev256.jsonl 128 GSM8K-train + 128 MATH-train practice problems the paper's agent never used:
where we measure every edit we try (same 50/50 mix as the agent's held-out sets)
build128.jsonl 64 + 64 more unused practice problems: where the capstone builds its own edit
"""
import argparse
import json
import os
import random
ap = argparse.ArgumentParser()
ap.add_argument("--exp8", required=True, help="path to lm_understand-logs/exp8_math")
ap.add_argument("--out", required=True)
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
def dump(name, rows):
with open(os.path.join(a.out, f"{name}.jsonl"), "w") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
print(name, len(rows))
from datasets import load_dataset # noqa: E402
dump("math500", [{"id": r["unique_id"], "source": "math500", "problem": r["problem"], "answer": r["answer"]}
for r in load_dataset("HuggingFaceH4/MATH-500", split="test")])
dump("gsm8k", [{"id": f"gsm8k_test_{i}", "source": "gsm8k", "problem": r["question"],
"answer": r["answer"].split("####")[-1].strip().replace(",", "")}
for i, r in enumerate(load_dataset("openai/gsm8k", "main", split="test"))])
dump("ifeval", [{"id": f"ifeval_{r['key']}", "source": "ifeval", "kind": "ifeval", "prompt": r["prompt"],
"instruction_id_list": r["instruction_id_list"], "kwargs": r["kwargs"]}
for r in load_dataset("google/IFEval", split="train")])
used = set()
for fn in ("splits.json", "fresh_splits.json"):
for v in json.load(open(os.path.join(a.exp8, "runs/p50c/work_notes", fn))).values():
used |= set(v)
def practice(name):
with open(os.path.join(a.exp8, "practice", f"{name}.jsonl")) as f:
return [json.loads(line) for line in f if line.strip()]
rng = random.Random(0)
pools = {}
for name in ("gsm8k_train", "math_train"):
rows = [r for r in practice(name) if r["id"] not in used]
rng.shuffle(rows)
pools[name] = rows
print(name, "unused by the paper's agent:", len(rows))
def pick(rows):
return [{"id": r["id"], "source": r["source"], "problem": r["problem"], "answer": r["answer"],
**({"level": r["level"], "subject": r["subject"]} if "level" in r else {})} for r in rows]
dump("dev256", pick(pools["gsm8k_train"][:128] + pools["math_train"][:128]))
dump("build128", pick(pools["gsm8k_train"][128:192] + pools["math_train"][128:192]))
|