| """Build a decontaminated, verifiable prompt pool for Beam -> Dory distillation. |
| |
| Sources (train splits only): |
| math zwhe99/DeepMath-103K difficulty >= 5, closed-form final answers |
| science nvidia/OpenScienceReasoning-2 multiple choice, single-letter answers |
| |
| Decontamination: drop any prompt sharing a 13-gram with MATH-500, AIME 2024/2025, |
| AMC 2023 or GPQA-style eval questions we can access without auth. |
| """ |
|
|
| import argparse |
| import hashlib |
| import json |
| import random |
| import re |
| from pathlib import Path |
|
|
| import duckdb |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| SOURCES = ROOT / "data" / "sources" |
|
|
| EVAL_SETS = [ |
| ("read_json_auto('hf://datasets/HuggingFaceH4/MATH-500/test.jsonl')", "problem"), |
| ("'hf://datasets/HuggingFaceH4/aime_2024/**/*.parquet'", "problem"), |
| ("'hf://datasets/yentinglin/aime_2025/**/*.parquet'", "problem"), |
| ("'hf://datasets/math-ai/amc23/**/*.parquet'", "question"), |
| ] |
|
|
| MCQ_PREFIX = ( |
| "Solve the following problem. Make sure to put the answer (and only answer) " |
| "inside \\boxed{}.\n\n" |
| ) |
| MATH_SUFFIX = "\n\nPut your final answer within \\boxed{}." |
|
|
|
|
| def norm_words(text): |
| return re.sub(r"[^a-z0-9 ]+", " ", text.lower()).split() |
|
|
|
|
| def ngrams(text, n=13): |
| w = norm_words(text) |
| return {" ".join(w[i : i + n]) for i in range(max(0, len(w) - n + 1))} |
|
|
|
|
| def eval_ngrams(con): |
| grams = set() |
| for src, col in EVAL_SETS: |
| repo = src.split("datasets/")[1].split("/", 2)[:2] |
| repo = "/".join(repo) |
| try: |
| rows = con.execute(f"SELECT {col} FROM {src}").fetchall() |
| except Exception as e: |
| print(f" decontam: skipped {repo}: {str(e)[:120]}") |
| continue |
| for (q,) in rows: |
| grams |= ngrams(q or "") |
| print(f" decontam: {repo}: {len(rows)} questions") |
| return grams |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--math", type=int, default=6000) |
| ap.add_argument("--science", type=int, default=4000) |
| ap.add_argument("--seed", type=int, default=20261022) |
| ap.add_argument("--out", default=str(ROOT / "data" / "pool.jsonl")) |
| args = ap.parse_args() |
|
|
| con = duckdb.connect() |
| con.execute("INSTALL httpfs; LOAD httpfs;") |
| rng = random.Random(args.seed) |
|
|
| print("loading eval sets for decontamination") |
| bad = eval_ngrams(con) |
|
|
| print("loading DeepMath-103K (difficulty >= 5)") |
| math_rows = con.execute( |
| """SELECT question, final_answer, difficulty, topic |
| FROM 'hf://datasets/zwhe99/DeepMath-103K/**/*.parquet' |
| WHERE difficulty >= 5""" |
| ).fetchall() |
| print(f" {len(math_rows)} candidates") |
|
|
| print("loading OpenScienceReasoning-2 (single-letter MCQ)") |
| sci_rows = con.execute( |
| """SELECT input, expected_answer |
| FROM 'hf://datasets/nvidia/OpenScienceReasoning-2/**/*.parquet' |
| WHERE length(expected_answer) = 1 AND expected_answer BETWEEN 'A' AND 'J'""" |
| ).fetchall() |
| print(f" {len(sci_rows)} candidates") |
|
|
| pool, seen, dropped = [], set(), 0 |
|
|
| def add(item): |
| nonlocal dropped |
| key = hashlib.sha1(" ".join(norm_words(item["prompt"])).encode()).hexdigest() |
| if key in seen: |
| return False |
| if ngrams(item["prompt"]) & bad: |
| dropped += 1 |
| return False |
| seen.add(key) |
| item["id"] = f"{item['source']}-{key[:12]}" |
| pool.append(item) |
| return True |
|
|
| |
| bands = {} |
| for q, a, d, t in math_rows: |
| if not a or len(a) > 80: |
| continue |
| bands.setdefault(min(int(d), 9), []).append((q, a, d, t)) |
| per_band = args.math // len(bands) |
| for band, rows in sorted(bands.items()): |
| rng.shuffle(rows) |
| taken = 0 |
| for q, a, d, t in rows: |
| if taken >= per_band: |
| break |
| taken += add( |
| { |
| "source": "deepmath", |
| "domain": "math", |
| "prompt": q.strip() + MATH_SUFFIX, |
| "answer": a, |
| "answer_type": "math", |
| "difficulty": d, |
| "topic": t, |
| } |
| ) |
| print(f" math band {band}: {taken}") |
|
|
| rng.shuffle(sci_rows) |
| taken = 0 |
| for q, a in sci_rows: |
| if taken >= args.science: |
| break |
| prompt = q if q.startswith(MCQ_PREFIX) else MCQ_PREFIX + q |
| taken += add( |
| { |
| "source": "openscience2", |
| "domain": "science", |
| "prompt": prompt.strip(), |
| "answer": a, |
| "answer_type": "mcq", |
| } |
| ) |
| print(f" science: {taken}") |
|
|
| rng.shuffle(pool) |
| with open(args.out, "w") as f: |
| for item in pool: |
| f.write(json.dumps(item, ensure_ascii=False) + "\n") |
| print(f"wrote {len(pool)} prompts to {args.out} (decontam dropped {dropped})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
