SPEC.md

SPEC.md — granule

This is the implementation contract. An agent working on this repo should treat it as authoritative and update it (with a note) rather than silently diverging.

Goal

A byte-level BPE tokenizer implemented from first principles, plus a benchmark harness comparing it to production tokenizers. Correctness and measurement quality matter more than raw speed.

Non-goals

• Beating tiktoken on encode throughput in pure Python.
• Training a language model. This repo stops at the tokenizer.
• Supporting every tokenizer family. Unigram/WordPiece are explicitly out of scope for v1.

Public API

class Tokenizer:
    @classmethod
    def train(cls, corpus: str | Path | Iterable[str], vocab_size: int,
              pattern: Literal["gpt2", "gpt4", "none"] = "gpt4",
              special_tokens: list[str] | None = None,
              fast: bool = True,
              verbose: bool = False) -> "Tokenizer": ...

    def encode(self, text: str, allowed_special: set[str] | None = None) -> list[int]: ...
    def decode(self, ids: list[int]) -> str: ...
    def save(self, path: str | Path) -> None: ...

    @classmethod
    def load(cls, path: str | Path) -> "Tokenizer": ...

    @property
    def vocab_size(self) -> int: ...

Invariants (enforce with tests)

1. decode(encode(s)) == s for all inputs, including emoji, mixed scripts, lone surrogates escaped as bytes, and control characters. This is the property test; run it over Hypothesis-generated text plus a fixture file of adversarial strings.
2. encode is deterministic across processes.
3. save → load → encode produces byte-identical output to the pre-save tokenizer.
4. Merges are applied in training order; ties in pair frequency are broken deterministically (lowest pair id first — document the choice).
5. Special tokens are matched only when explicitly allowed; otherwise their literal text is encoded as ordinary bytes.

Model file format

Human-inspectable, versioned, single file. Suggested: a header line with format version, pattern name, and vocab size, followed by the merge list in application order, followed by the special-token table. Include a .vocab sidecar writer for debugging that renders each token as printable text with unprintables escaped — this is for humans, never re-read by the loader.

Algorithms

Naive trainer (implement first)

Corpus → pre-tokenize → for each pre-token, a list of byte ids. Loop: count all adjacent pairs across the corpus, take the argmax, merge it everywhere, record it. Cost is O(N) per merge. This is the reference implementation and stays in the repo as a correctness oracle.

Fast trainer (implement second, must match naive exactly)

1. Pre-tokenize, then collapse to unique pre-tokens with occurrence counts. Most corpora shrink by 10–100x here.
2. Maintain pair_counts: dict[pair, int] and pair_locations: dict[pair, set[word_idx]].
3. Use a max-heap of (-count, pair) with lazy deletion — pop, verify the count still matches pair_counts, discard if stale.
4. On merging pair p: for each word containing p, rewrite the symbol sequence and apply delta updates to the counts of only the neighbouring pairs that changed. Never rescan the corpus.

A test must assert that naive and fast produce identical merge lists on a fixed corpus with a fixed vocab size. That test is the whole point of having both.

Pre-tokenization patterns

Do not transcribe split regexes from memory. Fetch the canonical gpt2 and cl100k_base patterns from tiktoken's published source, cite the source file in a code comment, and add unit tests covering: contractions, runs of digits, leading whitespace attachment, and newline handling.

Encoding

Split with the pattern, encode each pre-token independently, functools.lru_cache on pre-token → ids. Merges applied greedily in training order.

Benchmark harness

Corpora (src/granule/bench/corpora.py)

Each corpus is a small committed sample (≤5 MB) with a loader that records source, license, and SHA256. Required set: English web text, Tamil, Hindi, one more non-Latin script, Python source, and JSON/log lines. Document provenance in data/README.md. Do not commit anything license-ambiguous.

Metrics (src/granule/bench/metrics.py)

• chars/token and bytes/token — primary compression metrics, reported per corpus.
• Fertility — tokens per whitespace-delimited word. Note in the docs that this is misleading for scripts without whitespace word boundaries; report it anyway with the caveat.
• Encode throughput — MB/s, median of N≥5 runs, warm cache excluded, single core, with CPU model recorded.
• Train wall-clock and peak RSS, naive vs fast, across corpus sizes.
• Vocab overlap — fraction of our merges that also appear in the baseline's vocabulary. Interesting, cheap, rarely measured.

Baselines

tiktoken (cl100k_base, o200k_base) required. HuggingFace tokenizers optional behind an extra.

Reporting

granule bench --report writes results/results.json, renders Markdown tables, regenerates plots into results/, and splices the tables into README between the AUTOGEN:results-start/end markers. Committed results must be reproducible from a single command.

Quality bar

• Python 3.11+, ruff + mypy --strict clean on src/.
• pytest with coverage on src/granule ≥ 85%.
• GitHub Actions: lint, type-check, test on 3.11 and 3.12.
• Zero non-stdlib runtime dependencies for the core tokenizer. Benchmarking deps live under an optional extra.
• Every public function has a docstring stating what it does and what it assumes.

Milestones

ID	Deliverable	Done when
M0	Repo scaffold, pyproject.toml, CI, empty package importable	CI green on an empty test suite
M1	Naive trainer + encode/decode + save/load	Round-trip property tests pass
M2	Pre-tokenization patterns + special tokens	Pattern unit tests pass
M3	Fast trainer	Merge lists identical to naive; speedup measured
M4	Benchmark harness + corpora	granule bench runs end to end
M5	Report generation, plots, README autofill	results/ committed and reproducible
M6	Docs, demo GIF, limitations write-up	README readable by a stranger in 60 seconds