| """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 |
| 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 |
| 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 = 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)) |
|
|
