Add byte-level BPE tokenizer trained from scratch

Lossless on any string (256-byte base vocab), regex pre-tokenization so merges never
cross word boundaries, incremental pair-count training with deterministic tie-breaks,
special tokens, and JSON persistence. The trainer gains --tokenizer bpe and reports
bits/byte so tokenizers are comparable; old character checkpoints still load.

Tests check the incremental trainer against a naive reference, round-trips on random
unicode, and merge-priority semantics; mutation checks confirmed they fail on broken code.

README.md

@@ -2,11 +2,13 @@
 
 A GPT-style language model written from scratch in PyTorch, plus the inference techniques that
 matter in production LLM serving: a **KV cache**, **grouped-query attention**, **continuous
-batching** over variable-length requests, and **speculative decoding**. Every one of them is
+batching** over variable-length requests, **speculative decoding**, and a **byte-level BPE
-backed by tests that check correctness against a reference, not just "it runs".
+tokenizer**. Every one of them is backed by tests that check correctness against a reference,
+not just "it runs".
 
-No `nn.Transformer`, no Hugging Face. Attention, masking, caching, batching and sampling are all
+No `nn.Transformer`, no Hugging Face, no tokenizer library. Attention, masking, caching,
-written out explicitly (~900 lines of library code) so every line can be reasoned about.
+batching, sampling and tokenization are all written out explicitly (~1,100 lines of library code)
+so every line can be reasoned about.
 
 ## What's inside
 
@@ -17,6 +19,7 @@ written out explicitly (~900 lines of library code) so every line can be reasone
 | KV cache + sliding window | `model.py` | `O(T)` work per generated token instead of `O(T^2)` |
 | Continuous batching | `batching.py` | Left-padded shared cache, per-row RoPE positions, key mask, mid-flight admission, column trimming |
 | Speculative decoding | `speculative.py` | Exact target distribution (rejection sampling + residual), cache rollback |
+| Byte-level BPE tokenizer | `bpe.py` | Lossless on any string (256-byte base), incremental pair-count updates, deterministic ties, special tokens |
 | Sampling: temperature / top-k / top-p | `sampling.py` | `logits_to_probs` exposes the exact sampling distribution (greedy = one-hot) |
 | Trainer | `train.py` | AdamW with decoupled decay, warmup + cosine LR, grad clipping |
 
@@ -25,11 +28,13 @@ written out explicitly (~900 lines of library code) so every line can be reasone
 ```bash
 python3 -m venv .venv && source .venv/bin/activate
 pip install -e ".[dev]"
-pytest                                   # 49 tests, ~15s on CPU
+pytest                                   # 69 tests, ~15s on CPU
 
 mkdir -p data && curl -sSL -o data/shakespeare.txt \
   https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
 python -m microgpt.train --data data/shakespeare.txt --steps 1200 --batch-size 24
+python -m microgpt.train --data data/shakespeare.txt --tokenizer bpe --vocab-size 512 \
+    --steps 1200 --batch-size 24 --out checkpoints/bpe.pt   # BPE instead of characters
 python -m microgpt.train --data data/shakespeare.txt --n-layer 1 --n-embd 64 --n-head 2 \
     --steps 2000 --out checkpoints/draft.pt            # small draft model
 
@@ -40,6 +45,7 @@ python scripts/benchmark.py              # KV cache vs no cache
 python scripts/bench_gqa.py              # cache memory: MHA vs GQA vs MQA
 python scripts/bench_batching.py         # continuous batching throughput
 python scripts/bench_speculative.py      # speculative decoding speedup
+python scripts/bench_bpe.py              # tokenizer training time and compression
 ```
 
 ## How each feature is verified
@@ -65,6 +71,14 @@ checks 2,500 speculative samples against it. A power check asserts the draft rea
 different distribution, and I confirmed the test fails when the algorithm is deliberately broken
 (always accepting proposals, or resampling from `p` instead of the residual `max(0, p - q)`).
 
+**BPE tokenizer.** Round-trips random unicode (accents, CJK, emoji, control characters,
+underscores) exactly. The incremental trainer is checked against a naive "recount every pair each
+iteration" reference and must learn identical merges; encoding is checked against a reference that
+applies merges in learned order; the classic `aaabdaaabac` example gives the textbook merges. A
+boundary test asserts merges never cross chunk boundaries. As with speculative decoding, I
+confirmed the tests fail when the code is deliberately broken (stale pair counts, wrong
+tie-break, wrong merge priority in the encoder).
+
 ## Results
 
 All numbers were measured with this repo's scripts on an Intel iMac (6 CPU threads, no GPU).
@@ -83,6 +97,31 @@ CPU timings are noisy; treat small differences as approximate.
 
 Initial loss 4.24 is close to ln 65 = 4.17 (uniform over the 65-character vocabulary), as expected.
 
+### BPE tokenizer vs characters
+
+Same architecture, batch size (24) and step count (1200); only the tokenizer differs. Loss per
+token is not comparable across tokenizers, so the comparison is **bits per byte** of held-out text:
+
+| tokenizer | vocab | bytes / token | params | val bits/byte | train time |
+|:--|--:|--:|--:|--:|--:|
+| character | 65 | 1.00 | 0.81M | 2.303 | 5.6 min |
+| byte-level BPE | 512 | 1.88 | 0.87M | **2.127** | 7.4 min |
+
+BPE needs about 8% fewer bits per byte. Read this with care: a BPE batch covers ~1.9x more text
+(the same 128-token context spans more characters), the model is slightly larger (bigger
+embedding table), and it trained ~30% longer, so this is not a compute-matched comparison. It also
+overfits sooner: validation loss rose between steps 1000 and 1200 (2.116 -> 2.127 bits/byte) while
+training loss kept falling, on only ~0.5M training tokens.
+
+Tokenizer cost and compression on held-out text (`scripts/bench_bpe.py`, pure Python):
+
+| vocab | train time (s) | bytes/token | encode (MB/s) | round-trip |
+|---:|---:|---:|---:|:-:|
+| 300 | 0.5 | 1.38 | 2.73 | yes |
+| 512 | 1.1 | 1.88 | 2.08 | yes |
+| 1024 | 2.7 | 2.26 | 1.94 | yes |
+| 2048 | 6.5 | 2.56 | 1.80 | yes |
+
 ### KV cache
 
 | new tokens | no cache (tok/s) | KV cache (tok/s) | speedup | identical output |
@@ -155,6 +194,10 @@ GPU-scale target the same acceptance rates would give a larger speedup.
 * **Ties in parity tests (avoided).** The default init gives near-uniform logits, where greedy
   argmax could flip on float noise between code paths. Parity tests therefore use a wider init
   (`tests/helpers.py`) so they compare confident outputs rather than ties.
+* **A test that could not fail.** When I deliberately broke the BPE encoder (applying the
+  highest-rank merge first instead of the lowest), all 19 tests still passed: the random-string
+  comparison against the reference almost never hits two competing overlapping merges. I added a
+  deterministic test (`ab` vs `bc` competing for the same `b`) that does catch it.
 * **Wrong test, not wrong code.** One batching test failed because it sampled cache width *after*
   a row had already been retired in the same step; the engine was right, the assertion was not.
 
@@ -167,11 +210,14 @@ GPU-scale target the same acceptance rates would give a larger speedup.
 * The cache grows by `torch.cat` each step (simple and correct, but it reallocates); a production
   system would use preallocated or paged blocks.
 * Attention is the plain `O(T^2)`-memory formulation: no FlashAttention.
-* Character-level and small-scale: the point is correctness and clarity, not sample quality.
+* The BPE trainer and encoder are pure Python: fine for megabyte-scale corpora (seconds), not for
+  web-scale ones. Pre-tokenization is a simplified GPT-2-style regex, not the exact `tiktoken`
+  pattern, so token ids are not compatible with any public tokenizer.
+* Small scale (0.8M parameters, 1.1 MB of text): the point is correctness and clarity, not sample
+  quality.
 
 ## Roadmap
 
-- [ ] BPE tokenizer trained from scratch, with round-trip tests
 - [ ] Paged KV cache (fixed-size blocks, block table) to remove the `torch.cat` reallocation
 - [ ] Speculative decoding inside the batch engine
 - [ ] Load a public checkpoint into this architecture and match Hugging Face logits

scripts/bench_bpe.py

@@ -0,0 +1,42 @@
+"""BPE tokenizer: training time, compression on held-out text, and encode throughput.
+
+    python scripts/bench_bpe.py --data data/shakespeare.txt
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from microgpt.bpe import BPETokenizer
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--data", default="data/shakespeare.txt")
+    args = p.parse_args()
+
+    text = Path(args.data).read_text(encoding="utf-8")
+    cut = int(0.9 * len(text))
+    train_text, val_text = text[:cut], text[cut:]
+    val_bytes = len(val_text.encode("utf-8"))
+    print(f"{len(train_text):,} train chars, {val_bytes:,} held-out bytes\n")
+    print("| vocab | train time (s) | bytes/token (held-out) | encode (MB/s) | round-trip |")
+    print("|---:|---:|---:|---:|:-:|")
+    for vocab in (300, 512, 1024, 2048):
+        t0 = time.perf_counter()
+        tok = BPETokenizer.train(train_text, vocab)
+        train_s = time.perf_counter() - t0
+        t0 = time.perf_counter()
+        ids = tok.encode(val_text)
+        enc_s = time.perf_counter() - t0
+        ok = tok.decode(ids) == val_text
+        print(
+            f"| {vocab} | {train_s:.1f} | {val_bytes / len(ids):.2f} "
+            f"| {val_bytes / enc_s / 1e6:.2f} | {'yes' if ok else 'NO'} |"
+        )
+
+
+if __name__ == "__main__":
+    main()

src/microgpt/__init__.py

@@ -1,8 +1,9 @@
 """microgpt: a from-scratch GPT with KV-cache inference."""
 
+from microgpt.bpe import BPETokenizer
 from microgpt.data import CharTokenizer
 from microgpt.model import GPT, GPTConfig
 from microgpt.sampling import sample_next
 
-__all__ = ["GPT", "GPTConfig", "CharTokenizer", "sample_next"]
+__all__ = ["GPT", "GPTConfig", "BPETokenizer", "CharTokenizer", "sample_next"]
 __version__ = "0.1.0"

src/microgpt/bpe.py

@@ -0,0 +1,184 @@
+"""Byte-level BPE tokenizer, trained and applied from scratch.
+
+Design (the GPT-2 family recipe, minus the heavy dependencies):
+  * The base vocabulary is the 256 byte values, so *any* string round-trips losslessly:
+    there is no "unknown token".
+  * Text is first split into chunks (words, numbers, punctuation runs, whitespace) with a regex,
+    and merges never cross chunk boundaries. This stops the tokenizer from learning junk like
+    "dog." as one token and keeps training fast (we count unique chunks, weighted by frequency).
+  * Training repeatedly merges the most frequent adjacent pair into a new token. Pair counts are
+    updated incrementally: only the chunks that contain the merged pair are touched.
+  * Encoding applies merges in the order they were learned (lowest rank first).
+  * Ties are broken by the smallest pair, so training is fully deterministic.
+  * Optional special tokens (e.g. "<|endoftext|>") get ids after the merges and are never produced
+    by splitting ordinary text.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Dict, Iterable, List, Sequence, Set, Tuple
+
+Pair = Tuple[int, int]
+
+# Contractions | letters | digits | punctuation/symbols (incl. "_") | whitespace.
+# Every character matches exactly one alternative, so the chunks always re-join to the input.
+PRETOKENIZE = re.compile(
+    r"""'(?:[sdmt]|ll|ve|re)| ?[^\W\d_]+| ?\d+| ?(?:[^\s\w]|_)+|\s+(?!\S)|\s+"""
+)
+
+
+def pretokenize(text: str) -> List[str]:
+    return PRETOKENIZE.findall(text)
+
+
+def merge_pair(ids: List[int], pair: Pair, new_id: int) -> List[int]:
+    """Replace every non-overlapping, left-to-right occurrence of `pair` with `new_id`."""
+    out: List[int] = []
+    i, n = 0, len(ids)
+    while i < n:
+        if i < n - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
+            out.append(new_id)
+            i += 2
+        else:
+            out.append(ids[i])
+            i += 1
+    return out
+
+
+def _pair_multiset(ids: Sequence[int]) -> Counter:
+    return Counter(zip(ids, ids[1:]))
+
+
+class BPETokenizer:
+    def __init__(self, merges: Sequence[Pair] = (), special_tokens: Sequence[str] = ()) -> None:
+        self.merges: List[Pair] = [tuple(m) for m in merges]  # type: ignore[misc]
+        self.special_tokens: List[str] = list(special_tokens)
+        self._rank: Dict[Pair, int] = {m: i for i, m in enumerate(self.merges)}
+        self._vocab: List[bytes] = [bytes([i]) for i in range(256)]
+        for a, b in self.merges:
+            self._vocab.append(self._vocab[a] + self._vocab[b])
+        self._special_id = {s: len(self._vocab) + i for i, s in enumerate(self.special_tokens)}
+        self._special_re = (
+            re.compile("(" + "|".join(re.escape(s) for s in self.special_tokens) + ")")
+            if self.special_tokens
+            else None
+        )
+        self._cache: Dict[str, List[int]] = {}
+
+    # ------------------------------------------------------------------ training
+    @classmethod
+    def train(
+        cls, text: str, vocab_size: int, special_tokens: Sequence[str] = ()
+    ) -> BPETokenizer:
+        n_merges = vocab_size - 256 - len(special_tokens)
+        if n_merges < 0:
+            raise ValueError(
+                f"vocab_size {vocab_size} is too small: need at least "
+                f"{256 + len(special_tokens)} (256 bytes + {len(special_tokens)} special tokens)"
+            )
+        freq = Counter(pretokenize(text))
+        words: List[List[int]] = [list(w.encode("utf-8")) for w in freq]
+        weights: List[int] = list(freq.values())
+
+        counts: Dict[Pair, int] = defaultdict(int)
+        where: Dict[Pair, Set[int]] = defaultdict(set)
+        for wi, (w, c) in enumerate(zip(words, weights)):
+            for pair, k in _pair_multiset(w).items():
+                counts[pair] += c * k
+                where[pair].add(wi)
+
+        merges: List[Pair] = []
+        while len(merges) < n_merges and counts:
+            # Most frequent pair; ties go to the numerically smallest pair (deterministic).
+            best = max(counts.items(), key=lambda kv: (kv[1], -kv[0][0], -kv[0][1]))[0]
+            new_id = 256 + len(merges)
+            merges.append(best)
+            for wi in where.pop(best):
+                w, c = words[wi], weights[wi]
+                for pair, k in _pair_multiset(w).items():  # retract this word's old pairs
+                    counts[pair] -= c * k
+                    if counts[pair] == 0:
+                        del counts[pair]
+                w = words[wi] = merge_pair(w, best, new_id)
+                for pair, k in _pair_multiset(w).items():  # add its new pairs
+                    counts[pair] += c * k
+                    where[pair].add(wi)
+            counts.pop(best, None)
+        return cls(merges, special_tokens)
+
+    # ------------------------------------------------------------------ properties
+    @property
+    def vocab_size(self) -> int:
+        return len(self._vocab) + len(self.special_tokens)
+
+    # ------------------------------------------------------------------ encode / decode
+    def _encode_chunk(self, chunk: str) -> List[int]:
+        cached = self._cache.get(chunk)
+        if cached is not None:
+            return cached
+        ids = list(chunk.encode("utf-8"))
+        while len(ids) > 1:
+            best, best_rank = None, None
+            for pair in zip(ids, ids[1:]):
+                rank = self._rank.get(pair)
+                if rank is not None and (best_rank is None or rank < best_rank):
+                    best, best_rank = pair, rank
+            if best is None:
+                break
+            ids = merge_pair(ids, best, 256 + best_rank)  # type: ignore[operator]
+        self._cache[chunk] = ids
+        return ids
+
+    def encode_ordinary(self, text: str) -> List[int]:
+        """Encode text, treating special-token strings as plain text."""
+        out: List[int] = []
+        for chunk in pretokenize(text):
+            out.extend(self._encode_chunk(chunk))
+        return out
+
+    def encode(self, text: str, allow_special: bool = True) -> List[int]:
+        if not (allow_special and self._special_re):
+            return self.encode_ordinary(text)
+        out: List[int] = []
+        for part in self._special_re.split(text):
+            if part in self._special_id:
+                out.append(self._special_id[part])
+            elif part:
+                out.extend(self.encode_ordinary(part))
+        return out
+
+    def decode(self, ids: Iterable[int]) -> str:
+        """Bytes -> UTF-8. A token sequence cut mid-character decodes to U+FFFD, not an error."""
+        special_by_id = {i: s for s, i in self._special_id.items()}
+        parts: List[bytes] = []
+        for i in ids:
+            if 0 <= i < len(self._vocab):
+                parts.append(self._vocab[i])
+            elif i in special_by_id:
+                parts.append(special_by_id[i].encode("utf-8"))
+            else:
+                raise ValueError(f"token id {i} is outside the vocabulary (size {self.vocab_size})")
+        return b"".join(parts).decode("utf-8", errors="replace")
+
+    # ------------------------------------------------------------------ persistence
+    def to_dict(self) -> dict:
+        return {
+            "type": "bpe",
+            "merges": [list(m) for m in self.merges],
+            "special_tokens": self.special_tokens,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict) -> BPETokenizer:
+        return cls([tuple(m) for m in d["merges"]], d.get("special_tokens", []))
+
+    def save(self, path: str) -> None:
+        Path(path).write_text(json.dumps(self.to_dict()), encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: str) -> BPETokenizer:
+        return cls.from_dict(json.loads(Path(path).read_text(encoding="utf-8")))

src/microgpt/data.py

@@ -1,11 +1,13 @@
-"""Character-level tokenizer and batching."""
+"""Tokenizers and batching."""
 
 from __future__ import annotations
 
-from typing import List, Sequence, Tuple
+from typing import List, Sequence, Tuple, Union
 
 import torch
 
+from microgpt.bpe import BPETokenizer
+
 
 class CharTokenizer:
     def __init__(self, chars: Sequence[str]) -> None:
@@ -26,6 +28,21 @@ class CharTokenizer:
     def decode(self, ids: Sequence[int]) -> str:
         return "".join(self.chars[i] for i in ids)
 
+    def to_dict(self) -> dict:
+        return {"type": "char", "chars": self.chars}
+
+
+Tokenizer = Union[CharTokenizer, BPETokenizer]
+
+
+def tokenizer_from_dict(d: dict) -> Tokenizer:
+    kind = d.get("type")
+    if kind == "char":
+        return CharTokenizer(d["chars"])
+    if kind == "bpe":
+        return BPETokenizer.from_dict(d)
+    raise ValueError(f"unknown tokenizer type: {kind!r}")
+
 
 def get_batch(
     data: torch.Tensor, batch_size: int, block_size: int, generator: torch.Generator | None = None

src/microgpt/generate.py

@@ -7,7 +7,7 @@ import sys
 
 import torch
 
-from microgpt.data import CharTokenizer
+from microgpt.data import CharTokenizer, tokenizer_from_dict
 from microgpt.model import GPT, GPTConfig
 from microgpt.speculative import speculative_generate
 
@@ -17,7 +17,11 @@ def load_checkpoint(path: str):
     model = GPT(GPTConfig(**ckpt["config"]))
     model.load_state_dict(ckpt["model"])
     model.eval()
-    return model, CharTokenizer(ckpt["chars"])
+    if "tokenizer" in ckpt:
+        tok = tokenizer_from_dict(ckpt["tokenizer"])
+    else:  # checkpoints written before BPE support stored only the character list
+        tok = CharTokenizer(ckpt["chars"])
+    return model, tok
 
 
 def main() -> None:

src/microgpt/train.py

@@ -10,6 +10,7 @@ from pathlib import Path
 
 import torch
 
+from microgpt.bpe import BPETokenizer
 from microgpt.data import CharTokenizer, get_batch
 from microgpt.model import GPT, GPTConfig
 
@@ -37,6 +38,8 @@ def main() -> None:
     p = argparse.ArgumentParser(description=__doc__)
     p.add_argument("--data", required=True, help="path to a UTF-8 text file")
     p.add_argument("--out", default="checkpoints/model.pt")
+    p.add_argument("--tokenizer", choices=["char", "bpe"], default="char")
+    p.add_argument("--vocab-size", type=int, default=512, help="BPE vocabulary size (incl. bytes)")
     p.add_argument("--steps", type=int, default=1000)
     p.add_argument("--batch-size", type=int, default=32)
     p.add_argument("--block-size", type=int, default=128)
@@ -52,10 +55,18 @@ def main() -> None:
 
     torch.manual_seed(args.seed)
     text = Path(args.data).read_text(encoding="utf-8")
+    cut = int(0.9 * len(text))
+    train_text, val_text = text[:cut], text[cut:]
+    if args.tokenizer == "bpe":
+        tok = BPETokenizer.train(train_text, args.vocab_size)  # merges learned on train split only
+    else:
         tok = CharTokenizer.from_text(text)
-    data = torch.tensor(tok.encode(text), dtype=torch.long)
+    train_data = torch.tensor(tok.encode(train_text), dtype=torch.long)
-    split = int(0.9 * len(data))
+    val_data = torch.tensor(tok.encode(val_text), dtype=torch.long)
-    train_data, val_data = data[:split], data[split:]
+    # Losses are nats per *token*, which is not comparable across tokenizers. Bits per byte of
+    # raw text is: it is how many bits the model needs per byte of the held-out text.
+    bits_per_byte = math.log2(math.e) * len(val_data) / len(val_text.encode("utf-8"))
+    print(f"tokenizer={args.tokenizer} bytes/token={len(val_text.encode()) / len(val_data):.2f}")
 
     cfg = GPTConfig(
         vocab_size=tok.vocab_size,
@@ -66,7 +77,7 @@ def main() -> None:
         n_kv_head=args.n_kv_head,
     )
     model = GPT(cfg)
-    print(f"vocab={tok.vocab_size} params={model.num_params():,} train_chars={len(train_data):,}")
+    print(f"vocab={tok.vocab_size} params={model.num_params():,} train_tokens={len(train_data):,}")
 
     decay = [p for p in model.parameters() if p.dim() >= 2]
     no_decay = [p for p in model.parameters() if p.dim() < 2]
@@ -83,7 +94,10 @@ def main() -> None:
         if step % args.eval_every == 0:
             tr = estimate_loss(model, train_data, args.batch_size)
             va = estimate_loss(model, val_data, args.batch_size)
-            print(f"step {step:5d} | train {tr:.3f} | val {va:.3f} | {time.time() - t0:6.1f}s")
+            print(
+                f"step {step:5d} | train {tr:.3f} | val {va:.3f} "
+                f"| val bits/byte {va * bits_per_byte:.3f} | {time.time() - t0:6.1f}s"
+            )
         if step == args.steps:
             break
         x, y = get_batch(train_data, args.batch_size, args.block_size)
@@ -95,7 +109,8 @@ def main() -> None:
 
     out = Path(args.out)
     out.parent.mkdir(parents=True, exist_ok=True)
-    torch.save({"config": asdict(cfg), "model": model.state_dict(), "chars": tok.chars}, out)
+    ckpt = {"config": asdict(cfg), "model": model.state_dict(), "tokenizer": tok.to_dict()}
+    torch.save(ckpt, out)
     print(f"saved {out}")
 
 

tests/test_bpe.py

@@ -0,0 +1,215 @@
+import json
+import random
+import sys
+from collections import Counter
+
+import pytest
+import torch
+
+from microgpt.bpe import BPETokenizer, merge_pair, pretokenize
+from microgpt.data import CharTokenizer
+from microgpt.generate import load_checkpoint
+from microgpt.model import GPT, GPTConfig
+
+POOL = list("abcdefg hij.,!?'\n\t_0123456789  ") + ["é", "ü", "日", "本", "😀", "—", "ß"]
+PARAGRAPH = (
+    "the quick brown fox jumps over the lazy dog. the dog barks; the fox runs away! "
+    "it's 2024 and we're testing tokenizers, aren't we? yes, they'll work. "
+)
+
+
+def random_text(seed: int, n: int) -> str:
+    rng = random.Random(seed)
+    return "".join(rng.choice(POOL) for _ in range(n))
+
+
+# --------------------------------------------------------------------------- references
+def ref_train(text: str, n_merges: int):
+    """Naive BPE: recount every pair from scratch each iteration. Slow but obviously correct."""
+    freq = Counter(pretokenize(text))
+    seqs = {w: list(w.encode("utf-8")) for w in freq}
+    merges = []
+    for _ in range(n_merges):
+        counts: Counter = Counter()
+        for w, c in freq.items():
+            for pair in zip(seqs[w], seqs[w][1:]):
+                counts[pair] += c
+        if not counts:
+            break
+        best = max(counts.items(), key=lambda kv: (kv[1], -kv[0][0], -kv[0][1]))[0]
+        merges.append(best)
+        for w in seqs:
+            seqs[w] = merge_pair(seqs[w], best, 255 + len(merges))
+    return merges
+
+
+def ref_encode(chunk: str, merges):
+    """Apply every merge, in the order it was learned, to the whole chunk."""
+    ids = list(chunk.encode("utf-8"))
+    for rank, pair in enumerate(merges):
+        ids = merge_pair(ids, pair, 256 + rank)
+    return ids
+
+
+@pytest.fixture(scope="module")
+def tok():
+    return BPETokenizer.train(PARAGRAPH * 30 + random_text(0, 1500), 256 + 80)
+
+
+# --------------------------------------------------------------------------- basics
+def test_merge_pair_is_non_overlapping_left_to_right():
+    assert merge_pair([1, 1, 1], (1, 1), 9) == [9, 1]
+    assert merge_pair([1, 1, 1, 1], (1, 1), 9) == [9, 9]
+    assert merge_pair([1, 2, 1, 2, 3], (1, 2), 9) == [9, 9, 3]
+    assert merge_pair([], (1, 2), 9) == []
+
+
+def test_byte_vocabulary_roundtrips_anything_with_no_merges():
+    t = BPETokenizer()
+    assert t.vocab_size == 256
+    for s in ["", "hello", "héllo wörld 日本語 😀\n\t", "\x00\x01 control", "a" * 1000]:
+        assert t.decode(t.encode(s)) == s
+    assert t.encode("é") == list("é".encode())
+
+
+def test_classic_wikipedia_example():
+    t = BPETokenizer.train("aaabdaaabac", 259)
+    assert t.merges == [(97, 97), (97, 98), (256, 257)]  # aa, ab, (aa)(ab)
+    assert t.encode("aaabdaaabac") == [258, 100, 258, 97, 99]
+    assert t.decode([258, 100, 258, 97, 99]) == "aaabdaaabac"
+
+
+def test_training_is_deterministic():
+    text = PARAGRAPH * 5
+    assert BPETokenizer.train(text, 300).merges == BPETokenizer.train(text, 300).merges
+
+
+def test_training_stops_when_nothing_is_left_to_merge():
+    t = BPETokenizer.train("ab", 400)
+    assert t.merges == [(97, 98)] and t.vocab_size == 257
+
+
+# --------------------------------------------------------------------------- vs references
+@pytest.mark.parametrize("seed", [0, 1, 2])
+def test_incremental_training_matches_naive_reference(seed):
+    text = PARAGRAPH * 3 + random_text(seed, 2500)
+    assert BPETokenizer.train(text, 256 + 70).merges == ref_train(text, 70)
+
+
+def test_encoding_matches_reference_merge_order(tok):
+    for seed in range(60):
+        s = random_text(seed, 80) + PARAGRAPH[: seed + 5]
+        expected = [i for chunk in pretokenize(s) for i in ref_encode(chunk, tok.merges)]
+        assert tok.encode(s) == expected
+
+
+# --------------------------------------------------------------------------- losslessness
+def test_overlapping_merges_resolve_by_learned_order():
+    """When two merges compete for the same byte, the earlier-learned one must win."""
+    t = BPETokenizer(merges=[(97, 98), (98, 99)])  # ab = 256 (rank 0), bc = 257 (rank 1)
+    assert t.encode("abc") == [256, 99]  # NOT [97, 257]
+    t2 = BPETokenizer(merges=[(98, 99), (97, 98)])  # reversed priority
+    assert t2.encode("abc") == [97, 256]
+    assert t.decode(t.encode("abcabc")) == "abcabc"
+
+
+def test_roundtrip_random_unicode(tok):
+    for seed in range(400):
+        s = random_text(seed, random.Random(seed).randint(0, 70))
+        assert tok.decode(tok.encode(s)) == s
+
+
+def test_pretokenizer_is_lossless_and_splits_sensibly():
+    for seed in range(500):
+        s = random_text(seed, 60)
+        assert "".join(pretokenize(s)) == s
+    assert pretokenize("Hello, world!  It's 42_x") == [
+        "Hello", ",", " world", "!", " ", " It", "'s", " 42", "_", "x",
+    ]  # fmt: skip
+
+
+def test_merges_never_cross_chunk_boundaries():
+    t = BPETokenizer.train("the cat sat on the mat. " * 200, 256 + 100)
+    for rank in range(len(t.merges)):
+        token = t._vocab[256 + rank]
+        assert b" " not in token[1:], token  # a space may only *start* a token
+
+
+def test_tokenizer_compresses_text(tok):
+    ids = tok.encode(PARAGRAPH * 3)
+    assert len(PARAGRAPH * 3) / len(ids) > 1.8
+
+
+# --------------------------------------------------------------------------- special tokens
+def test_special_tokens():
+    eot = "<|endoftext|>"
+    t = BPETokenizer.train(PARAGRAPH * 10, 256 + 20 + 1, special_tokens=[eot])
+    assert t.vocab_size == 256 + 20 + 1
+    text = f"the fox{eot}the dog"
+    ids = t.encode(text)
+    assert ids.count(t.vocab_size - 1) == 1
+    assert t.decode(ids) == text
+    plain = t.encode(text, allow_special=False)  # treated as ordinary characters
+    assert t.vocab_size - 1 not in plain and t.decode(plain) == text
+    assert t.encode(eot) == [t.vocab_size - 1]
+
+
+# --------------------------------------------------------------------------- persistence / errors
+def test_save_load_roundtrip(tok, tmp_path):
+    path = tmp_path / "tok.json"
+    tok.save(str(path))
+    loaded = BPETokenizer.load(str(path))
+    assert loaded.merges == tok.merges and json.loads(path.read_text())["type"] == "bpe"
+    s = PARAGRAPH + random_text(3, 100)
+    assert loaded.encode(s) == tok.encode(s)
+
+
+def test_errors():
+    with pytest.raises(ValueError):
+        BPETokenizer.train("abc", 255)
+    with pytest.raises(ValueError):
+        BPETokenizer.train("abc", 256, special_tokens=["<x>"])
+    t = BPETokenizer()
+    for bad in (256, -1):
+        with pytest.raises(ValueError):
+            t.decode([bad])
+
+
+def test_decoding_a_split_character_does_not_crash():
+    t = BPETokenizer()
+    assert t.decode([0xE6]) == "�"  # first byte of a 3-byte character
+    assert t.decode(list("日".encode())) == "日"
+
+
+# --------------------------------------------------------------------------- integration
+TINY = dict(n_layer=1, n_head=2, n_embd=16)
+
+
+def test_train_cli_with_bpe_then_generate(tmp_path, monkeypatch):
+    from microgpt import train
+
+    data = tmp_path / "text.txt"
+    data.write_text(PARAGRAPH * 40, encoding="utf-8")
+    out = tmp_path / "m.pt"
+    argv = [
+        "train", "--data", str(data), "--out", str(out), "--tokenizer", "bpe",
+        "--vocab-size", "300", "--steps", "3", "--batch-size", "2", "--block-size", "16",
+        "--n-layer", "1", "--n-embd", "16", "--n-head", "2", "--eval-every", "3",
+    ]  # fmt: skip
+    monkeypatch.setattr(sys, "argv", argv)
+    train.main()
+
+    model, loaded = load_checkpoint(str(out))
+    assert isinstance(loaded, BPETokenizer)
+    assert model.cfg.vocab_size == loaded.vocab_size == 300
+    out_ids = model.generate(torch.tensor([loaded.encode("the ")]), 6, temperature=0)
+    assert isinstance(loaded.decode(out_ids[0].tolist()), str)
+
+
+def test_legacy_char_checkpoint_still_loads(tmp_path):
+    chars = list("abc ")
+    cfg = GPTConfig(vocab_size=len(chars), block_size=16, **TINY)
+    path = tmp_path / "old.pt"
+    torch.save({"config": cfg.__dict__, "model": GPT(cfg).state_dict(), "chars": chars}, path)
+    _, loaded = load_checkpoint(str(path))
+    assert isinstance(loaded, CharTokenizer) and loaded.decode(loaded.encode("ab c")) == "ab c"