Subword Tokenization with BPE

Byte Pair Encoding

Most LLM tokenizers are built with BPE (Byte Pair Encoding). Start with single characters; repeatedly find the most frequent neighbouring pair and merge it into a new token. After many merges, common words become one token and rare words split into reusable pieces. That way the vocabulary stays fixed (e.g. 100k tokens) but any text can be encoded.

A Tiny BPE Trainer

Example

Train 6 merges on a few words. _ marks the end of a word. Ties are broken alphabetically so the output is repeatable:

from collections import Counter

corpus = ["lower", "lowest", "newer", "newest", "wider", "low", "new"]
words = [list(w) + ["_"] for w in corpus]

def pair_counts(words):
    c = Counter()
    for w in words:
        for pair in zip(w, w[1:]):
            c[pair] += 1
    return c

def merge(word, pair):
    out, i = [], 0
    while i < len(word):
        if i + 1 < len(word) and (word[i], word[i + 1]) == pair:
            out.append(word[i] + word[i + 1]); i += 2
        else:
            out.append(word[i]); i += 1
    return out

merges = []
for step in range(1, 7):
    counts = pair_counts(words)
    best = max(counts, key=lambda pr: (counts[pr], [-ord(ch) for ch in pr[0] + pr[1]]))
    merges.append(best)
    words = [merge(w, best) for w in words]
    print(f"merge {step}: {best[0]} + {best[1]} -> {best[0] + best[1]}  (count {counts[best]})")

print()
for w in words:
    print(" ".join(w))

Output

merge 1: w + e -> we  (count 4)
merge 2: l + o -> lo  (count 3)
merge 3: n + e -> ne  (count 3)
merge 4: r + _ -> r_  (count 3)
merge 5: lo + we -> lowe  (count 2)
merge 6: ne + we -> newe  (count 2)

lowe r_
lowe s t _
newe r_
newe s t _
w i d e r_
lo w _
ne w _

Tokenizing a New Word

Example

To tokenize new text, apply the learned merges in order. A word never seen in training still gets encoded from known pieces:

# the 6 merges learned by the trainer above, in order
merges = [("w", "e"), ("l", "o"), ("n", "e"), ("r", "_"), ("lo", "we"), ("ne", "we")]

def tokenize(word):
    toks = list(word) + ["_"]
    for a, b in merges:
        i, out = 0, []
        while i < len(toks):
            if i + 1 < len(toks) and toks[i] == a and toks[i + 1] == b:
                out.append(a + b); i += 2
            else:
                out.append(toks[i]); i += 1
        toks = out
    return toks

for w in ["lower", "slower", "newest", "renew"]:
    t = tokenize(w)
    print(f"{w:8} -> {t}  ({len(t)} tokens)")

Output

lower    -> ['lowe', 'r_']  (2 tokens)
slower   -> ['s', 'lowe', 'r_']  (3 tokens)
newest   -> ['newe', 's', 't', '_']  (4 tokens)
renew    -> ['r', 'e', 'ne', 'w', '_']  (5 tokens)