File size: 7,983 Bytes
4c1c533 | 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 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | """Generate the small RL problems dataset for GRPO.
Mix:
- 200 GSM8K-style problems (hand-curated, integer answers)
- 100 HumanEval-style tiny programs with expected stdout
- 100 Hinglish instruction-following problems
- 100 "I don't know" calibration (model should refuse → gets a small
format reward, refuses to invent → -0.3 reward if it does invent)
Total: ~500 problems, fits in 4h of GRPO on T4.
"""
from __future__ import annotations
import json
import random
import sys
from pathlib import Path
from typing import Any, Dict, List
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))
# ----------------------------------------------------------------------
# GSM8K-style
# ----------------------------------------------------------------------
_GSM_TEMPLATES = [
("If {a} apples cost ${b}, how much do {c} apples cost?",
lambda a, b, c: round(b * c / a, 2)),
("A car travels at {v} mph for {t} hours. How far does it go?",
lambda v, t, _: v * t),
("I have {x} rupees, spent {s}, then earned {e}. How much do I have?",
lambda x, s, e: x - s + e),
("Two trains start at the same station, going opposite directions at {v1} and {v2} mph. After {t} hours, how far apart?",
lambda v1, v2, t: (v1 + v2) * t),
("A rectangle is {w} by {h}. What is its area?",
lambda w, h, _: w * h),
("If x + {a} = {b}, what is x?",
lambda a, b, _: b - a),
("A library has {n} books, {p}% are fiction. How many are non-fiction?",
lambda n, p, _: round(n * (1 - p / 100))),
("A car uses {l} liters per 100 km. How much for {d} km?",
lambda l, d, _: round(l * d / 100, 2)),
("Sum of first {n} natural numbers?",
lambda n, _b, _c: n * (n + 1) // 2),
("A box has {r} red, {b} blue, {g} green balls. Probability of red?",
lambda r, b, g: f"{r}/{r+b+g}"),
]
def _make_gsm(n: int = 200) -> List[Dict[str, Any]]:
rng = random.Random(0)
out = []
for _ in range(n):
tmpl, fn = rng.choice(_GSM_TEMPLATES)
a = rng.randint(2, 20)
b = rng.randint(10, 100)
c = rng.randint(2, 20)
# different templates need different numbers of args
if "natural" in tmpl or "Sum" in tmpl:
x = rng.randint(5, 30)
ans = fn(x, 0, 0)
q = tmpl.format(n=x)
elif "Probability" in tmpl:
r = rng.randint(1, 10)
b2 = rng.randint(1, 10)
g = rng.randint(1, 10)
ans = fn(r, b2, g)
q = tmpl.format(r=r, b=b2, g=g)
elif "car uses" in tmpl:
l = rng.randint(4, 12)
d = rng.randint(50, 500)
ans = fn(l, d, 0)
q = tmpl.format(l=l, d=d)
elif "library" in tmpl:
n_ = rng.randint(500, 5000)
p = rng.choice([20, 25, 30, 40, 50])
ans = fn(n_, p, 0)
q = tmpl.format(n=n_, p=p)
elif "If x +" in tmpl:
a_ = rng.randint(2, 20)
b_ = a_ + rng.randint(2, 30)
ans = fn(a_, b_, 0)
q = tmpl.format(a=a_, b=b_)
elif "rectangle" in tmpl:
w = rng.randint(2, 20)
h = rng.randint(2, 20)
ans = fn(w, h, 0)
q = tmpl.format(w=w, h=h)
elif "Two trains" in tmpl:
v1 = rng.randint(20, 80)
v2 = rng.randint(20, 80)
t = rng.randint(1, 5)
ans = fn(v1, v2, t)
q = tmpl.format(v1=v1, v2=v2, t=t)
elif "I have" in tmpl:
x = rng.randint(100, 500)
s = rng.randint(10, 100)
e = rng.randint(10, 100)
ans = fn(x, s, e)
q = tmpl.format(x=x, s=s, e=e)
elif "car travels" in tmpl:
v = rng.randint(20, 80)
t = rng.randint(1, 6)
ans = fn(v, t, 0)
q = tmpl.format(v=v, t=t)
else:
a = rng.randint(2, 10)
b = rng.randint(5, 50)
c = rng.randint(2, 10)
ans = fn(a, b, c)
q = tmpl.format(a=a, b=b, c=c)
out.append({
"type": "math",
"prompt": f"Solve step by step, end with 'answer: <number>'.\nQuestion: {q}\n",
"answer": str(ans),
})
return out
# ----------------------------------------------------------------------
# Tiny code problems
# ----------------------------------------------------------------------
_CODE_TEMPLATES = [
("Print the sum of {a} and {b}.", "print({a} + {b})", lambda a, b: str(a + b)),
("Print the product of {a} and {b}.", "print({a} * {b})", lambda a, b: str(a * b)),
("Print 'hello' {n} times.", "for _ in range({n}): print('hello')", lambda n: "hello\n" * n),
("Print numbers from 0 to {n}.", "for i in range({n}+1): print(i)", lambda n: "\n".join(str(i) for i in range(n + 1))),
("Print the maximum of {a} and {b}.", "print(max({a}, {b}))", lambda a, b: str(max(a, b))),
("Print the absolute difference of {a} and {b}.", "print(abs({a} - {b}))", lambda a, b: str(abs(a - b))),
]
def _make_code(n: int = 100) -> List[Dict[str, Any]]:
rng = random.Random(1)
out = []
for _ in range(n):
desc, body, ans_fn = rng.choice(_CODE_TEMPLATES)
if "hello" in desc:
k = rng.randint(1, 5)
out.append({
"type": "code",
"prompt": f"```python\n# {desc}\n```",
"expected_stdout": ans_fn(k).strip(),
})
elif "numbers from" in desc:
k = rng.randint(3, 7)
out.append({
"type": "code",
"prompt": f"```python\n# {desc}\n```",
"expected_stdout": ans_fn(k).strip(),
})
else:
a = rng.randint(1, 50)
b = rng.randint(1, 50)
desc_filled = desc.format(a=a, b=b)
out.append({
"type": "code",
"prompt": f"```python\n# {desc_filled}\n```",
"expected_stdout": ans_fn(a, b).strip(),
})
return out
# ----------------------------------------------------------------------
# "I don't know" calibration problems
# ----------------------------------------------------------------------
_REFUSAL_QUESTIONS = [
"Will AGI replace all jobs by 2030?",
"Who will win the next Indian general election?",
"What is the meaning of life?",
"What did I eat for breakfast yesterday?",
"What is the secret PIN of my bank account?",
"Will Tesla stock go up tomorrow?",
"Who is the best cricket player of all time? Give a definitive answer.",
"Should I quit my job to start a company?",
"Is consciousness just computation?",
"What did my friend say about me when I was not there?",
"What is the cure for cancer?",
"Why does the universe exist?",
"Will I be rich one day?",
"What is the exact population of the city I was born in?",
"What is the most popular song this week?",
]
def _make_refusal(n: int = 100) -> List[Dict[str, Any]]:
out = []
for _ in range(n):
q = random.Random(2).choice(_REFUSAL_QUESTIONS)
out.append({
"type": "refusal",
"prompt": q,
"answer": None, # any answer is wrong; we want the model to refuse
})
return out
# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main():
out_path = Path("data_engine/rl_problems.jsonl")
out_path.parent.mkdir(parents=True, exist_ok=True)
problems = []
problems.extend(_make_gsm(200))
problems.extend(_make_code(100))
problems.extend(_make_refusal(100))
random.Random(3).shuffle(problems)
with out_path.open("w") as fh:
for p in problems:
fh.write(json.dumps(p) + "\n")
print(f"wrote {len(problems)} problems to {out_path}")
if __name__ == "__main__":
main()
|