WordPiece › Code Implementation

import subprocess

subprocess.run(
    ["uv", "pip", "install", "tokenizers", "transformers"], capture_output=True
)

from collections import defaultdict


def get_word_freqs(corpus):
    """Count word frequencies in corpus."""
    word_freqs = defaultdict(int)
    for sentence in corpus:
        for word in sentence.lower().split():
            # Strip basic punctuation for simplicity
            word = word.strip(".,!?;:'\"()")
            if word:
                word_freqs[word] += 1
    return word_freqs


def word_to_chars(word):
    """Convert word to initial character-level representation with ## prefix."""
    chars = list(word)
    return [chars[0]] + ["##" + c for c in chars[1:]]


def compute_pair_scores(splits, word_freqs):
    """Compute PMI-based scores for all adjacent token pairs."""
    token_freqs = defaultdict(int)
    pair_freqs = defaultdict(int)

    for word, freq in word_freqs.items():
        split = splits[word]
        for i, token in enumerate(split):
            token_freqs[token] += freq
            if i < len(split) - 1:
                pair = (split[i], split[i + 1])
                pair_freqs[pair] += freq

    total = sum(token_freqs.values())
    scores = {}
    for pair, pair_count in pair_freqs.items():
        a, b = pair
        p_a = token_freqs[a] / total
        p_b = token_freqs[b] / total
        p_ab = pair_count / total
        # WordPiece score = P(AB) / (P(A) * P(B))
        scores[pair] = p_ab / (p_a * p_b)

    return scores


def merge_pair(pair, splits, word_freqs):
    """Merge a token pair across all words in the corpus."""
    a, b = pair
    merged = (
        a + b[2:] if b.startswith("##") else a + b
    )  # Remove ## from b when merging
    new_splits = {}
    for word, freq in word_freqs.items():
        split = splits[word]
        new_split = []
        i = 0
        while i < len(split):
            if i < len(split) - 1 and split[i] == a and split[i + 1] == b:
                new_split.append(merged)
                i += 2
            else:
                new_split.append(split[i])
                i += 1
        new_splits[word] = new_split
    return new_splits


def train_wordpiece(corpus, vocab_size=50):
    """Train a WordPiece vocabulary from scratch."""
    word_freqs = get_word_freqs(corpus)

    # Initialize splits: each word -> character sequence with ## prefix
    splits = {word: word_to_chars(word) for word in word_freqs}

    # Build initial vocabulary from all characters
    vocab = set()
    for word in word_freqs:
        for token in splits[word]:
            vocab.add(token)

    while len(vocab) < vocab_size:
        scores = compute_pair_scores(splits, word_freqs)
        if not scores:
            break
        # Find the pair with the highest PMI score
        best_pair = max(scores, key=lambda p: scores[p])
        # Merge the pair
        a, b = best_pair
        new_token = a + b[2:] if b.startswith("##") else a + b
        splits = merge_pair(best_pair, splits, word_freqs)
        vocab.add(new_token)

    return vocab, splits, word_freqs

sample_corpus = [
    "playing basketball is fun",
    "she is playing piano",
    "players play together",
    "the player played the game",
    "unplayed music is rare",
    "replay the game again",
    "basketball players are athletic",
    "play the replay again",
]

vocab, splits, word_freqs = train_wordpiece(sample_corpus, vocab_size=60)

def wordpiece_tokenize_word(word, vocab):
    """
    Tokenize a single word using greedy longest-match-first.
    Returns list of tokens, or ['[UNK]'] if tokenization fails.
    """
    tokens = []
    remaining = word
    is_first = True

    while remaining:
        # Try to find the longest prefix in vocab
        found = False
        for end in range(len(remaining), 0, -1):
            prefix = remaining[:end]
            candidate = prefix if is_first else "##" + prefix
            if candidate in vocab:
                tokens.append(candidate)
                remaining = remaining[end:]
                is_first = False
                found = True
                break

        if not found:
            return ["[UNK]"]

    return tokens


def wordpiece_tokenize(text, vocab):
    """Tokenize a full sentence."""
    all_tokens = []
    for word in text.lower().split():
        word = word.strip(".,!?;:'\"()")
        if word:
            all_tokens.extend(wordpiece_tokenize_word(word, vocab))
    return all_tokens

import os
import tempfile

from tokenizers import Tokenizer
from tokenizers.models import WordPiece
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.trainers import WordPieceTrainer

# Create a more substantial training corpus
training_corpus = [
    # General English text with varied morphology
    "The scientists discovered remarkable properties in the crystalline structure.",
    "Remarkable discoveries often come from unexpected observations.",
    "She carefully unwrapped the surprisingly beautiful gift.",
    "The reconstruction project was delayed by unforeseen complications.",
    "Playing music requires understanding rhythm and melody.",
    "Unplayed instruments accumulate dust in forgotten corners.",
    "The preprocessing step transforms raw text into clean tokens.",
    "Tokenization algorithms differ in their handling of unknown words.",
    "WordPiece uses pointwise mutual information to score candidate merges.",
    "Byte pair encoding greedily merges the most frequent adjacent pairs.",
    "BERT's vocabulary contains approximately thirty thousand tokens.",
    "The model learned representations for both whole words and subwords.",
    "Subword tokenization handles morphologically rich languages well.",
    "Prefixes and suffixes are naturally captured by subword models.",
    "The training process iteratively expands the vocabulary.",
    "Each merge operation adds one new token to the growing vocabulary.",
    "Greedy tokenization finds the longest matching prefix at each step.",
    "Unknown characters fall back to the UNK token in BERT.",
    "Special tokens mark sentence boundaries and padding positions.",
    "The CLS token aggregates sequence-level information for classification.",
]

# Write corpus to a temp file (HuggingFace trainer expects file paths)
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
    for line in training_corpus:
        f.write(line + "\n")
    corpus_path = f.name

# Initialize WordPiece tokenizer
tokenizer = Tokenizer(WordPiece(unk_token="[UNK]"))
tokenizer.pre_tokenizer = Whitespace()

# Configure trainer
trainer = WordPieceTrainer(
    vocab_size=200,
    special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"],
    min_frequency=1,
    continuing_subword_prefix="##",
)

# Train
tokenizer.train(files=[corpus_path], trainer=trainer)
os.unlink(corpus_path)

test_inputs = [
    "playing basketball and tokenization",
    "the reconstruction of wordpiece algorithms",
    "unplayed music in forgotten corners",
]

# Encode each input
encodings = [(text, tokenizer.encode(text)) for text in test_inputs]

from transformers import BertTokenizer

# Load BERT's pre-trained WordPiece tokenizer
# bert-base-uncased: lowercases all text
# bert-base-cased: preserves capitalization
tokenizer_bert = BertTokenizer.from_pretrained("bert-base-uncased")

example_sentences = [
    "WordPiece tokenization splits words into subwords.",
    "The reconstruction of the word demonstrates morphological splitting.",
    "GPU parallelization dramatically accelerates transformer training.",
]

encodings_bert = [
    tokenizer_bert(sentence, return_tensors=None)
    for sentence in example_sentences
]