How to Build a BPE Tokenizer from Scratch in Python and Rust

Building a BPE tokenizer from scratch requires initializing a 256-token byte alphabet, iteratively merging the most frequent adjacent pairs, and applying learned merges deterministically during encoding, as implemented in the rohitg00/ai-engineering-from-scratch repository.

The rohitg00/ai-engineering-from-scratch repository provides complete from-scratch implementations of byte-pair encoding (BPE) tokenization in both Python and Rust. Whether you're learning the algorithm for educational purposes or optimizing for production inference, understanding how to build a BPE tokenizer from scratch centers on three core components: byte-level vocabulary initialization, corpus statistics collection, and greedy merge operations.

Core Architecture of BPE Tokenization

Byte-Level Vocabulary Initialization

Every BPE tokenizer begins with a guaranteed-reversible byte alphabet. According to the source code in phases/19-capstone-projects/30-bpe-tokenizer-from-scratch/code/main.py, the Python implementation initializes the vocabulary in initialize() (lines 56-58) by mapping the first 256 token IDs directly to raw bytes 0-255. This design ensures any UTF-8 string can be represented without unknown tokens.

The Rust implementation in phases/10-llms-from-scratch/01-tokenizers/code/bpe.rs follows the same pattern in BPETokenizer::new() (lines 11-13), inserting these 256 entries into a HashMap<u32, Vec<u8>> to establish the foundation for all subsequent merges.

Special Token Handling

After the byte alphabet, the Python implementation reserves additional IDs for user-defined special tokens such as <|endoftext|> or <|pad|> within the same initialize() method (lines 58-63). The Rust example focuses on the core algorithm without built-in special token support, but they could be inserted into the vocab HashMap immediately after construction.

Training Pipeline Implementation

Pre-tokenization and Chunking

Before counting pairs, text must be split into independent training units that preserve positional information. The Python _pretokenize() function (lines 66-73 in the capstone implementation) splits input into chunks while preserving whitespace runs as distinct units. This separation ensures the decoder can reconstruct the original string exactly without ambiguity about word boundaries.

Counting Adjacent Byte Pairs

The core training statistic identifies the most frequent adjacent byte pairs across the corpus. In Python, _count_pairs() (lines 80-85) iterates over symbols in each chunk and updates a Counter with pair frequencies. The Rust equivalent, get_pairs() (lines 20-26 in bpe.rs), performs the same aggregation using a HashMap to track occurrence counts.

Greedy Merge Selection and Application

Each training iteration selects the highest-frequency pair for merging. In Python, _add_token() (lines 41-47) creates the new token with ID len(vocab) and records the merge operation in a separate table. The Rust implementation computes the new token ID as 256 + i at line 53, updating both the vocab mapping and the merges vector.

The training loop continues until reaching the target vocabulary size. Python's train() function (lines 18-32) and Rust's train() method (lines 43-71) both iterate through the corpus, applying each new merge to every chunk immediately to keep frequency statistics current for the next iteration.

Encoding and Decoding Operations

Encoding with Learned Merges

After training, encoding applies merges in rank order (most frequent first). The Python implementation uses _encode_chunk() (lines 55-78) and the public encode() function (lines 82-94) to convert text into token IDs by repeatedly scanning for mergeable pairs. In Rust, encode() (lines 75-80) walks through self.merges and calls merge_pair for each learned pair in sequence.

Lossless Decoding

Decoding reverses the process by mapping token IDs back to their byte sequences. Python's decode() (lines 20-30) and Rust's decode() (lines 83-89) concatenate the byte vectors from vocab[id], guaranteeing decode(encode(text)) == text for any valid input.

Serialization and Reusability

The Python capstone implementation includes JSON serialization via save() and load() (lines 33-65), enabling trained tokenizers to persist across sessions without retraining. The Rust implementation in bpe.rs focuses on in-memory operations with minimal standard-library dependencies, though the merges vector and vocab HashMap could be serialized using external crates like serde to achieve parity with the Python version.

Complete Working Examples

The following snippets demonstrate the full training, encoding, and decoding workflow for each implementation.

Python (from phases/19-capstone-projects/30-bpe-tokenizer-from-scratch/code/main.py):

from main import BPETokenizer, train, encode, decode

# 1️⃣  Build the tokenizer

tok = BPETokenizer()
train(tok, corpus="the quick brown fox jumps over the lazy dog", target_vocab_size=300)

# 2️⃣  Encode a sentence

text = "the fox is quick and the dog is lazy"
ids = encode(tok, text)                     # -> list of ints

print("Encoded IDs:", ids)

# 3️⃣  Decode back (lossless)

recovered = decode(tok, ids)
print("Decoded text:", recovered)
assert recovered == text

Rust (from phases/10-llms-from-scratch/01-tokenizers/code/bpe.rs):

fn main() {
    // 1️⃣  Initialise tokenizer
    let mut tokenizer = BPETokenizer::new();

    // 2️⃣  Train on a small corpus (30 merges)
    let corpus = "the quick brown fox jumps over the lazy dog";
    tokenizer.train(corpus, 30);

    // 3️⃣  Encode a sentence
    let sentence = "the fox is quick and the dog is lazy";
    let ids = tokenizer.encode(sentence);
    println!("Encoded IDs: {:?}", ids);

    // 4️⃣  Decode back (lossless)
    let decoded = tokenizer.decode(&ids);
    println!("Decoded text: {}", decoded);
    assert_eq!(decoded, sentence);
}

Both implementations produce deterministic merge tables and guarantee round-trip fidelity.

Summary

• BPE tokenizers start with 256 byte tokens to guarantee universal UTF-8 coverage without unknown tokens, as implemented in initialize() (Python) and BPETokenizer::new() (Rust).
• Training iteratively merges the most frequent adjacent pairs until reaching the target vocabulary size, updating the corpus representation after each merge.
• Pre-tokenization preserves whitespace boundaries to ensure lossless reconstruction of the original text spacing.
• Encoding applies learned merges deterministically by rank order, while decoding maps IDs back to byte sequences for exact reconstruction.
• The Python implementation includes JSON serialization for model persistence, while the Rust version offers a minimal-dependency alternative suitable for embedded deployment.

Frequently Asked Questions

What is the advantage of starting with 256 initial tokens?

Starting with 256 tokens (one per byte) ensures that any valid UTF-8 string can be tokenized without encountering unknown characters. This byte-level approach makes BPE language-agnostic and robust to rare Unicode characters or typos, as every possible input decomposes into known byte sequences that gradually form longer tokens during training.

How does BPE handle unknown words or characters?

BPE does not encounter truly unknown characters because the initial vocabulary covers all 256 byte values. Rare characters or misspellings simply decompose into their constituent bytes initially, then gradually form new tokens as the merge table grows during training on domain-specific corpora. This inherent fallback to bytes eliminates the need for <unk> tokens common in word-level approaches.

Why are whitespace runs preserved as separate chunks during pre-tokenization?

Preserving whitespace as distinct pre-tokenization chunks prevents the merge algorithm from conflating word boundaries with internal character patterns. This separation ensures that decoding can reconstruct the exact original spacing, maintaining lossless round-trip fidelity between raw text and tokenized representations while allowing the model to learn specific whitespace-related patterns separately.

Can the Rust implementation save and load tokenizers like the Python version?

The current Rust implementation in phases/10-llms-from-scratch/01-tokenizers/code/bpe.rs focuses on the core algorithm without serialization logic, but the merges vector and vocab HashMap could easily be serialized using libraries like serde_json to achieve parity with Python's save() and load() methods shown in the capstone project.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}