"""Build a verifiable task set whose answers are correct by construction. Purpose: the training and evaluation problems for the GRPO lab and the reality check. Every problem is generated from a template together with its answer, so the reference answer is not a label somebody typed: it is the number the generator computed before it wrote the question. Two families are produced. "arith" is multi-step arithmetic word problems, the family the lab trains on. "logic" is counting, calendar and ordering puzzles, the different-kind held-out family the reality check uses to ask whether anything generalised. Platform: all (pure Python; --source gsm8k additionally needs the datasets package) Minimum memory: 8 GB Assumes: Python 3.10 or newer. For --source gsm8k, `datasets` installed and network access to the Hugging Face Hub; the GSM8K card gives the licence as MIT. Usage: python3 make-tasks.py --out-dir tasks --seed 0 python3 make-tasks.py --out-dir tasks --train 240 --heldout 60 --seed 0 python3 make-tasks.py --out-dir tasks --source gsm8k --train 240 --heldout 60 """ from __future__ import annotations import argparse import hashlib import json import random from pathlib import Path from typing import Callable INSTRUCTION = ( "Solve the problem. Put your working inside and tags, then end " "with a single line of the form 'Answer: N' where N is the final number and " "nothing follows it." ) DAYS = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"] NAMES = ["Ada", "Bela", "Chi", "Dara", "Eero", "Fen", "Gita", "Hugo", "Ines", "Jonas", "Kira", "Liam", "Mira", "Nils", "Oona", "Pia", "Rafa", "Sena", "Tomas", "Ulla"] GOODS = [("pencil", "pencils"), ("notebook", "notebooks"), ("cable", "cables"), ("mug", "mugs"), ("tile", "tiles"), ("filter", "filters"), ("bulb", "bulbs")] # -------------------------------------------------------------------------------------- # Family: arithmetic word problems # -------------------------------------------------------------------------------------- def arith_purchase(rng: random.Random) -> tuple[str, int]: boxes = rng.randint(3, 12) per_box = rng.randint(4, 15) used = rng.randint(2, boxes * per_box - 2) singular, plural = rng.choice(GOODS) who = rng.choice(NAMES) question = (f"{who} buys {boxes} boxes of {plural}, with {per_box} {plural} in each box. " f"{used} {plural if used != 1 else singular} are used during the week. " f"How many {plural} are left?") return question, boxes * per_box - used def arith_rate(rng: random.Random) -> tuple[str, int]: per_hour = rng.randint(6, 30) hours = rng.randint(2, 9) extra = rng.randint(1, 40) who = rng.choice(NAMES) question = (f"A machine that {who} operates produces {per_hour} parts an hour. It runs for " f"{hours} hours, and {extra} parts from yesterday are added to the same crate. " f"How many parts are in the crate?") return question, per_hour * hours + extra def arith_split(rng: random.Random) -> tuple[str, int]: people = rng.randint(3, 9) each = rng.randint(2, 12) leftover = rng.randint(1, people - 1) singular, plural = rng.choice(GOODS) question = (f"A crate of {plural} is shared between {people} people. Each person gets " f"{each} {plural}, and {leftover} are left over. How many {plural} were in " f"the crate?") return question, people * each + leftover def arith_growth(rng: random.Random) -> tuple[str, int]: start = rng.randint(5, 40) gain = rng.randint(2, 15) days = rng.randint(3, 12) loss = rng.randint(1, 20) who = rng.choice(NAMES) question = (f"{who} starts a log with {start} entries and adds {gain} entries every day for " f"{days} days. At the end {loss} duplicate entries are deleted. How many entries " f"does the log hold?") return question, start + gain * days - loss def arith_area_cost(rng: random.Random) -> tuple[str, int]: width = rng.randint(2, 15) length = rng.randint(2, 15) price = rng.randint(2, 20) spare = rng.randint(1, 10) question = (f"A floor is {width} metres by {length} metres. Tiles cost {price} units per " f"square metre, and {spare} extra units are spent on adhesive. What is the total " f"cost in units?") return question, width * length * price + spare def arith_two_step_difference(rng: random.Random) -> tuple[str, int]: a_rate = rng.randint(5, 20) b_rate = rng.randint(1, a_rate - 1) minutes = rng.randint(4, 25) question = (f"Two pumps fill the same tank. The first adds {a_rate} litres a minute and the " f"second removes {b_rate} litres a minute. Both run for {minutes} minutes. How " f"many litres are in the tank at the end, if it started empty?") return question, (a_rate - b_rate) * minutes ARITH_TEMPLATES: list[Callable[[random.Random], tuple[str, int]]] = [ arith_purchase, arith_rate, arith_split, arith_growth, arith_area_cost, arith_two_step_difference, ] # -------------------------------------------------------------------------------------- # Family: counting, calendar and ordering puzzles # -------------------------------------------------------------------------------------- def logic_inclusion_exclusion(rng: random.Random) -> tuple[str, int]: both = rng.randint(2, 12) only_a = rng.randint(3, 25) only_b = rng.randint(3, 25) neither = rng.randint(0, 10) question = (f"In a group, {only_a + both} people read the newsletter, {only_b + both} people " f"attend the meeting, {both} do both, and {neither} do neither. How many people " f"are in the group?") return question, only_a + only_b + both + neither def logic_calendar(rng: random.Random) -> tuple[str, int]: start = rng.randrange(7) ahead = rng.randint(9, 400) question = (f"A task starts on a {DAYS[start]}. Counting the start day as day 1, day {ahead} " f"falls on which day of the week? Answer with the position of that day in the " f"week where Monday is 1 and Sunday is 7.") return question, (start + ahead - 1) % 7 + 1 def logic_ordering(rng: random.Random) -> tuple[str, int]: people = rng.sample(NAMES, 5) positions = list(range(1, 6)) rng.shuffle(positions) placed = dict(zip(people, positions)) front = min(placed, key=placed.get) back = max(placed, key=placed.get) subject = rng.choice([p for p in people if p not in (front, back)]) clues = [] for person in people: if person == subject: continue ahead = placed[person] < placed[subject] clues.append(f"{person} is {'ahead of' if ahead else 'behind'} {subject}") rng.shuffle(clues) question = ("Five people stand in a queue, numbered 1 at the front to 5 at the back. " + "; ".join(clues) + f". What position is {subject} in?") return question, placed[subject] def logic_handshakes(rng: random.Random) -> tuple[str, int]: people = rng.randint(6, 20) absent = rng.randint(1, 4) present = people - absent question = (f"{people} people are invited to a meeting and {absent} do not come. Everyone " f"who comes shakes hands with everyone else exactly once. How many handshakes " f"take place?") return question, present * (present - 1) // 2 def logic_comparison_chain(rng: random.Random) -> tuple[str, int]: base = rng.randint(10, 60) step_one = rng.randint(2, 20) step_two = rng.randint(2, 20) a, b, c = rng.sample(NAMES, 3) question = (f"{a} is {step_one} years older than {b}, and {b} is {step_two} years younger " f"than {c}. {c} is {base} years old. How old is {a}?") return question, base - step_two + step_one def logic_parity(rng: random.Random) -> tuple[str, int]: total = rng.randint(20, 120) step = rng.choice([3, 4, 5, 6, 7]) question = (f"The whole numbers from 1 to {total} are written down. How many of them are " f"multiples of {step}?") return question, total // step LOGIC_TEMPLATES: list[Callable[[random.Random], tuple[str, int]]] = [ logic_inclusion_exclusion, logic_calendar, logic_ordering, logic_handshakes, logic_comparison_chain, logic_parity, ] FAMILIES = {"arith": ARITH_TEMPLATES, "logic": LOGIC_TEMPLATES} # -------------------------------------------------------------------------------------- # Assembly # -------------------------------------------------------------------------------------- def make_row(index: int, family: str, question: str, answer: int) -> dict: return { "id": f"{family}-{index:04d}", "family": family, "prompt": [{"role": "user", "content": f"{question}\n\n{INSTRUCTION}"}], "question": question, "answer": str(answer), } def generate(family: str, count: int, rng: random.Random, seen: set[str], start_index: int) -> list[dict]: """Distinct problems from the family's templates, cycling so each is used equally.""" templates = FAMILIES[family] rows: list[dict] = [] attempts = 0 while len(rows) < count: attempts += 1 if attempts > count * 200: raise SystemExit( f"could not generate {count} distinct {family} problems; the templates repeat " f"themselves before that. Ask for fewer, or add a template." ) template = templates[len(rows) % len(templates)] question, answer = template(rng) if question in seen: continue seen.add(question) rows.append(make_row(start_index + len(rows), family, question, answer)) return rows def from_gsm8k(count: int, split: str, offset: int) -> list[dict]: """A slice of GSM8K, whose card gives the licence as MIT and marks answers with ####.""" try: from datasets import load_dataset # noqa: PLC0415 - optional dependency except ImportError: raise SystemExit("--source gsm8k needs the datasets package: pip install datasets") from None data = load_dataset("openai/gsm8k", "main", split=split) rows = [] for i in range(offset, min(offset + count, len(data))): item = data[i] answer = item["answer"].split("####")[-1].strip().replace(",", "") rows.append(make_row(i, "gsm8k", item["question"].strip(), int(answer))) return rows def write_jsonl(path: Path, rows: list[dict]) -> str: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False) + "\n") return hashlib.sha256(path.read_bytes()).hexdigest() def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--out-dir", default="tasks") parser.add_argument("--train", type=int, default=240, help="training problems, same family") parser.add_argument("--heldout", type=int, default=60, help="held-out problems, per family; one set of the same kind and one of a different kind") parser.add_argument("--family", default="arith", choices=sorted(FAMILIES), help="the family the lab trains on") parser.add_argument("--other-family", default="logic", choices=sorted(FAMILIES), help="the different-kind family the reality check evaluates on") parser.add_argument("--source", default="generated", choices=["generated", "gsm8k"], help="generated problems, or a slice of GSM8K (MIT licence, per its dataset card)") parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() if args.family == args.other_family: raise SystemExit("--family and --other-family must differ; the point is a different kind of task") out = Path(args.out_dir) rng = random.Random(args.seed) seen: set[str] = set() manifest: dict[str, object] = {"seed": args.seed, "source": args.source, "files": {}} if args.source == "gsm8k": train = from_gsm8k(args.train, "train", 0) same = from_gsm8k(args.heldout, "test", 0) print(f"GSM8K, licence MIT per the dataset card; {len(train)} train, {len(same)} held-out") else: train = generate(args.family, args.train, rng, seen, 0) same = generate(args.family, args.heldout, rng, seen, args.train) different = generate(args.other_family, args.heldout, rng, seen, 0) for name, rows in (("train", train), ("heldout-same", same), ("heldout-different", different)): digest = write_jsonl(out / f"{name}.jsonl", rows) manifest["files"][name] = {"rows": len(rows), "sha256": digest} print(f"{name:18s} {len(rows):5d} rows sha256 {digest[:12]}...") (out / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") print(f"\nwritten to {out}/ (train, heldout-same, heldout-different, manifest.json)") print("Read ten of them before you train on them; a generator with a bug produces a " "reward function that is confidently wrong.") sample = train[0] print(f"\nexample problem ({sample['family']}):\n {sample['question']}\n reference answer: {sample['answer']}") if __name__ == "__main__": main()