← all writing

2025-06-14 · tokenizers · 6 min read

A BPE tokenizer in Rust, with pyo3 bindings

When the fine-tune was almost ready but the tokenizer kept splitting our domain vocabulary, I wrote my own. Here's what that meant.

We were fine-tuning Qwen2.5-Coder-3B for an agent that edits code inside a web IDE. The model was fine. The pipeline was fine. But the tokenizer kept mangling things that matter to us: library names like Unsloth, internal function names, and compound tokens like QLoRA were getting split into pieces the model had to stitch back together every single time.

You can live with that in a chatbot. In a code agent, it shows up as repeated, expensive mistakes: the model generates broken identifiers because the token stream never contains the identifier whole.

How BPE works, quickly

Byte Pair Encoding starts with a vocabulary of single bytes (256 entries) and iteratively merges the most frequent adjacent pair of tokens until you reach your target vocabulary size. Each merge is recorded: ("St", "r") → "Str".

Training looks like this:

// train loop, simplified
let mut vocab = byte_vocab();
while vocab.len() < target_size {
    let mut counts: HashMap<(u32, u32), u64> = HashMap::new();
    for seq in &corpus {
        for pair in seq.windows(2) {
            *counts.entry((pair[0], pair[1])).or_default() += 1;
        }
    }
    let ((a, b), _) = counts.into_iter().max_by_key(|(_, c)| *c).unwrap();
    let new = next_id();
    vocab.insert(new, format!("{}{}", vocab[&a], vocab[&b]));
    for seq in &mut corpus {
        merge_pairs(seq, a, b, new);
    }
}

The naive version is O(merges × corpus). The trick is to count pairs once, then only recount around the pairs you actually changed after each merge, the same delta-tracking idea as incremental parsers. With that, training over a 2 GB corpus of code and docs took us from ~6 hours of wall time down to ~40 minutes on a single machine.

Encoding is a lookup, not a search

The classic mistake is re-running the merge logic at encode time. You don't need to: merges form a deterministic decision tree, so encoding is a greedy left-to-right pass where each position asks "what's the longest merge in my set that starts here?"

// encode: greedy longest-match over the merge table
let mut out = Vec::new();
let mut i = 0;
while i < bytes.len() {
    let mut best = None;
    for len in (1..=max_len).rev() {
        if let Some(&id) = merges.get(&bytes[i..i + len]) {
            best = Some((id, len));
            break;
        }
    }
    match best {
        Some((id, len)) => { out.push(id); i += len; }
        None => { out.push(bytes[i] as u32); i += 1; }
    }
}

With the merge table keyed on byte slices, encode hits roughly 1.2M tokens/second single-threaded. Decode is just string concatenation. That's fast enough to run tokenization inside the fine-tuning data pipeline without it ever becoming the bottleneck.

Domain vocabulary injection

The reason this whole project existed: we pre-seed merges from a domain token list. Before training, we take the vocabulary from the base tokenizer plus our known-critical tokens, build their byte-pair decomposition, and pin those merges so training never splits them again.

For our code corpus this cut average token counts by roughly 30% on code files. Every token saved is context window saved, and for an agent editing long files, that's real money.

pyo3 bindings

The tokenizer is Rust; the training pipeline is Python. Instead of shuttling JSON across a process boundary, I exposed it with pyo3:

#[pymodule]
fn bpetok(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<Tokenizer>()?;
    m.add_function(wrap_pyfunction!(train, m)?)?;
    Ok(())
}

#[pyclass]
struct Tokenizer { inner: BpeTokenizer }

#[pymethods]
impl Tokenizer {
    #[new]
    fn new(model_path: &str) -> PyResult<Self> { /* ... */ }

    fn encode(&self, text: &str) -> Vec<u32> { self.inner.encode(text) }

    fn decode(&self, ids: Vec<u32>) -> String { self.inner.decode(&ids) }

    fn vocab_size(&self) -> usize { self.inner.vocab_size() }
}

Python calls it like any library, no subprocesses, no serialization overhead. The whole crate builds in under 30 seconds and the bindings add maybe 150 lines.

Lessons

I'd do it again only where tokens are money: code, biomedicine, legal. For a general-purpose project, no. The 30% context savings paid for the whole thing in GPU-hours within a month.