Implement a Toy BPE Tokenizer

What is Byte Pair Encoding (BPE)?

Traditional word level tokenization creates massive vocabularies and fails on new compound words. Character level tokenization creates extremely long sequences.

Byte Pair Encoding (BPE) is a subword tokenization algorithm that balances vocabulary size and sequence length:

Initial Characters: ['l', 'o', 'w', 'e', 'r', ' ', 'n', 'e', 'w', 'e', 's', 't']
Iteration 1:        Merge most frequent pair ('e', 's') ──► 'es'
Iteration 2:        Merge ('es', 't')                 ──► 'est'
Iteration 3:        Merge ('l', 'o')                  ──► 'lo'

Common words become single tokens (for example "the"), while rare words are split into subword fragments (for example "unhelpfulness" $\to$ ["un", "help", "ful", "ness"]).

BPE Algorithm Steps

1. Initialize Vocabulary: Split text corpus into individual characters, adding an end of word symbol </w>.
2. Count Pair Frequencies: Iterate through the corpus to count frequencies of adjacent symbol pairs.
3. Merge Most Frequent Pair: Find the single most frequent pair $(A, B)$ and replace all occurrences with new token $AB$.
4. Repeat: Repeat steps 2 and 3 until target vocabulary size or merge iteration count is reached.

Python Implementation from Scratch

from collections import defaultdict


class ToyBPETokenizer:

  def __init__(self, num_merges=10):
    self.num_merges = num_merges
    self.merges = {}
    self.vocab = set()

  def _get_stats(self, vocab):
    pairs = defaultdict(int)
    for word, freq in vocab.items():
      symbols = word.split()
      for i in range(len(symbols) - 1):
        pairs[symbols[i], symbols[i + 1]] += freq
    return pairs

  def _merge_vocab(self, pair, v_in):
    v_out = {}
    bigram = " ".join(pair)
    replacement = "".join(pair)
    for word in v_in:
      # Replace adjacent pair with merged representation
      w_out = word.replace(bigram, replacement)
      v_out[w_out] = v_in[word]
    return v_out

  def train(self, corpus):
    # Initialize vocabulary with character spaces
    words = corpus.split()
    vocab = defaultdict(int)
    for word in words:
      vocab[" ".join(list(word)) + " </w>"] += 1

    for i in range(self.num_merges):
      pairs = self._get_stats(vocab)
      if not pairs:
        break

      # Find most frequent symbol pair
      best_pair = max(pairs, key=pairs.get)
      vocab = self._merge_vocab(best_pair, vocab)
      self.merges[best_pair] = i
      print(f'Merge {i + 1}: {best_pair} -> {"".join(best_pair)}')

    return vocab


# Example Usage
corpus = 'low lower lowest newest widest'
tokenizer = ToyBPETokenizer(num_merges=5)
final_vocab = tokenizer.train(corpus)

Say this out loud

BPE tokenization builds subword vocabularies by merging frequent character pairs. It starts with individual characters and iteratively merges the most frequent adjacent pair across a corpus. BPE provides subword representations for LLMs, handling rare words without out of vocabulary errors.

Followups to expect

1. What is Byte-Level BPE (used in GPT-2 / GPT-4)? Running BPE over raw UTF-8 bytes (256 base byte tokens) rather than Unicode characters, allowing the tokenizer to handle any text language, code, or binary data without unknown tokens.
2. What is WordPiece vs SentencePiece? WordPiece (used in BERT) merges pairs based on likelihood maximization rather than frequency. SentencePiece trains directly on raw unsegmented text streams without requiring pre-tokenized white space splits.