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.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
|
||||
@@ -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")))
|
||||
+19
-2
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
+21
-6
@@ -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)
|
||||
split = int(0.9 * len(data))
|
||||
train_data, val_data = data[:split], data[split:]
|
||||
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}")
|
||||
|
||||
|
||||
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user