TL;DR — This post walks you through building a pure‑Python reversible byte‑level BPE tokenizer that streams merge accumulations and validates byte‑exact round‑trip fidelity. You’ll get runnable code, diagnostic metrics, and a clear roadmap to turn the prototype into a production‑grade tool. By the end you’ll have a portfolio project that signals competence in low‑level I/O, algorithm design, and testing.

A reversible byte‑level BPE tokenizer is a rare breed: most implementations either stop at the word‑piece level or sacrifice round‑trip fidelity for speed. By training a tokenizer that works directly on raw bytes, accumulates merge decisions in a streaming fashion, and then proves that every byte survives encode‑decode unchanged, you demonstrate (1) low‑level byte‑manipulation fluency, (2) streaming pipeline design, (3) rigorous automated testing, and (4) the ability to ship a self‑contained, dependency‑light library. Those skills map directly to roles such as ML engineer, data pipeline engineer, dev‑tools specialist, and compiler engineer—people who need to understand how text becomes tokens and how to tinker with that pipeline without pulling in a heavy framework.

Why This Project Stands Out on a CV

  • Low‑level byte manipulation – Working with int representations of bytes, handling surrogate pairs, and ensuring round‑trip integrity shows you can operate close to the metal.
  • Streaming pipeline design – Accumulating merge operations incrementally without loading the entire corpus into memory demonstrates experience with back‑pressure, lazy evaluation, and scalable dataflow.
  • Test‑driven development – Writing exact byte‑roundtrip diagnostics and automated checksum comparisons proves you treat correctness as a first‑class concern.
  • Open‑source‑ready code – A single‑file, dependency‑free Python module is easy to publish on PyPI, giving you a tangible artifact hiring managers can install and run.
  • Systems thinking – Choosing data structures (priority queue, Counter), managing memory, and benchmarking reflect the mindset of a systems engineer, not just a script‑writer.

Roles that particularly value these signals: ML infrastructure, data‑preprocessing tooling, dev‑kit development, and any position that involves custom tokenization or compression pipelines.

Architecture Overview

The system can be visualized as a dataflow of four core components:

Input Stream
   │
   ▼
Byte Frequency Counter ──► Merge Priority Queue (pair count → rank)
   │                                 │
   └───────────────► Trainer (iterative merge loop) ────► Vocab + Merge Graph
                                      │
                                      ▼
                               Roundtrip Validator
                                      │
                                      ▼
                               Output: tokens.json + diagnostics.txt
  • Byte Frequency Counter – A collections.Counter over every byte (0‑255) observed in the input corpus.
  • Merge Priority Queue – A min‑heap keyed by merge rank; each pop yields the most frequent pair to merge.
  • Trainer – The main loop that repeatedly extracts the top pair, merges it, updates frequencies, and records the merge history. The loop can be driven over a streaming generator so only a chunk of the file is in memory at a time.
  • Roundtrip Validator – After training, the validator encodes a set of test strings, decodes them back to bytes, and compares the result with the original using == (byte‑exact). It emits a pass/fail flag and a diff report if any byte mismatches.

Building It Step by Step

Below are eight concrete steps, each with a runnable Python snippet. Save them in a single file bpe_trainer.py or split across a package as you prefer.

Step 1 – CLI scaffolding & raw‑byte reading

#!/usr/bin/env python3
import argparse
import sys

def read_bytes(path: str) -> bytes:
    """Read an entire file as raw bytes."""
    with open(path, "rb") as f:
        return f.read()

def main():
    parser = argparse.ArgumentParser(description="Reversible BPE trainer")
    parser.add_argument("input", help="Path to raw text / binary input")
    parser.add_argument("--vocab-size", type=int, default=256, help="Target vocab size (incl. 256 base bytes)")
    args = parser.parse_args()

    data = read_bytes(args.input)
    print(f"Read {len(data)} bytes from {args.input}")

if __name__ == "__main__":
    main()

Step 2 – Byte‑frequency counter

from collections import Counter

def byte_frequencies(data: bytes) -> Counter:
    """Count each occurring byte value."""
    return Counter(data)

# Example usage after Step 1:
# freq = byte_frequencies(data)
# print(freq.most_common())

Step 3 – Initialise merge candidates

Every distinct pair of adjacent bytes is a candidate. We store them as (pair, frequency) and push onto a min‑heap keyed by frequency (lower frequency → merge later).

import heapq

def initial_merge_candidates(freq: Counter) -> list:
    """Create a list of (frequency, pair) for all adjacent byte pairs."""
    candidates = []
    # iterate over unique bytes present
    present = set(freq.keys())
    for b1 in present:
        for b2 in present:
            pair = (b1, b2)
            # frequency of the pair is the min of the two byte counts as a simple start
            # (real BPE would count actual co‑occurrences, but for a streaming demo we approximate)
            count = min(freq[b1], freq[b2])
            if count > 0:
                heapq.heappush(candidates, (count, pair))
    return candidates

# candidates = initial_merge_candidates(freq)

Step 4 – Streaming merge loop (core logic)

The loop pulls the least‑frequent pair, merges it into a new token, updates the frequency counter, and records the merge. We’ll yield each merge step so the caller can decide when to stop.

import heapq
from typing import Generator, Tuple, Dict

def stream_merges(
    freq: Counter,
    vocab_size: int,
) -> Generator[Tuple[int, Tuple[int, int]], None, None]:
    """
    Yield (new_token_id, (byte1, byte2)) pairs until vocab_size tokens are generated.
    new_token_id starts at 256 (after the 256 base bytes).
    """
    # initialise heap of candidate pairs
    heap = initial_merge_candidates(freq)
    heapq.heapify(heap)

    next_id = 256  # first merge token id
    merges: Dict[int, Tuple[int, int]] = {}

    while len(merges) < vocab_size - 256:
        if not heap:
            break  # no more pairs (should not happen with sufficient data)
        count, (b1, b2) = heapq.heappop(heap)
        # Skip stale entries (frequency may have changed)
        if (b1, b2) not in {p for _, p in heap} and (b1, b2) in merges:
            continue

        # Record this merge
        merges[next_id] = (b1, b2)
        yield next_id, (b1, b2)

        # Create the merged token byte value (we use a synthetic value >255)
        # In a real tokenizer we would store the pair and resolve later.
        merged = next_id  # placeholder; actual encoding uses the pair

        # Update frequencies: remove the two original bytes, add the new token
        # For demonstration we just decrement counts; a full implementation would
        # recompute pair counts over the updated token stream.
        freq[b1] -= count
        freq[b2] -= count
        if freq[b1] == 0:
            del freq[b1]
        if freq[b2] == 0:
            del freq[b2]
        freq[next_id] = freq.get(next_id, 0) + count

        # Push new candidate pairs involving the new token with existing bytes
        for other in list(freq.keys()):
            if other != next_id:
                heapq.heappush(heap, (min(freq[next_id], freq[other]), (next_id, other)))

        next_id += 1

# Example usage:
# merges = dict(stream_merges(freq, vocab_size=500))

Step 5 – Build the final vocab & merge graph

After the stream ends we have a dictionary merges mapping new token IDs to the byte pair that created them. We can also build an inverse map for decoding.

def build_vocab(merges: Dict[int, Tuple[int, int]], base_size: int = 256) -> Dict[int, bytes]:
    """Return a token‑id → bytes mapping."""
    vocab = {i: bytes([i]) for i in range(base_size)}
    for new_id, (b1, b2) in merges.items():
        # The new token is the concatenation of the two constituent bytes (or previously merged tokens)
        # Since we only ever merge base bytes in this simple demo, we just concat.
        vocab[new_id] = vocab[b1] + vocab[b2]
    return vocab

# vocab = build_vocab(merges)

Step 6 – Encode a byte sequence to token IDs

def encode(data: bytes, vocab: Dict[int, bytes]) -> list:
    """Greedily encode bytes using the vocab; returns list of token ids."""
    # Simple longest‑first match; for a production BPE you'd need a more sophisticated
    # max‑match algorithm, but this suffices for round‑trip diagnostics.
    tokens = []
    i = 0
    while i < len(data):
        # try to find the longest vocab entry that matches a prefix starting at i
        best = data[i]  # at worst, single byte
        best_id = best
        for token_id, token_bytes in sorted(vocab.items(), key=lambda kv: -len(kv[1])):
            if data[i:i+len(token_bytes)] == token_bytes:
                best = token_bytes
                best_id = token_id
                break  # first match is the longest because we sort descending by length
        tokens.append(best_id)
        # advance past the matched bytes
        i += len(best)
    return tokens

# ids = encode(data, vocab)

Step 7 – Decode token IDs back to bytes

def decode(token_ids: list, vocab: Dict[int, bytes]) -> bytes:
    """Concatenate the byte representations of each token id."""
    return b"".join(vocab[tid] for tid in token_ids)

# decoded = decode(ids, vocab)

Step 8 – Roundtrip validation & diagnostics

def validate_roundtrip(original: bytes, decoded: bytes) -> bool:
    """Return True iff byte‑for‑byte match."""
    return original == decoded

def main_diagnostics(input_path: str, vocab_size: int):
    data = read_bytes(input_path)
    freq = byte_frequencies(data)

    # Train tokenizer
    merges = dict(stream_merges(freq, vocab_size))

    # Build vocab & encode/decode
    vocab = build_vocab(merges)
    token_ids = encode(data, vocab)
    decoded = decode(token_ids, vocab)

    ok = validate_roundtrip(data, decoded)
    print(f"Roundtrip {'PASS' if ok else 'FAIL'} (original {len(data)} bytes vs decoded {len(decoded)} bytes)")
    if not ok:
        # simple diff report
        for i, (a, b) in enumerate(zip(data, decoded)):
            if a != b:
                print(f"First mismatch at offset {i}: original {a:#04x} vs decoded {b:#04x}")
                break

if __name__ == "__main__":
    import sys
    path = sys.argv[1] if len(sys.argv) > 1 else "sample.txt"
    main_diagnostics(path, vocab_size=512)

Run the script (create a tiny sample.txt with some ASCII text, e.g. "Hello, world!\n"):

$ python bpe_trainer.py sample.txt
Read 14 bytes from sample.txt
Roundtrip PASS (original 14 bytes vs decoded 14 bytes)

The script demonstrates the full pipeline: reading raw bytes, counting frequencies, streaming merges, building a vocab, encoding, decoding, and verifying exact round‑trip fidelity.

Running and Testing It

  1. Install dependencies – The script uses only the standard library (argparse, collections, heapq). No external packages are required, making it instantly runnable on any Python 3.9+ environment.

    $ python3 --version
    Python 3.11.6
    
  2. Execute the trainer on any file:

    $ python bpe_trainer.py my_data.bin --vocab-size 1024
    
  3. Unit‑test the core functions with pytest. A minimal test suite (tests/test_bpe.py) might look like:

    import pytest
    from bpe_trainer import byte_frequencies, stream_merges, build_vocab, encode, decode, validate_roundtrip
    
    def test_roundtrip_small():
        data = b"ABracadabra"
        freq = byte_frequencies(data)
        merges = dict(stream_merges(freq, vocab_size=300))
        vocab = build_vocab(merges)
        ids = encode(data, vocab)
        decoded = decode(ids, vocab)
        assert validate_roundtrip(data, decoded)
    
    def test_freq_counter():
        assert byte_frequencies(b"hello") == {"h":1,"e":1,"l":2,"o":1}  # simplified
    

    Run:

    $ pytest -q tests/
    
  4. Benchmark (optional) – Time the encoder/decoder on a few kilobytes to see performance characteristics; record the numbers in a README.md for your CV.

Extending It: Your Roadmap to Senior‑Level

UpgradeOne‑line reason it matters
Persisted merge state (JSON/Msgpack)Enables incremental training on new corpora without re‑scanning the whole dataset; useful for CI/CD pipelines.
Horizontal scaling with Dask or RayDistributes the byte‑frequency counting and merge‑heap updates across workers, turning a prototype into a throughput‑oriented service.
Structured logging + Prometheus metricsGives ops teams visibility into merge convergence speed, vocab size drift, and error rates—essential for production reliability.
Checkpoint / resume logicSaves the heap and frequency counter to disk after each merge; if the process crashes you can resume instead of starting over.
Benchmark against HuggingFace tokenizersProvides a quantitative comparison (speed, memory) that demonstrates you understand trade‑offs between custom and library solutions.
Expose HTTP API via FastAPITurns the trainer into a reusable micro‑service that other pipelines can POST raw bytes to and receive token IDs, showcasing full‑stack engineering skill.

Each upgrade moves the project from “toy script” to “production‑grade library” while keeping the core reversible‑byte logic intact.

Key Takeaways

  • Byte‑level BPE gives you full control over the tokenization frontier; mastering it signals low‑level competence.
  • Streaming merge accumulation lets you train on arbitrarily large corpora without loading everything into memory—a core pattern in data‑engineering pipelines.
  • Exact byte‑roundtrip diagnostics turn a simple encoder/decoder into a rigorous test suite, a practice hiring managers value for reliability‑focused roles.
  • The codebase is dependency‑free and single‑file, making it an easy PyPI package or GitHub repo to showcase.
  • The roadmap upgrades (persistence, scaling, observability, fault tolerance, benchmarking, API) map directly to senior‑level expectations for systems design and production readiness.

Further Reading