#!/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 (``) becomes ``. - **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 ` ` (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 = "", "", "", "" 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 ); 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 / , so # char-split ONLY the unknown part's surface substring and KEEP known morphemes (e.g. an # unknown stem with a real -ing suffix -> [, 'ing'], not one opaque ). pieces = [] for m in morphs: if m.startswith(" 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 () — that would create a # second copy of the special (`▁`) 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 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" tokens: {unk} ({100 * unk / total:.3f}%)") print(f"wrote {out.relative_to(REPO)}") if __name__ == "__main__": _main()