Lab 01 — BPE Tokenizer + Chat-Template Renderer

Phase: 01 — Tokenization & the Multimodal Input Pipeline Difficulty: ⭐⭐⭐☆☆ (the algorithm is small; the invariants are where seniors are made) Time: 3–4 hours

The model never sees your text. It sees integer IDs, and it speaks integer IDs back. The tokenizer is the entire boundary between the human world and the tensor world — and it is the single most under-respected component in the stack, the place where round-trips break, where chat quality silently dies, and where prompt-injection-flavored bugs hide. This lab builds the real thing: learn a byte-level BPE merge table from a corpus, encode/decode with a UTF-8 byte fallback so the codec is closed (no <unk> for normal text, a true decode(encode(x)) == x round-trip), reserve the special/control tokens chat models live and die by, and render a ChatML chat template byte-for-byte the way the model was trained to expect.

What you build

  • train_bpe(corpus, num_merges, *, lowercase=False) — learn the ordered merge list and the vocabulary. Greedy and deterministic: pre-tokenize into words (\S+|\s+), represent each word as a tuple of base byte-tokens, then repeatedly count adjacent pairs, merge the most frequent (tie-break: lexicographically smallest pair), and rebuild. Returns (merges, vocab).
  • BPETokenizer — the trained codec (a @dataclass). Validates that the vocab is closed (all 4 specials + all 256 base byte-tokens present, IDs unique) and builds the id_to_token inverse and the merge_ranks priority table.
    • BPETokenizer.from_training(corpus, num_merges, *, lowercase=False) — train + construct.
    • vocab_size, bos_id, eos_id, pad_id, unk_id — properties.
    • encode(text, *, add_bos=False, add_eos=False) — text → list of IDs, applying merges per word in learned-priority order, with the byte fallback so nothing is ever OOV.
    • decode(ids) — IDs → exact original text; specials are skipped, byte-tokens are mapped back to bytes and the byte stream is UTF-8 decoded.
  • render_chat(messages, *, add_generation_prompt=True) — render a conversation into ChatML: <|im_start|>{role}\n{content}<|im_end|>\n per message, plus an open <|im_start|>assistant\n turn for generation. Must match training byte-for-byte.
  • Module constantsSPECIAL_TOKENS (<pad>, <unk>, <bos>, <eos>), IM_START / IM_END (the ChatML control strings), BYTE_ENCODER / BYTE_DECODER (the GPT-2 byte↔printable-unicode bijection).

Key concepts

ConceptWhat to understand
Subword (BPE)between char-level (long sequences) and word-level (huge vocab, <unk> on OOV); merge frequent pairs to balance both
Byte-level base alphabetstarting from the 256 byte values makes the vocab closed — every possible input decomposes into base tokens, so no <unk> and a true round-trip
Merge priority = encodingthe ordered merge list is the rank table; encoding applies the lowest-rank (highest-priority) present pair first, repeatedly
Deterministic tie-breakties on pair frequency are broken by the lexicographically smallest pair, so the result never depends on dict iteration order
Special tokens are reservedstructural IDs added by add_bos/add_eos; the literal text <bos> must never tokenize to the reserved id — the byte path guarantees this
Whitespace is a tokensplitting on \S+|\s+ keeps space/newline runs as their own words, which is what makes spaces survive decode(encode(x))
Chat template is exactChatML boundary tokens must be byte-identical to training; a stray space or missing newline silently degrades quality

Files

FilePurpose
lab.pyskeleton with # TODO markers — your implementation
solution.pycomplete reference; python solution.py runs a worked example
test_lab.pythe proof — run it red, make it green
requirements.txtpytest only (pure stdlib otherwise)

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

  • All tests pass against your implementation (and LAB_MODULE=solution).
  • decode(encode(x)) == x holds for ASCII, multiple spaces, newlines/tabs, accents (café), CJK (你好世界), and emoji — the soul test. You can explain why the byte fallback makes this total.
  • You can hand-count the first merge of the lab corpus and explain why ('a','t') (freq 8) beats ('t','h') (freq 6) — counting the actual most-frequent pair, not the eye-catching one.
  • You can explain why train_bpe(reversed(CORPUS), …) learns the identical merges (the tie-break, not dict order, is the source of determinism).
  • You can explain why encoding the literal string <bos> does not produce the reserved bos_id, and why that property matters for safety.
  • You can explain why render_chat must match the training format byte-for-byte, and name the failure mode when it doesn't.

How this maps to the real stack

The miniatureThe production mechanismWhere to verify it
train_bpe greedy mergesthe BPE training in HF tokenizers and Google SentencePiece (--model_type=bpe); the merge-frequency loop is the same algorithmtokenizers.models.BPE + a BpeTrainer; spm_train
byte-level base + byte fallbackGPT-2/GPT-3/GPT-4 byte-level BPE and tiktoken's closed vocab — no <unk>, any byte string round-tripstiktoken.get_encoding("cl100k_base"); the GPT-2 bytes_to_unicode
merges as the rank tablethe merges.txt file shipped with every HF BPE tokenizer; encoding applies them by rankAutoTokenizer.from_pretrained(...).backend_tokenizer
SPECIAL_TOKENS reserved IDstokenizer.add_special_tokens, bos_token_id, eos_token_id, pad_token_idAutoTokenizer special-token attributes
render_chat (ChatML)tokenizer.apply_chat_template(messages, add_generation_prompt=True), driven by a Jinja chat_template in tokenizer_config.jsonHF chat templates; OpenAI ChatML

Limits of the miniature (be honest in the interview): real tokenizers use a more elaborate GPT-2/GPT-4 pre-tokenization regex (splitting digits, contractions, and leading spaces specially) instead of plain \S+|\s+; production trainers prune rare words and cap vocabulary at the target size (32k–128k) rather than a fixed merge count; SentencePiece operates on a unigram language model as an alternative to BPE and treats the input as a raw character stream with a meta-space; and special tokens in real systems get dedicated embedding rows the model attends to, which this lab models only as reserved string IDs.

Extensions (build these on real hardware / real libs)

  • Swap the \S+|\s+ splitter for the actual GPT-2 / cl100k_base pre-tokenization regex and re-measure compression (tokens per character) on a paragraph of English.
  • Add a train_to_vocab_size(corpus, target) that stops when len(vocab) == target instead of a fixed num_merges, and plot compression vs vocab size — find the knee.
  • Train on a multilingual corpus and measure the fertility (tokens per word) gap between English and a low-resource language; explain the fairness/cost implication.
  • Build the <image> placeholder path: reserve an image special token, expand it to k patch IDs, and show how a vision tokenizer slots patch embeddings into the same sequence (forward reference to Phase 04).
  • Compare your merges against tiktoken on the same text: tiktoken.get_encoding("gpt2").

Interview / resume

  • Talking points: "char vs word vs subword — what does BPE actually buy you?" "Why does byte-level BPE have no <unk>?" "Walk me through one BPE merge step and the tie-break." "Why must a chat template match the training format byte-for-byte, and what breaks if it doesn't?" "What is a glitch token and where do they come from?"
  • Resume bullet: Implemented a byte-level BPE tokenizer from scratch — deterministic merge-training with lexicographic tie-break, a closed UTF-8 byte-fallback codec with a verified decode(encode(x)) == x round-trip across ASCII/CJK/emoji, reserved special/control tokens, and a byte-exact ChatML chat-template renderer.