Problem 101 · Tokenization & Sampling

Byte-Pair Encoding: Train Merges and Encode

HardAI engineering

01Problem

Implement BPE from scratch. `train_bpe(corpus, num_merges)` splits each text on whitespace, represents every word as characters plus an end-of-word marker `</w>`, and repeatedly merges the most frequent adjacent pair (ties broken by the lexicographically smallest pair) until `num_merges` merges are learned or no pairs remain. `encode(word, merges)` tokenises a new word by applying the learned merges in rank order.

02Key insight

BPE is a greedy compression algorithm, and the rank order **is** the model: at encode time you must apply merges by their training rank (lowest rank first), not left-to-right and not by frequency in the new word, or you get token ids the model never saw in training. Count pairs weighted by word frequency over a `Counter` of unique words rather than over the raw corpus; that is the difference between minutes and hours on real data. Deterministic tie-breaking matters too: two training runs that disagree on a tie produce different vocabularies and incompatible checkpoints.

03Solution

solution.pypython51 lines

from collections import Counter

EOW = "</w>"
Pair = tuple[str, str]


def _merge(symbols: tuple[str, ...], pair: Pair) -> tuple[str, ...]:
    out: list[str] = []
    i = 0
    while i < len(symbols):
        if i + 1 < len(symbols) and (symbols[i], symbols[i + 1]) == pair:
            out.append(symbols[i] + symbols[i + 1])
            i += 2
        else:
            out.append(symbols[i])
            i += 1
    return tuple(out)


def train_bpe(corpus: list[str], num_merges: int) -> list[Pair]:
    words: Counter[tuple[str, ...]] = Counter()
    for text in corpus:
        for w in text.split():
            words[tuple(w) + (EOW,)] += 1
    merges: list[Pair] = []
    for _ in range(num_merges):
        pairs: Counter[Pair] = Counter()
        for symbols, freq in words.items():
            for a, b in zip(symbols, symbols[1:]):
                pairs[(a, b)] += freq
        if not pairs:
            break
        best = min(pairs, key=lambda p: (-pairs[p], p))
        merges.append(best)
        merged: Counter[tuple[str, ...]] = Counter()
        for symbols, freq in words.items():
            merged[_merge(symbols, best)] += freq
        words = merged
    return merges


def encode(word: str, merges: list[Pair]) -> list[str]:
    rank = {p: i for i, p in enumerate(merges)}
    symbols = tuple(word) + (EOW,)
    while len(symbols) > 1:
        candidates = [(rank[p], p) for p in zip(symbols, symbols[1:]) if p in rank]
        if not candidates:
            break
        _, best = min(candidates)
        symbols = _merge(symbols, best)
    return list(symbols)

Try it first.

TimeTrain O(M * U * L) for M merges over U unique words of length L; encode O(L^2) per word

SpaceO(U * L + M)

tests.py · 8 passing

tests.pypython11 lines

corpus = ["low " * 5 + "lower " * 2 + "newest " * 6 + "widest " * 3]
merges = train_bpe(corpus, 5)
assert merges[:4] == [("e", "s"), ("es", "t"), ("est", EOW), ("l", "o")], merges
assert merges[4] == ("lo", "w")
assert encode("lowest", merges) == ["low", "est</w>"]
assert encode("low", merges) == ["low", EOW]
assert encode("zz", merges) == ["z", "z", EOW]
toks = encode("newest", merges)
assert "".join(toks).replace(EOW, "") == "newest"
assert train_bpe(["a"], 10) == [("a", EOW)]
assert train_bpe([""], 3) == []

04Staff-level follow-up

How does production BPE stay fast? Index pair -> words so each merge only updates affected words (O(total changes) instead of rescanning), pre-tokenise with a regex so merges never cross word or digit boundaries, and run byte-level so there is no unknown-token path. At encode time cache per-word results; most traffic is a few thousand frequent words.