spb/wpai0001_uqo

1"""Micro-world B — FSA: deterministic finite automaton execution over an invented alphabet.23``<bos> s_1 ... s_n = q_1 ... q_n`` (state after each symbol; the final answer is ``q_n``). The automaton is drawn once from ``automaton_seed``; the4transition table is part of the split config so the task is fully verifiable.5"""6from __future__ import annotations78import random910from wp_ai_0001.config import FSATaskConfig11from wp_ai_0001.datasets.common import BOS, EQ, Example, Split, Vocab1213GENERATOR_VERSION = "fsa-1.0"141516def transition_table(cfg: FSATaskConfig) -> dict[tuple[int, str], int]:17    rng = random.Random(f"{GENERATOR_VERSION}|automaton|{cfg.automaton_seed}")18    return {(q, a): rng.randrange(cfg.n_states) for q in range(cfg.n_states) for a in cfg.alphabet}192021def run(table: dict[tuple[int, str], int], symbols: list[str]) -> list[int]:22    q, out = 0, []23    for s in symbols:24        q = table[(q, s)]25        out.append(q)26    return out272829def vocab_for(cfg: FSATaskConfig) -> Vocab:30    return Vocab(list(cfg.alphabet) + [f"q{i}" for i in range(cfg.n_states)])313233def generate(cfg: FSATaskConfig, n: int, seed: int, name: str, *, min_len: int | None = None, max_len: int | None = None,34             exclude_hashes: set[str] | None = None) -> Split:35    rng = random.Random(f"{GENERATOR_VERSION}|{name}|{seed}")36    table = transition_table(cfg)37    lo = cfg.min_len if min_len is None else min_len38    hi = cfg.max_len if max_len is None else max_len39    seen = set(exclude_hashes or set())40    out: list[Example] = []41    attempts = 042    while len(out) < n:43        attempts += 144        if attempts > 200 * n + 10_000:45            raise RuntimeError(f"could not generate {n} unique FSA examples for {name}")46        k = rng.randint(lo, hi)47        syms = [rng.choice(cfg.alphabet) for _ in range(k)]48        states = [f"q{q}" for q in run(table, syms)]49        ex = Example((BOS, *syms, EQ, *states), n_answer=k, cluster=f"len{k}", meta={"len": k})50        if ex.item_hash in seen:51            continue52        seen.add(ex.item_hash)53        out.append(ex)54    conf = cfg.model_dump() | {"table": {f"{q},{a}": v for (q, a), v in table.items()}}55    return Split(name, out, "fsa", GENERATOR_VERSION, seed, conf)565758def validate(split: Split, cfg: FSATaskConfig) -> None:59    table = transition_table(cfg)60    for e in split.examples:61        syms = list(e.prompt[1:-1])62        assert tuple(f"q{q}" for q in run(table, syms)) == e.answers, e63