our_tokenizer.py

#!/usr/bin/env python3
"""Our morphological tokenizer, wrapped for real use (training + eval + intrinsic comparison).

The project has an ANALYZER (`scripts/morph_tokenizer.analyze`: one WORD -> morpheme sequence, MorphyNet +
curated supplements). This module turns it into a TOKENIZER over raw text:

  - **Pre-tokenization** (fast regex, GPT-2 style; NOT nltk/spaCy): split a line into `word | number |
    punctuation` units, tracking whether each unit was preceded by whitespace.
  - **Morphemes:** each word/number is analysed into morphemes; each punctuation mark is its own token.
    Anything the analyzer can't handle (`<UNK...>`) becomes `<unk>`.
  - **Word boundary:** fused `▁` on the WORD-INITIAL morpheme of a space-preceded unit (SentencePiece style,
    option A in docs/tokenizer-design-decisions.md §1; user-approved 2026-07-18). `▁` == "preceded by a space".
  - **Vocab:** every distinct emitted token over the 10M corpus + specials `<pad> <unk> <bos> <eos>`
    (ids 0..3; id 0 = pad so it is safe to be the dropped id in the SimpleStories harness).

Design notes: the corpus is normalised (curly->ASCII apostrophes, lowercased) exactly as the analyzer expects.
Per-unit results are memoised, so tokenizing 10M words costs ~one analyze() per unique word type.

Usage:
    python our_tokenizer.py build   # build vocab from data/strict-small -> our_vocab.json (+ stats)
    from our_tokenizer import OurTokenizer
    tok = OurTokenizer.load("comparisons/common/our_vocab.json")
    tok.tokenize("The smarter birds flew.")   # ['▁the','▁smart','er','▁bird','s','▁fly', ... , '.']
"""
from __future__ import annotations

import json
import re
import sys
from collections import Counter
from pathlib import Path


# HF-shipped variant: morph_tokenizer.py is a sibling module; resources resolve from the
# vocab file directory (the model snapshot root), NOT a repo checkout.
REPO = Path(__file__).parent
try:
    from .morph_tokenizer import load_resources as _hf_copy_probe  # noqa: F401  (forces HF dynamic-module file copy)
    from . import morph_tokenizer as mt  # noqa: E402
except ImportError:  # direct, non-package use (importlib evades HF check_imports regex)
    import importlib as _il
    mt = _il.import_module("morph_tokenizer")

BOUNDARY = "▁"  # ▁ , the SentencePiece word-boundary marker (fused onto the word-initial token)
PAD, UNK, BOS, EOS = "<pad>", "<unk>", "<bos>", "<eos>"
SPECIALS = [PAD, UNK, BOS, EOS]  # ids 0,1,2,3

# One unit = a word (letters, with internal '/-), OR a run of digits, OR a single non-word/non-space char
# (punctuation/symbol), OR a run of underscores. Whitespace is NOT captured (we infer boundaries from it).
PRETOK_RE = re.compile(r"[^\W\d_]+(?:['-][^\W\d_]+)*|\d+|[^\s\w]|_+")


class OurTokenizer:
    """Morphological tokenizer over raw text. `res` is the analyzer resource dict from load_resources()."""

    def __init__(self, res: dict, vocab: dict[str, int] | None = None, char_backoff: bool = True):
        self.res = res
        self._cache: dict[str, list[str]] = {}
        self.vocab = vocab
        self.char_backoff = char_backoff  # OOV token -> its characters (not a single <unk>); needs char tokens in vocab
        self.ids_to_tokens = {i: t for t, i in vocab.items()} if vocab else None

    # -- core: text -> morpheme tokens ------------------------------------------------------------
    def _pieces(self, unit: str) -> list[str]:
        """Morpheme pieces for one pre-token unit (word/number analysed; punctuation kept literal)."""
        cached = self._cache.get(unit)
        if cached is not None:
            return cached
        ch = unit[0]
        if ch.isalpha() or ch.isdigit():  # word or number -> analyze (numbers digit-split)
            morphs, _ = mt.analyze(unit, self.res)
            if self.char_backoff:
                # Per-morpheme backoff: the analyzer marks unknowns as <UNK:surface> / <UNK_cls:surface>, so
                # char-split ONLY the unknown part's surface substring and KEEP known morphemes (e.g. an
                # unknown stem with a real -ing suffix -> [<stem chars>, 'ing'], not one opaque <unk>).
                pieces = []
                for m in morphs:
                    if m.startswith("<UNK"):
                        pieces.extend(m[m.index(":") + 1:-1] if ":" in m else m)  # its chars
                    else:
                        pieces.append(m)
            else:
                pieces = [UNK if m.startswith("<UNK") else m for m in morphs]
        else:  # punctuation / symbol / underscore run -> literal token
            pieces = [unit]
        self._cache[unit] = pieces
        return pieces

    def tokenize_with_spans(self, text: str) -> list[tuple[str, tuple[int, int]]]:
        """Like tokenize() but pairs each token with the (start, end) CHAR span of its source word in `text`.
        Morphemes are analysis (not segmentation), so all pieces of a word share the word's span — word-
        granularity offsets, which is what the eval harness needs to locate the whitespace-bounded completion.
        (Normalisation is length-preserving for ASCII, so norm spans == original spans for English text.)"""
        norm = text.translate(mt.CURLY).lower()
        out: list[tuple[str, tuple[int, int]]] = []
        for m in PRETOK_RE.finditer(norm):
            start, end = m.start(), m.end()
            space_preceded = start == 0 or norm[start - 1].isspace()
            pieces = self._pieces(m.group())
            if not pieces:
                continue
            # fuse ▁ on the word-initial piece, but NEVER on a special (<unk>) — that would create a
            # second copy of the special (`▁<unk>`) and fragment its embedding.
            first = pieces[0]
            if space_preceded and first not in SPECIALS:
                first = BOUNDARY + first
            out.append((first, (start, end)))
            out.extend((p, (start, end)) for p in pieces[1:])
        return out

    def tokenize(self, text: str) -> list[str]:
        """Raw text -> list of tokens (fused ▁ on space-preceded word-initial pieces)."""
        return [t for t, _ in self.tokenize_with_spans(text)]

    # -- ids ---------------------------------------------------------------------------------------
    @staticmethod
    def char_pieces(token: str) -> list[str]:
        """Split an OOV token into character tokens, keeping a leading ▁ fused onto the first character.
        e.g. '▁wug' -> ['▁w','u','g']; 'wug' -> ['w','u','g']. Used for char backoff (and vocab coverage)."""
        if token.startswith(BOUNDARY):
            body = token[len(BOUNDARY):]
            return [BOUNDARY + body[0], *body[1:]] if body else [token]
        return list(token)

    def encode_with_spans(self, text: str) -> tuple[list[int], list[tuple[int, int]]]:
        """Token ids AND an aligned (start, end) source-word char span per id (word-granularity; char-backoff
        pieces inherit their word's span). Used by the HF wrapper to emit offset_mapping."""
        assert self.vocab is not None, "no vocab loaded"
        unk = self.vocab[UNK]
        ids: list[int] = []
        spans: list[tuple[int, int]] = []
        for t, span in self.tokenize_with_spans(text):
            tid = self.vocab.get(t)
            if tid is not None:
                ids.append(tid); spans.append(span)
            elif self.char_backoff:  # OOV -> characters (each char->id, or <unk> if that char is unseen)
                for c in self.char_pieces(t):
                    ids.append(self.vocab.get(c, unk)); spans.append(span)
            else:
                ids.append(unk); spans.append(span)
        return ids, spans

    def encode(self, text: str, add_special: bool = False) -> list[int]:
        ids = self.encode_with_spans(text)[0]
        if add_special:
            ids = [self.vocab[BOS], *ids, self.vocab[EOS]]
        return ids

    def decode(self, ids: list[int]) -> str:
        """Best-effort surface reconstruction (realization rules are deferred; fine for debug/sanity)."""
        assert self.ids_to_tokens is not None
        parts: list[str] = []
        for i in ids:
            t = self.ids_to_tokens.get(i, UNK)
            if t in (PAD, BOS, EOS):
                continue
            parts.append(" " + t[len(BOUNDARY):] if t.startswith(BOUNDARY) else t)
        return "".join(parts).strip()

    @property
    def vocab_size(self) -> int:
        assert self.vocab is not None
        return len(self.vocab)

    # -- persistence -------------------------------------------------------------------------------
    def save(self, path: str | Path) -> None:
        assert self.vocab is not None
        Path(path).write_text(json.dumps(self.vocab, ensure_ascii=False, indent=0), encoding="utf-8")

    @classmethod
    def load(cls, vocab_path: str | Path, res: dict | None = None) -> OurTokenizer:
        res = res if res is not None else mt.load_resources(Path(vocab_path).resolve().parent)
        vocab = json.loads(Path(vocab_path).read_text(encoding="utf-8"))
        return cls(res, vocab)


def build_vocab(tok: OurTokenizer, corpus_dir: str | Path, min_count: int = 1) -> tuple[dict[str, int], Counter]:
    """Tokenize every *.train.txt line and build the vocab (specials first, then by descending frequency)."""
    counts: Counter = Counter()
    for fp in sorted(Path(corpus_dir).glob("*.train.txt")):
        for line in mt._lines(fp):
            counts.update(tok.tokenize(line))
    vocab = {s: i for i, s in enumerate(SPECIALS)}
    for tokn, cnt in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0])):
        if tokn in vocab or cnt < min_count:
            continue
        vocab[tokn] = len(vocab)
    return vocab, counts


def _main() -> None:
    corpus = REPO / "data" / "strict-small"
    print(f"loading analyzer resources ...", file=sys.stderr)
    res = mt.load_resources(REPO)
    tok = OurTokenizer(res)
    print(f"building vocab over {corpus} ...", file=sys.stderr)
    vocab, counts = build_vocab(tok, corpus)
    tok.vocab = vocab
    out = Path(__file__).parent / "our_vocab.json"
    tok.save(out)
    total = sum(counts.values())
    unk = counts.get(UNK, 0)
    boundary_toks = sum(1 for t in vocab if t.startswith(BOUNDARY))
    print(f"vocab size (with specials):     {len(vocab)}")
    print(f"distinct tokens over corpus:    {len(counts)}")
    print(f"  word-initial (▁) tokens:      {boundary_toks}")
    print(f"total tokens emitted:           {total}")
    print(f"<unk> tokens:                   {unk}  ({100 * unk / total:.3f}%)")
    print(f"wrote {out.relative_to(REPO)}")


if __name__ == "__main__":
    _main()