1"""Verifiable synthetic task generators with programmatic ground truth.23Families (each item is a dict: family, template, depth, prompt, answer, text, meta):4 arith multi-step arithmetic chains (depth = number of operations)5 bind variable binding + dependent substitution (depth = chain length)6 fsm register / finite-state updates (depth = number of updates)7 excl conjunctive constraints with distractors (depth = number of constraints)8 sched rule-based selection under a budget and exclusions (depth = number of candidates)9 compose held-out composition: bind chain feeding an arithmetic expression (depth = total steps)1011Splitting is done BEFORE sampling: by (family, template id, depth regime, seed stream). Train uses templates12T0-T2 and depths <= max_train_depth; held-out "template" items use templates T3-T4; held-out "depth" items use13depths > max_train_depth; the compose family is entirely held out from training unless explicitly requested.14Text format: "<prompt> Answer: <answer>" ; the answer is the only supervised span for exact-match evaluation.15"""16from __future__ import annotations1718import hashlib19import json20import random21from dataclasses import dataclass, asdict22from typing import Callable2324OPS = {"+": lambda a, b: a + b, "-": lambda a, b: a - b, "*": lambda a, b: a * b}2526# Phrasing templates: index 0..2 = train, 3..4 = held-out paraphrases.27TEMPLATES = {28 "arith": ["Compute {expr}.", "What is {expr}?", "Evaluate the expression {expr}.", "Calculate {expr} step by step and give the result.", "Find the value of {expr}."],29 "bind": ["Let {chain}. What is {var}?", "Given {chain}, compute {var}.", "{chain}. Find {var}.", "Suppose {chain}. Then {var} equals what?", "With {chain}, the value of {var} is?"],30 "fsm": ["Start with register {init}. {ops}. What is the register?", "The counter begins at {init}. {ops}. Final counter?", "Initial value {init}. {ops}. Result?", "A register holds {init}. Apply: {ops}. What does it hold now?", "Begin at {init}; {ops}; final value?"],31 "excl": ["Items: {items}. Keep the items that are {rules}. Which items remain?", "Consider {items}. Select those that are {rules}. List them.", "{items}. Filter: {rules}. Remaining items?", "Given the list {items}, keep only the ones {rules}. Answer with the names.", "From {items}, which satisfy: {rules}?"],32 "sched": ["Budget {budget}. Tasks: {tasks}. Rule: {rule}. Choose the affordable task with the highest value.", "You have {budget} units. Options: {tasks}. Constraint: {rule}. Which affordable option has the largest value?", "Capacity {budget}. Candidates: {tasks}. Exclusion: {rule}. Pick the best affordable candidate.", "With a budget of {budget}, among {tasks}, under the rule {rule}, which task maximises value while staying affordable?", "Budget is {budget}. {tasks}. {rule}. Best affordable task by value?"],33 "compose": ["Let {chain}. Compute {expr}.", "Given {chain}, what is {expr}?", "{chain}. Evaluate {expr}.", "Suppose {chain}. Then find {expr}.", "With {chain}, calculate {expr}."],34}353637def _rng(family: str, template: int, depth: int, seed: int, idx: int) -> random.Random:38 key = f"{family}|{template}|{depth}|{seed}|{idx}"39 return random.Random(int(hashlib.sha256(key.encode()).hexdigest()[:16], 16))404142def gen_arith(r: random.Random, depth: int, template: int):43 vals = [r.randint(1, 20)]44 ops = []45 expr = str(vals[0])46 cur = vals[0]47 for _ in range(depth):48 op = r.choice(["+", "-", "*"])49 b = r.randint(1, 9) if op == "*" else r.randint(1, 30)50 cur = OPS[op](cur, b) # left-to-right evaluation (stated below)51 expr += f" {op} {b}"52 ops.append(op)53 prompt = TEMPLATES["arith"][template].format(expr=expr) + " Evaluate left to right."54 return prompt, str(cur), {"ops": "".join(ops)}555657def gen_bind(r: random.Random, depth: int, template: int):58 names = r.sample(list("abcdefghjkmnpqrstuvwxyz"), depth + 1)59 vals = {names[0]: r.randint(1, 15)}60 parts = [f"{names[0]} = {vals[names[0]]}"]61 for i in range(1, depth + 1):62 op = r.choice(["+", "-", "*"])63 b = r.randint(1, 5) if op == "*" else r.randint(1, 12)64 src = names[i - 1]65 vals[names[i]] = OPS[op](vals[src], b)66 parts.append(f"{names[i]} = {src} {op} {b}")67 # distractor: an unused variable68 d = r.choice([n for n in "ABCDEFG"])69 parts.insert(r.randint(1, len(parts)), f"{d} = {r.randint(1, 40)}")70 chain = ", ".join(parts)71 target = names[depth]72 prompt = TEMPLATES["bind"][template].format(chain=chain, var=target)73 return prompt, str(vals[target]), {"target": target}747576def gen_fsm(r: random.Random, depth: int, template: int):77 init = r.randint(0, 10)78 cur = init79 steps = []80 for _ in range(depth):81 op = r.choice(["add", "sub", "double", "reset", "halve"])82 if op == "add":83 b = r.randint(1, 9); cur += b; steps.append(f"add {b}")84 elif op == "sub":85 b = r.randint(1, 9); cur -= b; steps.append(f"subtract {b}")86 elif op == "double":87 cur *= 2; steps.append("double it")88 elif op == "halve":89 cur = cur // 2; steps.append("halve it (round down)")90 else:91 b = r.randint(0, 9); cur = b; steps.append(f"reset to {b}")92 prompt = TEMPLATES["fsm"][template].format(init=init, ops=", then ".join(steps))93 return prompt, str(cur), {}949596def gen_excl(r: random.Random, depth: int, template: int):97 n = r.randint(4, 7)98 names = r.sample(["apple", "pear", "plum", "fig", "kiwi", "lime", "date", "yam", "nut", "bean"], n)99 items = {nm: (r.randint(1, 30), r.choice(["red", "green", "blue"])) for nm in names}100 rules = []101 def keep(v, c):102 return True103 preds: list[Callable] = []104 for _ in range(depth):105 kind = r.choice(["gt", "lt", "color", "even"])106 if kind == "gt":107 t = r.randint(5, 25); rules.append(f"greater than {t}"); preds.append(lambda v, c, t=t: v > t)108 elif kind == "lt":109 t = r.randint(5, 25); rules.append(f"smaller than {t}"); preds.append(lambda v, c, t=t: v < t)110 elif kind == "color":111 col = r.choice(["red", "green", "blue"]); rules.append(f"not {col}"); preds.append(lambda v, c, col=col: c != col)112 else:113 rules.append("even"); preds.append(lambda v, c: v % 2 == 0)114 remaining = [nm for nm in names if all(p(*items[nm]) for p in preds)]115 items_s = ", ".join(f"{nm} ({v}, {c})" for nm, (v, c) in items.items())116 prompt = TEMPLATES["excl"][template].format(items=items_s, rules=" and ".join(rules))117 ans = " ".join(remaining) if remaining else "none"118 return prompt, ans, {"n_items": n, "n_remaining": len(remaining)}119120121def gen_sched(r: random.Random, depth: int, template: int):122 names = [chr(ord("A") + i) for i in range(depth)]123 tasks = {nm: (r.randint(1, 12), r.randint(1, 20)) for nm in names}124 budget = r.randint(4, 12)125 kind = r.choice(["cost_gt", "value_lt", "none"])126 if kind == "cost_gt":127 t = r.randint(3, 10); rule = f"exclude tasks with cost greater than {t}"; ok = lambda c, v: c <= t128 elif kind == "value_lt":129 t = r.randint(3, 10); rule = f"exclude tasks with value below {t}"; ok = lambda c, v: v >= t130 else:131 rule = "no exclusion"; ok = lambda c, v: True132 feas = [(v, nm) for nm, (c, v) in tasks.items() if c <= budget and ok(c, v)]133 if not feas:134 ans = "none"135 else:136 best = max(v for v, _ in feas)137 ans = sorted(nm for v, nm in feas if v == best)[0]138 tasks_s = ", ".join(f"{nm} (cost {c}, value {v})" for nm, (c, v) in tasks.items())139 prompt = TEMPLATES["sched"][template].format(budget=budget, tasks=tasks_s, rule=rule)140 return prompt, ans, {"n_feasible": len(feas)}141142143def gen_compose(r: random.Random, depth: int, template: int):144 # bind chain of length depth-2 feeding a 2-op arithmetic expression over two chain variables145 k = max(1, depth - 2)146 names = r.sample(list("abcdefghjkmnpqrstuvwxyz"), k + 1)147 vals = {names[0]: r.randint(1, 10)}148 parts = [f"{names[0]} = {vals[names[0]]}"]149 for i in range(1, k + 1):150 op = r.choice(["+", "-"]); b = r.randint(1, 9); src = names[i - 1]151 vals[names[i]] = OPS[op](vals[src], b); parts.append(f"{names[i]} = {src} {op} {b}")152 v1, v2 = r.sample(names, 2) if len(names) > 1 else (names[0], names[0])153 op1, op2 = r.choice(["+", "-", "*"]), r.choice(["+", "-"])154 c = r.randint(1, 9)155 val = OPS[op2](OPS[op1](vals[v1], vals[v2]), c)156 expr = f"{v1} {op1} {v2} {op2} {c} (left to right)"157 prompt = TEMPLATES["compose"][template].format(chain=", ".join(parts), expr=expr)158 return prompt, str(val), {}159160161GENERATORS = {"arith": gen_arith, "bind": gen_bind, "fsm": gen_fsm, "excl": gen_excl, "sched": gen_sched, "compose": gen_compose}162163164def make_item(family: str, template: int, depth: int, seed: int, idx: int) -> dict:165 r = _rng(family, template, depth, seed, idx)166 prompt, answer, meta = GENERATORS[family](r, depth, template)167 text = f"{prompt} Answer: {answer}"168 return {"family": family, "template": template, "depth": depth, "seed": seed, "idx": idx,169 "prompt": prompt + " Answer:", "answer": " " + answer, "text": text, "meta": meta,170 "uid": hashlib.sha256(text.encode()).hexdigest()[:16]}171172173@dataclass174class SplitSpec:175 name: str176 families: tuple[str, ...]177 templates: tuple[int, ...]178 depths: tuple[int, ...]179 seed: int180 n_per_cell: int181182183def generate_split(spec: SplitSpec) -> list[dict]:184 items = []185 for fam in spec.families:186 for tpl in spec.templates:187 for dep in spec.depths:188 for i in range(spec.n_per_cell):189 items.append(make_item(fam, tpl, dep, spec.seed, i))190 return items191192193def default_splits(n_train_per_cell: int = 400, n_eval_per_cell: int = 60) -> dict[str, SplitSpec]:194 """Train: 5 families, templates 0-2, depths 1-4, seed 1. Dev: seed 2 same cells (IID generalisation).195 Held-out template: templates 3-4. Held-out depth: depths 5-6. Compose: fully held-out family.196 Final test: seed 9 across all eval regimes (LOCKED until preregistration is frozen)."""197 fams = ("arith", "bind", "fsm", "excl", "sched")198 return {199 "train": SplitSpec("train", fams, (0, 1, 2), (1, 2, 3, 4), 1, n_train_per_cell),200 "calib": SplitSpec("calib", fams, (0, 1, 2), (1, 2, 3, 4), 3, n_eval_per_cell // 2),201 "dev_iid": SplitSpec("dev_iid", fams, (0, 1, 2), (1, 2, 3, 4), 2, n_eval_per_cell),202 "dev_template": SplitSpec("dev_template", fams, (3, 4), (1, 2, 3, 4), 2, n_eval_per_cell),203 "dev_depth": SplitSpec("dev_depth", fams, (0, 1, 2), (5, 6), 2, n_eval_per_cell),204 "dev_compose": SplitSpec("dev_compose", ("compose",), (0, 1, 2), (3, 4, 5), 2, n_eval_per_cell),205 "test_iid": SplitSpec("test_iid", fams, (0, 1, 2), (1, 2, 3, 4), 9, n_eval_per_cell),206 "test_template": SplitSpec("test_template", fams, (3, 4), (1, 2, 3, 4), 9, n_eval_per_cell),207 "test_depth": SplitSpec("test_depth", fams, (0, 1, 2), (5, 6), 9, n_eval_per_cell),208 "test_compose": SplitSpec("test_compose", ("compose",), (0, 1, 2), (3, 4, 5), 9, n_eval_per_cell),209 }210211212def write_splits(out_dir: str, n_train_per_cell: int = 400, n_eval_per_cell: int = 60) -> dict:213 from pathlib import Path214 out = Path(out_dir); out.mkdir(parents=True, exist_ok=True)215 manifest = {}216 all_uids: dict[str, str] = {}217 for name, spec in default_splits(n_train_per_cell, n_eval_per_cell).items():218 items = generate_split(spec)219 # cross-split exact-duplicate removal: an item whose full text already appeared in an earlier split220 # (train is written first) is DROPPED from the later split; within-split duplicates are dropped too.221 kept, dups = [], 0222 for it in items:223 if it["uid"] in all_uids:224 dups += 1225 continue226 all_uids[it["uid"]] = name227 kept.append(it)228 items = kept229 p = out / f"{name}.jsonl"230 with open(p, "w") as f:231 for it in items:232 f.write(json.dumps(it) + "\n")233 manifest[name] = {"spec": asdict(spec), "n": len(items), "duplicates_dropped": dups,234 "sha256": hashlib.sha256(p.read_bytes()).hexdigest(), "locked": name.startswith("test")}235 (out / "manifest.json").write_text(json.dumps(manifest, indent=2))236 return manifest237