From bd349c61931d10b9cec10d816aca7c07ad3b2758 Mon Sep 17 00:00:00 2001 From: jzzliu Date: Sat, 7 Nov 2026 20:48:09 -0700 Subject: [PATCH] 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 | 60 +++++++++-- scripts/bench_bpe.py | 42 ++++++++ src/microgpt/__init__.py | 3 +- src/microgpt/bpe.py | 184 +++++++++++++++++++++++++++++++++ src/microgpt/data.py | 21 +++- src/microgpt/generate.py | 8 +- src/microgpt/train.py | 29 ++++-- tests/test_bpe.py | 215 +++++++++++++++++++++++++++++++++++++++ 8 files changed, 543 insertions(+), 19 deletions(-) create mode 100644 scripts/bench_bpe.py create mode 100644 src/microgpt/bpe.py create mode 100644 tests/test_bpe.py diff --git a/README.md b/README.md index 980100f..40bafb7 100644 --- a/README.md +++ b/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 -backed by tests that check correctness against a reference, not just "it runs". +batching** over variable-length requests, **speculative decoding**, and a **byte-level BPE +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 -written out explicitly (~900 lines of library code) so every line can be reasoned about. +No `nn.Transformer`, no Hugging Face, no tokenizer library. Attention, masking, caching, +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 diff --git a/scripts/bench_bpe.py b/scripts/bench_bpe.py new file mode 100644 index 0000000..ce62e7c --- /dev/null +++ b/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() diff --git a/src/microgpt/__init__.py b/src/microgpt/__init__.py index 88d78be..94ae8bf 100644 --- a/src/microgpt/__init__.py +++ b/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" diff --git a/src/microgpt/bpe.py b/src/microgpt/bpe.py new file mode 100644 index 0000000..7d5b36c --- /dev/null +++ b/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"))) diff --git a/src/microgpt/data.py b/src/microgpt/data.py index 27a1b85..6e9bc08 100644 --- a/src/microgpt/data.py +++ b/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 diff --git a/src/microgpt/generate.py b/src/microgpt/generate.py index 75421b7..26e0abf 100644 --- a/src/microgpt/generate.py +++ b/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: diff --git a/src/microgpt/train.py b/src/microgpt/train.py index 2f3b2a6..83879a3 100644 --- a/src/microgpt/train.py +++ b/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") - tok = CharTokenizer.from_text(text) - data = torch.tensor(tok.encode(text), dtype=torch.long) - split = int(0.9 * len(data)) - train_data, val_data = data[:split], data[split:] + 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) + train_data = torch.tensor(tok.encode(train_text), dtype=torch.long) + val_data = torch.tensor(tok.encode(val_text), dtype=torch.long) + # 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}") diff --git a/tests/test_bpe.py b/tests/test_bpe.py new file mode 100644 index 0000000..bd8eaac --- /dev/null +++ b/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=[""]) + 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"