File size: 5,309 Bytes
67ebcef | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """Builds the pool of verifiable training problems (math with reference answers, code with unit tests),
decontaminated against every prompt in the frozen eval suite (13-gram overlap on normalized text).
Output: data/train/pool_math.jsonl, data/train/pool_code.jsonl, data/train/dev.jsonl (held-out dev set used for
all intermediate decisions, so the test suite is only touched for the baseline and the final comparison).
"""
import glob
import json
import os
import random
import re
from datasets import load_dataset
ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
OUT = os.path.join(ROOT, "data", "train")
os.makedirs(OUT, exist_ok=True)
rng = random.Random(7)
MATH_SUFFIX = "\n\nPlease reason step by step, and put your final answer within \\boxed{}."
def words(s):
return re.sub(r"[^a-z0-9 ]", " ", s.lower()).split()
def ngrams(s, n=13):
w = words(s)
return {" ".join(w[i:i + n]) for i in range(max(0, len(w) - n + 1))} or ({" ".join(w)} if w else set())
eval_grams = set()
eval_short = set()
for path in glob.glob(os.path.join(ROOT, "data", "eval", "*.jsonl")):
with open(path, encoding="utf-8") as f:
for l in f:
r = json.loads(l)
text = r.get("code_prompt") or r["prompt"].replace(MATH_SUFFIX, "")
eval_grams |= ngrams(text)
eval_short.add(" ".join(words(text))[:200])
def contaminated(text):
if " ".join(words(text))[:200] in eval_short:
return True
return bool(ngrams(text) & eval_grams)
def last_boxed(text):
idx = text.rfind("\\boxed")
if idx < 0:
return None
i = text.find("{", idx)
depth = 0
for j in range(i, len(text)):
depth += text[j] == "{"
depth -= text[j] == "}"
if depth == 0:
return text[i + 1: j]
return None
math_rows, seen, dropped = [], set(), 0
def add_math(src, problem, answer, level=None):
global dropped
if not answer or len(answer) > 60 or "\\text{" in answer and len(answer) > 25:
return # keep answers that a symbolic checker can verify reliably
key = " ".join(words(problem))[:300]
if key in seen:
return
seen.add(key)
if contaminated(problem):
dropped += 1
return
math_rows.append({"id": f"{src}/{len(math_rows)}", "task": "math", "src": src, "level": level,
"prompt": problem.strip() + MATH_SUFFIX, "answer": answer.strip()})
ds = load_dataset("DigitalLearningGmbH/MATH-lighteval", split="train")
for r in ds:
lvl = int(r["level"].split()[-1]) if r["level"].split()[-1].isdigit() else None
if lvl and lvl >= 3:
add_math("math_train", r["problem"], last_boxed(r["solution"]), lvl)
n_math = len(math_rows)
ds = load_dataset("agentica-org/DeepScaleR-Preview-Dataset", split="train")
for r in ds:
add_math("deepscaler", r["problem"], r["answer"])
print(f"math: {n_math} from MATH train, {len(math_rows) - n_math} from DeepScaleR, dropped {dropped} contaminated")
code_rows, dropped_code = [], 0
ds = load_dataset("google-research-datasets/mbpp", "full")
for split in ("train", "validation", "prompt"):
for r in ds[split]:
if contaminated(r["text"]):
dropped_code += 1
continue
code_rows.append({
"id": f"mbpp_train/{r['task_id']}", "task": "code_fn", "src": "mbpp_train",
"prompt": r["text"].strip() + "\nYour code should pass this test:\n" + r["test_list"][0]
+ "\n\nWrite the solution as a Python function in a single ```python code block.",
"test": "\n".join(r["test_setup_code"].splitlines() + r["test_list"] + r.get("challenge_test_list", []))})
n_mbpp = len(code_rows)
ds = load_dataset("likaixin/TACO-verified", split="train")
for r in ds:
q = r["question"]
if len(q) > 6000 or contaminated(q):
dropped_code += 1
continue
try:
io = json.loads(r["input_output"])
except Exception:
continue
if io.get("fn_name") or not io.get("inputs") or len(io["inputs"]) < 2:
continue # keep stdin/stdout problems with at least two tests
if any(not isinstance(x, str) for x in io["inputs"] + io["outputs"]):
continue
if sum(len(x) for x in io["inputs"]) > 200_000:
continue
code_rows.append({
"id": f"taco/{r['id'] if 'id' in r else len(code_rows)}", "task": "code_io", "src": "taco",
"difficulty": r.get("difficulty"),
"prompt": q.strip() + "\n\nWrite a Python 3 program that reads from standard input and writes to standard "
"output. Put the complete program in a single ```python code block.",
"inputs": io["inputs"][:20], "outputs": io["outputs"][:20]})
print(f"code: {n_mbpp} MBPP, {len(code_rows) - n_mbpp} TACO, dropped {dropped_code} contaminated/long")
rng.shuffle(math_rows)
rng.shuffle(code_rows)
# dev: 300 math + 150 code problems, never trained on, used for every intermediate decision
dev = math_rows[:300] + code_rows[:150]
for name, rows in [("dev", dev), ("pool_math", math_rows[300:]), ("pool_code", code_rows[150:])]:
with open(os.path.join(OUT, f"{name}.jsonl"), "w", encoding="utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
print(name, len(rows))
|