data_spell.py

"""spell — a code-generated supervised task: a misspelled word -> the correct word.

The honest nano-task test: correcting a typo needs a learned lexicon AND a sense of
which real word a misspelling most likely meant. The realistic script baseline is
"nearest dictionary word by edit distance" — but it has no frequency prior, so when a
typo lands as close to a wrong word as to the right one (`bcak` -> back? buck? beck?)
it guesses, and it breaks on transpositions and double-letters. The model packs a
frequency-weighted vocabulary into ~1M params and ranks the correction.

Each example is one line::

    recieve => receive
    teh => the
    definately => definitely

Generation is answer-first and correct by construction: sample a real word from a
fixed, frequency-ordered vocabulary (Zipf-weighted so common words dominate), then
apply 1-2 realistic typo edits (delete, insert, substitute a keyboard neighbour,
transpose adjacent, double a letter). Because we start from the real word, the label
is ground truth. ~15% of examples are the identity (already-correct word) so the
model learns to leave correct words alone. The prompt is everything up to and
including ``" => "``; the target is the correct word. Byte-level vocab (256).
``spell_pairs`` is imported by both the dataset and the eval so train and test share
one data path. The vocabulary is the same frequency-ordered ~480-word list used
across the nano-* lexicon models.
"""

from __future__ import annotations

import numpy as np

_VOWELS = frozenset("aeiou")

# A fixed, roughly frequency-ordered vocabulary of common English words (most
# common first). This list IS the lexicon prior the model learns; earlier words
# are sampled more often (Zipf), so corrections resolve toward the common word.
_RAW = """
the be to of and a in that have it for not on with he as you do at this but his by
from they we say her she or an will my one all would there their what so up out if
about who get which go me when make can like time no just him know take people into
year your good some could them see other than then now look only come its over think
also back after use two how our work first well way even new want because any these
give day most us man find here thing tell very big small great little own under last
right move place such again off play move where turn need start show point try kind
hand high old life feel three world school still state never become between high
really something most family leave word both water side without head house seem
end ask group never word turn problem child city week company system program question
during number course company point case fact night area money story result job book
word eye door health person art war history party result change morning reason
research girl guy moment air teacher force education foot boy age policy music
market sense nation plan college interest death course experience effort water late
example heart paper space ground form event official matter center couple site
project activity star table court produce stock paint deep base camp signal critic
gold storm phone north union deal author cover plant flag worth wide class
green field add light cross hard early hold listen body keep watch human seven send
build stay fall reach kill remain suggest raise pass sell require report decide
pull return explain hope develop carry break receive agree support hit produce
eat cover catch draw choose cause point talk lead serve die remember love
consider appear buy wait serve die send expect build stay fall cut reach
read spend grow open walk win offer remember love consider appear wear
green blue red black white brown gray pink quick brave clean clear close
deep dry fair fine flat fresh full glad hot huge late loud nice plain proud
rare rich ripe rough safe sharp short slow soft sour sweet tall thick thin tight
warm weak wet wild wise young above across along among around behind below beneath
beside beyond inside outside toward within against during except toward
animal apple bird boat bread brother bridge button camera candle carpet castle
cattle circle cloud coast color corner cotton cousin danger dinner doctor dollar
dragon dream eagle earth engine farmer feather finger flower forest friend garden
ghost glass grass guard heaven honey island jacket jungle ladder leather letter
lion magic market mirror monkey mother mountain needle nephew object ocean office
orange palace parent pencil pepper picture pillow planet pocket potato prince queen
rabbit rather record river rocket sailor school season silver singer sister snake
spider spring statue summer sunset symbol table teacher temple thread throat thunder
ticket tiger toast tower travel turtle uncle valley village voice wagon weather
window winter wizard wonder yellow zebra ocean orange forest planet flower garden
""".split()

# rough QWERTY left/right neighbours for substitution typos
_NEIGHBORS = {
    "a": "sq", "b": "vn", "c": "xv", "d": "sf", "e": "wr", "f": "dg", "g": "fh",
    "h": "gj", "i": "uo", "j": "hk", "k": "jl", "l": "k", "m": "n", "n": "bm",
    "o": "ip", "p": "o", "q": "wa", "r": "et", "s": "ad", "t": "ry", "u": "yi",
    "v": "cb", "w": "qe", "x": "zc", "y": "tu", "z": "x",
}


def _build_vocab() -> list[str]:
    """De-duplicate `_RAW` keeping first (most-frequent) occurrence, drop words with
    no consonant skeleton (matches the shared nano-* lexicon build exactly, so the
    Zipf indices line up with the training data). Order = rough frequency rank."""
    seen: set[str] = set()
    vocab: list[str] = []
    for w in _RAW:
        if w in seen:
            continue
        seen.add(w)
        if len("".join(c for c in w if c not in _VOWELS)) >= 1:  # need a consonant
            vocab.append(w)
    return vocab


_WORDS = _build_vocab()
_RANKS = np.arange(1, len(_WORDS) + 1, dtype=np.float64)
_WEIGHTS = 1.0 / (_RANKS + 5.0)
_WEIGHTS /= _WEIGHTS.sum()


def _typo(rng, w: str) -> str:
    """Apply one realistic single-character typo to `w`."""
    if len(w) < 2:
        return w
    kind = rng.integers(5)
    i = int(rng.integers(len(w)))
    if kind == 0:  # delete
        return w[:i] + w[i + 1:]
    if kind == 1:  # insert a random letter
        c = chr(int(rng.integers(26)) + 97)
        return w[:i] + c + w[i:]
    if kind == 2:  # substitute a keyboard neighbour
        opts = _NEIGHBORS.get(w[i], "")
        if not opts:
            return w[:i] + chr(int(rng.integers(26)) + 97) + w[i + 1:]
        return w[:i] + opts[int(rng.integers(len(opts)))] + w[i + 1:]
    if kind == 3 and i < len(w) - 1:  # transpose adjacent
        return w[:i] + w[i + 1] + w[i] + w[i + 2:]
    # double a letter
    return w[:i] + w[i] + w[i:]


def spell_pairs(seed: int, n: int) -> list[tuple[str, str]]:
    """`n` deterministic (prompt, correct-word) pairs from `seed`.

    The prompt is ``<misspelling> => `` and the target is the correct word. ~15% are
    the identity (already-correct) so the model learns to leave good words alone.
    """
    rng = np.random.default_rng(seed)
    idx = rng.choice(len(_WORDS), size=n, p=_WEIGHTS)
    out: list[tuple[str, str]] = []
    for i in idx:
        word = _WORDS[int(i)]
        if rng.random() < 0.15:
            typo = word  # already correct
        else:
            typo = _typo(rng, word)
            if rng.random() < 0.3:  # sometimes a second edit
                typo = _typo(rng, typo)
        out.append((f"{typo} => ", word))
    return out