BPE Tokenizer From Scratch

A tokenizer is the map from a string to a list of integers and back. A language model never sees text, only those integers, so a mistake in this layer means every loss curve in the training run is measuring the wrong thing.

Byte-Pair Encoding (BPE) is the dominant family of subword tokenizers for general text models. It starts from a known alphabet, finds the adjacent symbol pair that appears most often in the training corpus, and merges that pair into a new symbol. It repeats until the vocabulary reaches the target size. Encoding new text reuses the same merge list in the same order.

We'll build the byte-level variant. The alphabet is the 256 raw bytes, not Unicode code points, which is what lets the tokenizer handle any UTF-8 input without falling back to an unknown token.

Training and inference share one merge table, and that sharing is the contract. Change the merge order at inference and you decode a different stream of ids.

The first 256 ids are reserved for the raw bytes 0x00 through 0xFF, so every input string can be expressed in the vocabulary before any merge happens. A small range after the byte block holds the special tokens. The training loop never proposes those ids as merge targets, because we keep them out of the pretokenized stream entirely.

The pretokenizer splits the corpus into runs of non-whitespace and runs of whitespace before training sees it, so punctuation stays attached to its word. Skip that split and the merge step happily learns merges that cross word boundaries (the space between two words included), so the vocabulary fills up with whole common phrases. With the split, merges stay inside a word and the result generalizes.

Each training step does three things. It walks every word in the corpus and counts how often each adjacent pair of current symbols appears, weighted by how often the word itself appears. It picks the pair with the highest count. Then it rewrites every occurrence of that pair into a single new symbol, whose id is the next free slot in the vocabulary, and records the merge.

The cost of each step is linear in the size of the corpus expressed as a list of symbol sequences. For a million words and a target vocabulary of ten thousand ids, the loop runs to completion in seconds because the symbol sequences shrink as merges land.

Encoding applies the merge table in the same order it was learned, and it never calls the pair counter. For a fresh word, the encoder starts from the byte split and scans the current sequence for the lowest-ranked merge, meaning the earliest one that applies. It performs that merge and scans again. The loop ends when no merge in the table applies to the current sequence.

Ordering by rank is what makes encoding deterministic and consistent with training on the same input. A merge learned first sits at the top of the table and gets applied first. If two merges could apply at the same position, the lower-rank one wins.

Special tokens are reserved ids that the byte stream can never produce. We assign them by hand, and two are enough for this lesson.

• <|endoftext|> separates documents during pretraining. It tells the model "a new document starts here, do not let the previous one's context leak in."
• <|pad|> fills out short sequences so a batch can be a rectangular tensor. The loss mask hides it during training.

The encoder takes a flag that allows special tokens in the input. With the flag off, the strings <|endoftext|> and <|pad|> are tokenized as the bytes that spell them out. With the flag on, the literal strings map to their reserved ids and no merge touches them.

Encoding then decoding must return the input bytes exactly. The decoder concatenates the byte expansion of every id in order. Every id is either a raw byte or the concatenation of two previously known ids, so the recursive expansion always terminates in raw bytes, and those bytes spell the UTF-8 string you started with.

The test suite checks this on an unseen sentence, on a sentence with a Unicode emoji, and on a sentence containing a literal <|endoftext|> token.

There's no regex-driven pretokenizer in the style of the largest production tokenizers here. The pretokenizer is a single \S+|\s+ split on whitespace runs, which is enough to produce sensible merges on a small training corpus, and the contract with the rest of the lesson chain stays the same. The next lesson treats the tokenizer as a black box and builds the sliding-window dataset on top of it.

The pair counter isn't parallelized either. A Python loop over a corpus of a few thousand words finishes in well under a second. For larger corpora, the obvious move is to count pairs per word in parallel and reduce.

main.py is built around one class and three entry points. BPETokenizer holds the vocabulary, the merge table, and the special-token table. train is the training loop, encode is the inference path, and decode is the byte concatenation. The demo at the bottom trains a small tokenizer on a built-in corpus, encodes a held-out sentence, decodes the ids back, and prints both. The tests in code/tests/test_bpe.py pin the round-trip property, the special-token reservation, and the merge ordering.

Run the demo, then change the demo's target (320 in the file) to a larger value and watch the encoded length of the held-out sentence drop. Training stops early once no pair occurs at least twice, so on the small built-in corpus the final vocabulary can fall short of the target. That curve is the BPE compression curve.