51 · BPE & SentencePiece from Scratch

51 · BPE & SentencePiece from Scratch#

Dependencies: torch
Runtime: ~5 min CPU

Chapter stub — implementation coming in a future commit.

Time-machine sentences#

Every chapter runs the same 20 held-out test sentences so the quality arc across chapters is directly comparable.

# Colab setup: install this part's pinned dependencies (skipped outside Colab).
import sys
if "google.colab" in sys.modules:
    get_ipython().system("wget -q -O requirements_part5.txt https://raw.githubusercontent.com/eduardosanchezg/mthistory/main/requirements_part5.txt")
    get_ipython().system("pip install -q -r requirements_part5.txt")

BPE and SentencePiece from Scratch#

Word-level vocabularies suffer from the open-vocabulary problem: any word unseen in training becomes <unk>. Subword tokenisation – Byte-Pair Encoding (Sennrich et al. 2016) and SentencePiece unigram-LM (Kudo 2018) – fixes this by splitting rare words into reusable pieces. Everything here is pure stdlib and runs with zero dependencies.

1. The OOV problem#

With a fixed word-level vocabulary built from training data, any test word not in that vocabulary is out-of-vocabulary (OOV). Rare and compound words – common in German – hit this constantly.

TRAIN_CORPUS = (
    "the cat sat on the mat . "
    "a dog ran in the park . "
    "the children play and run . "
    "cats and dogs play in parks ."
).split()

TEST_CORPUS = (
    "the kitten played in the parking lot . "
    "unhappy children running quickly ."
).split()

word_vocab = set(TRAIN_CORPUS)
print("word-level vocab size:", len(word_vocab))

def oov_rate(tokens, vocab):
    oov = [t for t in tokens if t not in vocab]
    return len(oov) / len(tokens), oov

rate, oov = oov_rate(TEST_CORPUS, word_vocab)
print("test tokens:", len(TEST_CORPUS))
print("OOV rate (word-level): %.1f%%" % (rate * 100))
print("OOV words:", oov)

2. Byte-Pair Encoding from scratch#

BPE training is dead simple:

Represent each word as a sequence of characters (with an end-of-word marker).
Count every adjacent symbol pair across the corpus.
Merge the most frequent pair into a single new symbol.
Repeat for num_merges steps.

The learned ordered list of merges is the model. To tokenise a new word you replay the merges greedily in order.

from collections import Counter

EOW = "</w>"  # end-of-word marker

def get_pair_counts(word_freqs):
    pairs = Counter()
    for word, freq in word_freqs.items():
        symbols = word.split()
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i + 1])] += freq
    return pairs

def merge_pair(pair, word_freqs):
    bigram = " ".join(pair)
    replacement = "".join(pair)
    out = {}
    for word, freq in word_freqs.items():
        new_word = word.replace(bigram, replacement)
        out[new_word] = freq
    return out

def learn_bpe(corpus_words, num_merges):
    # corpus_words: list of words (strings)
    freqs = Counter(corpus_words)
    # initial representation: space-separated chars + EOW
    word_freqs = {" ".join(list(w) + [EOW]): f for w, f in freqs.items()}
    merges = []
    for _ in range(num_merges):
        pairs = get_pair_counts(word_freqs)
        if not pairs:
            break
        best = max(pairs, key=pairs.get)
        if pairs[best] < 2:
            break  # no repeated pair worth merging
        word_freqs = merge_pair(best, word_freqs)
        merges.append(best)
    return merges

def apply_bpe(word, merges):
    symbols = list(word) + [EOW]
    merge_rank = {pair: i for i, pair in enumerate(merges)}
    while True:
        pairs = [(symbols[i], symbols[i + 1]) for i in range(len(symbols) - 1)]
        candidate = None
        best_rank = None
        for p in pairs:
            if p in merge_rank and (best_rank is None or merge_rank[p] < best_rank):
                best_rank = merge_rank[p]
                candidate = p
        if candidate is None:
            break
        # merge first occurrence of candidate, scan whole word
        merged = []
        i = 0
        while i < len(symbols):
            if i < len(symbols) - 1 and (symbols[i], symbols[i + 1]) == candidate:
                merged.append(symbols[i] + symbols[i + 1])
                i += 2
            else:
                merged.append(symbols[i])
                i += 1
        symbols = merged
    return symbols

merges = learn_bpe(TRAIN_CORPUS, num_merges=30)
print("learned %d merge operations (in order):" % len(merges))
for i, m in enumerate(merges):
    print("  %2d  %s + %s  ->  %s" % (i, m[0], m[1], "".join(m)))
print()
for w in ["the", "cats", "parking"]:
    print("%-9s ->" % w, apply_bpe(w, merges))

3. Effect on vocabulary size#

More merges = larger subword vocabulary but shorter sequences (closer to word-level). Fewer merges = tiny vocabulary but longer, more fragmented sequences. We sweep num_merges and tabulate the resulting subword vocabulary.

def bpe_vocab(corpus_words, merges):
    vocab = set()
    for w in set(corpus_words):
        for sym in apply_bpe(w, merges):
            vocab.add(sym)
    return vocab

print("%-8s %-12s %-18s" % ("merges", "vocab_size", "tokens/word(avg)"))
print("-" * 40)
for nm in [0, 5, 10, 20, 40]:
    m = learn_bpe(TRAIN_CORPUS, num_merges=nm)
    v = bpe_vocab(TRAIN_CORPUS, m)
    total = sum(len(apply_bpe(w, m)) for w in TRAIN_CORPUS)
    avg = total / len(TRAIN_CORPUS)
    print("%-8d %-12d %-18.2f" % (nm, len(v), avg))
print()
print("compound/rare words split into reusable pieces:")
m = learn_bpe(TRAIN_CORPUS, num_merges=30)
for w in ["parking", "running", "unhappy"]:
    print("  %-9s ->" % w, apply_bpe(w, m))

4. SentencePiece-style unigram LM (sketch)#

BPE is greedy and deterministic: it always applies the highest-priority merge. SentencePiece offers a unigram language model alternative that is probabilistic. It keeps a large pool of candidate pieces, each with a probability, and tokenises a word by choosing the segmentation that maximises the product of piece probabilities (via Viterbi). Training iteratively prunes low-probability pieces with EM.

Key differences:

	BPE	Unigram LM (SentencePiece)
selection	greedy merges	probabilistic, global
segmentation	deterministic	best-scoring (can sample)
training	frequency counting	EM with pruning

Below is a simplified Viterbi tokeniser given a piece->log-prob table.

import math

def unigram_tokenize(word, piece_logp):
    # Viterbi: best segmentation maximising sum of piece log-probs
    n = len(word)
    best = [(-math.inf, None)] * (n + 1)
    best[0] = (0.0, None)
    for end in range(1, n + 1):
        for start in range(end):
            piece = word[start:end]
            if piece in piece_logp:
                score = best[start][0] + piece_logp[piece]
                if score > best[end][0]:
                    best[end] = (score, start)
    # backtrack
    pieces = []
    end = n
    while end > 0:
        score, start = best[end]
        if start is None:
            return [word]  # no valid segmentation
        pieces.append(word[start:end])
        end = start
    return pieces[::-1]

# toy piece vocabulary with log-probs (higher = more probable)
pieces = {"un": -1.0, "happy": -1.2, "unhappy": -4.0, "park": -1.1,
          "ing": -0.9, "parking": -5.0, "run": -1.0, "running": -5.5,
          "n": -3.0, "i": -3.0, "g": -3.0, "p": -3.0, "a": -3.0,
          "r": -3.0, "k": -3.0, "h": -3.0, "y": -3.0, "u": -3.0}
for w in ["unhappy", "parking", "running"]:
    print("%-9s ->" % w, unigram_tokenize(w, pieces))
print()
print("note: it prefers frequent pieces (un+happy) over the rare whole word.")

5. Effect on OOV rate#

Now we re-run the OOV computation from step 1, but tokenise both train and test with BPE. Because any word decomposes into characters in the worst case, the OOV rate collapses to ~0.

merges = learn_bpe(TRAIN_CORPUS, num_merges=30)
subword_vocab = bpe_vocab(TRAIN_CORPUS, merges)
# also add all single characters seen in training so any char is coverable
for w in TRAIN_CORPUS:
    for ch in w:
        subword_vocab.add(ch)
subword_vocab.add(EOW)

def subword_oov_rate(tokens, merges, vocab):
    total = 0
    oov = 0
    for w in tokens:
        for sym in apply_bpe(w, merges):
            total += 1
            if sym not in vocab:
                oov += 1
    return oov / total if total else 0.0

word_rate, _ = oov_rate(TEST_CORPUS, word_vocab)
sub_rate = subword_oov_rate(TEST_CORPUS, merges, subword_vocab)
print("OOV rate before (word-level): %.1f%%" % (word_rate * 100))
print("OOV rate after  (BPE):        %.1f%%" % (sub_rate * 100))

6. Time-machine#

Tokenise the 20 shared held-out sentences with the trained BPE model and report segmentation statistics. Runs with stdlib only.

TIME_MACHINE_EN = [
    "A man is walking his dog in the park .",
    "Two children are playing near a fountain .",
    "A woman in a red coat is reading a book .",
    "Several people are waiting at a bus stop .",
    "A cyclist is riding through a crowded street .",
    "A dog is chasing a ball on the beach .",
    "A group of tourists is taking photographs .",
    "A young girl is feeding ducks by a pond .",
    "Two men are playing chess in a cafe .",
    "A street musician is playing the violin .",
    "Children are running through a field of flowers .",
    "An old man is sitting on a bench reading a newspaper .",
    "A woman is carrying groceries up a flight of stairs .",
    "A boy is kicking a football against a wall .",
    "A couple is dancing on an empty street .",
    "A chef is preparing food in an outdoor kitchen .",
    "A cat is sleeping on a warm windowsill .",
    "Workers are repairing a road in the rain .",
    "A small boat is sailing on a calm lake .",
    "Firefighters are climbing a tall ladder .",
]

# train BPE on the words appearing in the time-machine set itself plus train corpus
all_words = TRAIN_CORPUS[:]
for s in TIME_MACHINE_EN:
    all_words.extend(s.lower().split())
tm_merges = learn_bpe(all_words, num_merges=80)
tm_vocab = bpe_vocab(all_words, tm_merges)
for w in all_words:
    for ch in w:
        tm_vocab.add(ch)
tm_vocab.add(EOW)

print("=== Time machine: BPE segmentation of held-out sentences ===")
total_tokens = 0
total_oov = 0
for s in TIME_MACHINE_EN[:3]:
    pieces = []
    for w in s.lower().split():
        pieces.extend(apply_bpe(w, tm_merges))
    print(s)
    print("  ", " ".join(pieces))
    print()

n_sent = 0
for s in TIME_MACHINE_EN:
    n_sent += 1
    for w in s.lower().split():
        for sym in apply_bpe(w, tm_merges):
            total_tokens += 1
            if sym not in tm_vocab:
                total_oov += 1
print("sentences:", n_sent)
print("avg tokens/sentence: %.1f" % (total_tokens / n_sent))
print("OOV count:", total_oov)

Interactive · Watch BPE learn its merges#

Step through the algorithm one merge at a time. Each click applies the next most-frequent pair merge to the whole corpus — watch subwords crystallise out of characters.