spb/wpai0001_uqo

1"""Micro-world C — MORPH: context-dependent suffix selection on pseudo-word stems.23``<bos> c_1 ... c_m = suffix`` where the suffix class depends on the stem-final phoneme class4(vowel / nasal / other consonant). Stems are random pseudo-words so lexical priors do not apply.5"""6from __future__ import annotations78import random910from wp_ai_0001.config import MorphTaskConfig11from wp_ai_0001.datasets.common import BOS, EQ, Example, Split, Vocab1213GENERATOR_VERSION = "morph-1.0"14VOWELS = list("aeiou")15NASALS = list("mn")16OTHERS = list("bdgklprstv")17LETTERS = VOWELS + NASALS + OTHERS181920def phon_class(ch: str) -> str:21    if ch in VOWELS:22        return "vowel"23    if ch in NASALS:24        return "nasal"25    return "consonant"262728def vocab_for(cfg: MorphTaskConfig) -> Vocab:29    return Vocab(LETTERS + list(cfg.rules.values()))303132def generate(cfg: MorphTaskConfig, n: int, seed: int, name: str, *, exclude_hashes: set[str] | None = None) -> Split:33    rng = random.Random(f"{GENERATOR_VERSION}|{name}|{seed}")34    seen = set(exclude_hashes or set())35    out: list[Example] = []36    attempts = 037    while len(out) < n:38        attempts += 139        if attempts > 200 * n + 10_000:40            raise RuntimeError(f"could not generate {n} unique MORPH examples for {name}")41        m = rng.randint(cfg.min_stem, cfg.max_stem)42        stem = [rng.choice(LETTERS) for _ in range(m)]43        ans = cfg.rules[phon_class(stem[-1])]44        ex = Example((BOS, *stem, EQ, ans), 1, cluster=f"{phon_class(stem[-1])}|len{m}", meta={"len": m})45        if ex.item_hash in seen:46            continue47        seen.add(ex.item_hash)48        out.append(ex)49    return Split(name, out, "morph", GENERATOR_VERSION, seed, cfg.model_dump())505152def validate(split: Split, cfg: MorphTaskConfig) -> None:53    for e in split.examples:54        stem = list(e.prompt[1:-1])55        assert cfg.rules[phon_class(stem[-1])] == e.answer, e56