ML Grind

BPE Tokenization - Transformer Architecture

Implement BPE Tokenization in Python with NumPy. medium exercise with examples, hints, and instant feedback. No setup or signup required.

Byte-Pair Encoding (BPE) Training

What you'll build

The training loop of BPE: starting from characters, repeatedly merge the most frequent adjacent pair of tokens and record the merge rules in order.

Background

Tokenization is the first step of every LLM pipeline (text -> token ids -> embeddings). Word-level vocabularies can't handle unseen words; character-level sequences are long and slow. BPE (Sennrich et al., 2016; byte-level in GPT-2, LLaMA and most LLMs) sits in between: frequent words become single tokens and rare words split into known subwords, so there is never an unknown token. The ordered merge list is the tokenizer: to encode new text you apply the merges in the same order.

Algorithm

  1. Split every word into its tokens (initially characters, plus the end-of-word marker </w>).
  2. Count every adjacent pair of tokens across the corpus, weighted by the word's frequency.
  3. Pick the most frequent pair (a, b) and append it to merges.
  4. Replace every occurrence of a b with the single token ab in every word, scanning left to right (so a a a with the pair (a, a) becomes aa a).
  5. Repeat num_merges times, or stop early when no word has two or more tokens left.

Inputs

  • word_freqs: dict mapping a space-separated token sequence to its count, e.g. {"h u g </w>": 10, "p u g </w>": 5}
  • num_merges: int >= 0, the maximum number of merges

Expected Return Format

A list of (token_a, token_b) tuples of strings, in the order they were merged. Its length is num_merges, or less if the corpus ran out of pairs.

Edge cases and conventions

  • Ties: when several pairs share the top count, pick the one seen first when you scan the words in dict order and each word left to right. max(counter, key=counter.get) on a collections.Counter filled in that order does exactly this. (In the example, ('u', 'g') and ('g', '</w>') are both 17; ('u', 'g') appears first.)
  • Recount the pairs after every merge: merges change which pairs exist.

Constraints

Plain Python (collections.Counter is fine); no tokenizer libraries.

Follow-up

  • Write encode(word, merges): apply the merges in order. Why must encoding use the training order rather than "longest token first" (see Greedy Longest-Match Tokenizer)?
  • Byte-level BPE starts from 256 bytes instead of characters. What problem does that solve?

Examples

Input: bpe_train({"h u g </w>": 10, "p u g </w>": 5, "b u g </w>": 2}, 3)
Output: [('u', 'g'), ('ug', '</w>'), ('h', 'ug</w>')]

('u','g') and ('g','</w>') tie at 17 and ('u','g') is seen first; then 'ug </w>' (17), then 'h ug</w>' (10).

Input: bpe_train({"a b </w>": 1, "c d </w>": 1}, 5)
Output: [('a', 'b'), ('ab', '</w>'), ('c', 'd'), ('cd', '</w>')]

Every word is a single token after 4 merges, so training stops early.

Practice in Python

from collections import Counter

def bpe_train(word_freqs, num_merges):
    """
    Train a BPE tokenizer by learning merge rules.

    Args:
        word_freqs: Dict mapping space-separated token sequences to frequency,
                    e.g. {"h u g </w>": 10, "p u g </w>": 5}
        num_merges: Maximum number of merge operations

    Returns:
        merges: List of (token_a, token_b) tuples, in the order merged
                (ties: the pair seen first, scanning words in dict order)
    """
    # Your code here
    pass

The interactive editor loads with Python and NumPy. Run your implementation against 6 tests, or use the progressive hints.

More exercises in this section →