spb/wpai0001_uqo

1"""Micro-world A — OPSEQ: compositional operator sequences over Z_p.23A program is ``<bos> v0 op_1 ... op_k = r_1 ... r_k`` where each invented operator is a fixed map on4Z_p and ``r_j`` is the value after applying ``op_j`` (scratchpad execution trace; the final answer is5``r_k``). Operator names are meaningless strings so that no pre-training prior applies.6Donor competence / recipient weakness is created by *excluding* chosen operators from a model's7training distribution (``excluded_operators``). Compositional generalization is measured on8adjacent operator pairs that never co-occur in training (``held_out_pairs``), and length9generalization on programs longer than ``max_len``.1011Each split is generated from its own seed; disjointness of item hashes across splits is enforced12by the caller (``datasets.common.assert_disjoint``).13"""14from __future__ import annotations1516import random17from typing import Callable1819from wp_ai_0001.config import OpSeqTaskConfig20from wp_ai_0001.datasets.common import BOS, EQ, Example, Split, Vocab2122GENERATOR_VERSION = "opseq-1.0"2324# Invented operator semantics (affine maps x -> (a*x + b) mod p, or reflections).25OPERATOR_TABLE: dict[str, Callable[[int, int], int]] = {26    "zib": lambda x, p: (x + 1) % p,27    "zob": lambda x, p: (x + 3) % p,28    "qua": lambda x, p: (2 * x + 1) % p,29    "lar": lambda x, p: (3 * x) % p,30    "neg": lambda x, p: (-x) % p,31    "flip": lambda x, p: (p - 1 - x) % p,32    "dro": lambda x, p: (5 * x) % p,33    "mik": lambda x, p: (x * x) % p,34    "vex": lambda x, p: (4 * x + 3) % p,  # reserved for the unrelated-skill donor; absent from default configs35}363738def trace(v0: int, ops: list[str], p: int) -> list[int]:39    x, out = v0, []40    for op in ops:41        x = OPERATOR_TABLE[op](x, p)42        out.append(x)43    return out444546def apply_program(v0: int, ops: list[str], p: int) -> int:47    return trace(v0, ops, p)[-1] if ops else v0484950def vocab_for(cfg: OpSeqTaskConfig) -> Vocab:51    """Vocabulary over ALL known operators (fixed order) so that models trained on different operator52    subsets share token ids and output heads; ``cfg.operators`` only restricts generation."""53    return Vocab([str(i) for i in range(cfg.modulus)] + list(OPERATOR_TABLE))545556def _has_held_out_pair(ops: list[str], pairs: set[tuple[str, str]]) -> bool:57    return any((a, b) in pairs for a, b in zip(ops, ops[1:]))585960def visible_trace(v0: int, ops: list[str], p: int, stride: int = 1) -> list[int]:61    full = trace(v0, ops, p)62    idx = list(range(stride - 1, len(full), stride))63    if not idx or idx[-1] != len(full) - 1:64        idx.append(len(full) - 1)65    return [full[i] for i in idx]666768def make_example(v0: int, ops: list[str], p: int, stride: int = 1) -> Example:69    tr = [str(r) for r in visible_trace(v0, ops, p, stride)]70    tokens = (BOS, str(v0), *ops, EQ, *tr)71    return Example(tokens=tokens, n_answer=len(tr), cluster=" ".join(ops), meta={"len": len(ops), "v0": v0, "stride": stride})727374def generate(75    cfg: OpSeqTaskConfig,76    n: int,77    seed: int,78    name: str,79    *,80    mode: str = "train",81    require_operator: str | None = None,82    forbid_operator: str | None = None,83    min_len: int | None = None,84    max_len: int | None = None,85    exclude_hashes: set[str] | None = None,86) -> Split:87    """Generate ``n`` unique examples.8889    mode="train": no held-out pairs, no excluded operators.90    mode="comp":  every program contains at least one held-out adjacent pair.91    mode="any":   unconstrained with respect to held-out pairs.92    ``require_operator`` / ``forbid_operator`` restrict the operator content (target vs control).93    """94    rng = random.Random(f"{GENERATOR_VERSION}|{name}|{seed}")95    p = cfg.modulus96    pairs = set(map(tuple, cfg.held_out_pairs))97    allowed = [o for o in cfg.operators if o not in cfg.excluded_operators]98    if forbid_operator:99        allowed = [o for o in allowed if o != forbid_operator]100    lo = cfg.min_len if min_len is None else min_len101    hi = cfg.max_len if max_len is None else max_len102    seen: set[str] = set(exclude_hashes or set())103    out: list[Example] = []104    attempts = 0105    while len(out) < n:106        attempts += 1107        if attempts > 200 * n + 10_000:108            raise RuntimeError(f"could not generate {n} unique examples for {name} (got {len(out)})")109        k = rng.randint(lo, hi)110        ops = [rng.choice(allowed) for _ in range(k)]111        if require_operator and require_operator not in ops:112            continue113        held = _has_held_out_pair(ops, pairs)114        if mode == "train" and held:115            continue116        if mode == "comp" and not held:117            continue118        v0 = rng.randrange(p)119        ex = make_example(v0, ops, p, cfg.trace_stride)120        if ex.item_hash in seen:121            continue122        seen.add(ex.item_hash)123        out.append(ex)124    return Split(name, out, "opseq", GENERATOR_VERSION, seed, cfg.model_dump())125126127def validate(split: Split, cfg: OpSeqTaskConfig) -> None:128    """Independent recomputation of every answer (exact validator)."""129    for e in split.examples:130        assert e.tokens[0] == BOS and e.prompt[-1] == EQ131        v0 = int(e.tokens[1])132        ops = list(e.prompt[2:-1])133        assert all(o in OPERATOR_TABLE for o in ops), e134        assert tuple(str(r) for r in visible_trace(v0, ops, cfg.modulus, cfg.trace_stride)) == e.answers, e135