File size: 5,078 Bytes
f0975ca | 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 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """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: # missing/gated eval set: report, keep going
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
# Math: stratify by difficulty band so harder problems aren't swamped.
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()
|