Chapter 4 · Part 2

Train the tokenizer

Training BPE is just count-and-merge, in a loop. Each round: find the most common pair, merge it into the next new token id, and remember the rule. The magic is that later rounds merge pairs that include tokens from earlier rounds — so the tokenizer bootstraps bigger and bigger chunks out of smaller ones.

The training loop

You decide the final vocabulary size; everything above 256 is a learned merge. We'll do 20 merges (vocab of 276). The merges dict records every rule — pair → new id — which is the entire trained model. Add to bpe.py:

bpe.py (continued)
vocab_size = 276               # 256 bytes + 20 learned merges
num_merges = vocab_size - 256

ids = list(text.encode("utf-8"))
merges = {}                    # (int, int) -> new id  — the trained tokenizer

for i in range(num_merges):
  stats = get_stats(ids)
  pair = max(stats, key=stats.get)   # most common pair right now
  new_id = 256 + i
  ids = merge(ids, pair, new_id)     # fuse it everywhere
  merges[pair] = new_id              # remember the rule
  print(f"merge {i:2d}: {pair} -> {new_id}")

Decode the vocabulary

The merges dict is ids merging ids — to read it, build a vocab that maps each id back to the actual bytes it stands for. Each new token is just its two parts concatenated:

bpe.py (continued)
vocab = {i: bytes([i]) for i in range(256)}
for (a, b), idx in merges.items():
  vocab[idx] = vocab[a] + vocab[b]   # earlier ids are already in vocab

for (a, b), idx in merges.items():
  print(f"{idx} = {vocab[idx].decode('utf-8', errors='replace')!r}")

Run it — and watch words appear

terminal
python bpe.py
256 = ' t'
257 = 'he'
258 = 'at'
259 = ' the'
260 = 's '
261 = ' c'
262 = ' to'
263 = ' tok'
264 = ' toke'
265 = ' token'
...
269 = ' te'
270 = ' tex'
271 = ' text'

Look at what it discovered on its own. Token 256 is " t". Token 262 fuses that with o into " to". Then 263 = " tok", 264 = " toke", and 265 = " token" — the tokenizer built the word " token" one merge at a time, each step reusing the token before it. Tokens 269–271 do the same to assemble " text". Nobody told it these were words; it just kept merging what was common, and language fell out. That's BPE. Now let's use it on new text.