S106 · Tokenization — BPE, WordPiece, SentencePiece

# mini_bpe.py — a minimal BPE trainer + encoder.
# Trains on a text file, then encodes/decodes sample text.
# Also compares tokens-per-word against tiktoken (GPT-4 tokenizer).
import sys
import re
from collections import Counter
from pathlib import Path
 
VOCAB_TARGET = 500  # small for demo; production uses 30K-100K
END_OF_WORD = "</w>"
 
 
def tokenise_words(text: str) -> list[str]:
    """Simple pre-tokenizer: split on whitespace + punctuation boundaries."""
    return re.findall(r"[a-zA-Z]+|[^a-zA-Z\s]", text.lower())
 
 
def get_stats(splits: dict[tuple[str, ...], int]) -> Counter:
    """Count all adjacent symbol pairs across all words, weighted by word frequency."""
    pairs: Counter = Counter()
    for word, freq in splits.items():
        for i in range(len(word) - 1):
            pairs[(word[i], word[i + 1])] += freq
    return pairs
 
 
def merge_pair(splits: dict[tuple[str, ...], int], pair: tuple[str, str]) -> dict[tuple[str, ...], int]:
    """Merge all occurrences of `pair` in every word split."""
    a, b = pair
    merged = a + b
    new_splits: dict[tuple[str, ...], int] = {}
    for word, freq in splits.items():
        new_word: list[str] = []
        i = 0
        while i < len(word):
            if i < len(word) - 1 and word[i] == a and word[i + 1] == b:
                new_word.append(merged)
                i += 2
            else:
                new_word.append(word[i])
                i += 1
        new_splits[tuple(new_word)] = freq
    return new_splits
 
 
def train_bpe(corpus_path: str, vocab_target: int) -> tuple[list[str], list[tuple[str, str]]]:
    text = Path(corpus_path).read_text(encoding="utf-8")
    words = tokenise_words(text)
    counts = Counter(words)
    print(f"unique words: {len(counts):,}")
 
    # Represent each word as tuple(chars) + end-of-word marker.
    splits: dict[tuple[str, ...], int] = {
        tuple(list(w) + [END_OF_WORD]): freq for w, freq in counts.items()
    }
 
    # Base vocab: all chars seen + EOW.
    vocab: set[str] = {END_OF_WORD}
    for word in splits:
        vocab.update(word)
    print(f"base vocab (chars): {len(vocab)}")
 
    merges: list[tuple[str, str]] = []
    while len(vocab) < vocab_target:
        stats = get_stats(splits)
        if not stats:
            break
        best_pair, best_count = stats.most_common(1)[0]
        if best_count < 2:
            break
        splits = merge_pair(splits, best_pair)
        merges.append(best_pair)
        vocab.add(best_pair[0] + best_pair[1])
        if len(merges) % 50 == 0:
            print(f"  merge {len(merges):4d}: {best_pair} ×{best_count}  vocab={len(vocab)}")
    return sorted(vocab), merges
 
 
def encode(text: str, merges: list[tuple[str, str]]) -> list[str]:
    """Apply merges in order to a new piece of text."""
    words = tokenise_words(text)
    encoded: list[str] = []
    for word in words:
        pieces: list[str] = list(word) + [END_OF_WORD]
        # Apply merges greedily in order.
        for a, b in merges:
            merged = a + b
            i = 0
            while i < len(pieces) - 1:
                if pieces[i] == a and pieces[i + 1] == b:
                    pieces = pieces[:i] + [merged] + pieces[i + 2 :]
                else:
                    i += 1
        encoded.extend(pieces)
    return encoded
 
 
def compare_to_gpt4(text: str, our_tokens: list[str]) -> None:
    try:
        import tiktoken
    except ImportError:
        print("(pip install tiktoken to compare against GPT-4 tokenizer)")
        return
    enc = tiktoken.get_encoding("cl100k_base")  # GPT-4
    gpt4_tokens = enc.encode(text)
    words = len(text.split())
    print(f"\nText: {text[:60]!r}...")
    print(f"  words:           {words}")
    print(f"  our tokens:      {len(our_tokens)}  ({len(our_tokens)/words:.2f} per word)")
    print(f"  GPT-4 tokens:    {len(gpt4_tokens)}  ({len(gpt4_tokens)/words:.2f} per word)")
 
 
def main(corpus_path: str) -> None:
    vocab, merges = train_bpe(corpus_path, VOCAB_TARGET)
    print(f"\nfinal vocab={len(vocab)}  merges={len(merges)}")
 
    demo = "The quick brown fox jumps over the lazy dog. Tokenization matters."
    tokens = encode(demo, merges)
    print(f"\nDemo encode: {demo}")
    print(f"  → {tokens[:30]}{'...' if len(tokens) > 30 else ''}")
    print(f"  → {len(tokens)} tokens for {len(demo.split())} words")
    compare_to_gpt4(demo, tokens)
 
    # Try something non-English if the corpus has it.
    hindi = "नमस्ते दुनिया"
    print(f"\nNon-English test: {hindi!r}")
    compare_to_gpt4(hindi, encode(hindi, merges))
 
 
if __name__ == "__main__":
    main(sys.argv[1])