Start from a 256-byte base vocab, count adjacent-pair frequencies, repeatedly merge the most frequent pair into a new token, and record merges in order. Encoding replays merges by rank.
from collections import Counter
def train(corpus: str, vocab_size: int):
ids = list(corpus.encode("utf-8"))
vocab = {i: bytes([i]) for i in range(256)}
merges = {}
next_id = 256
while next_id < vocab_size:
pairs = Counter(zip(ids, ids[1:]))
if not pairs:
break
pair = pairs.most_common(1)[0][0]
merges[pair] = next_id
vocab[next_id] = vocab[pair[0]] + vocab[pair[1]]
ids = merge(ids, pair, next_id)
next_id += 1
return merges, vocab
def merge(ids, pair, new_id):
out, i = [], 0
while i < len(ids):
if i < len(ids) - 1 and (ids[i], ids[i + 1]) == pair:
out.append(new_id)
i += 2
else:
out.append(ids[i])
i += 1
return out
def encode(text, merges):
ids = list(text.encode("utf-8"))
while len(ids) >= 2:
pairs = set(zip(ids, ids[1:]))
candidates = [p for p in pairs if p in merges]
if not candidates:
break
pair = min(candidates, key=lambda p: merges[p]) # lowest merge rank first
ids = merge(ids, pair, merges[pair])
return ids
def decode(ids, vocab):
return b"".join(vocab[i] for i in ids).decode("utf-8", errors="replace")
Narrate the two things interviewers listen for: (1) encode must apply merges in training-rank order, not greedily by frequency in the new text; (2) byte-level start means no <UNK> token is ever needed.
Gotchas that sink candidates: O(N²) per merge if you rescan the whole corpus (use a heap of pair counts + incremental updates); using a Python set instead of a dict for merges; forgetting that merges are ordered; losing byte-level determinism on UTF-8 boundaries. Strong candidates account for GPT-2-style regex pre-tokenization ('s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+) and call out that this is what makes subword tokenization non-trivial.
Follow-ups: Why BPE over word-level or char-level vocab? (OOV handling vs sequence length.) What breaks with multilingual text if you train on English-heavy data? (Token fertility — the same sentence costs 3–5× more tokens in low-resource languages.)