Lab 02 — BPE Tokenization & the generate() Decode Loop

Phase 30 · Lab 02 · Phase README · Warmup

The problem

Two of the most cargo-culted lines in applied LLM work are:

tokenizer = AutoTokenizer.from_pretrained("gpt2")
out = model.generate(**inputs, temperature=0.7, top_p=0.9, repetition_penalty=1.2)

Most engineers can call both. Very few can say what a BPE merge actually is, why [PAD] needs an attention mask, what temperature does to the logits mathematically, or where top-k and top-p differ — and those are exactly the questions that separate "used HF once" from "understands the stack" in an interview. This lab makes you build both halves:

1. A trained BPE tokenizer. Byte-Pair Encoding starts from characters and repeatedly merges the most frequent adjacent symbol pair into a new vocabulary entry, up to a merge budget. Frequent words collapse into single tokens; rare words split into learnable pieces; nothing is ever out-of-vocabulary at the character level. You train the merges (with a deterministic tie-break), then implement encode/decode with the four canonical special tokens ([BOS]/[EOS]/[PAD]/[UNK]), padding + truncation + the attention mask (so a model can ignore pad positions), and the offset mapping back to source characters.

2. The generate() loop. Decoding is a loop over next-token logits, and every sampling knob is a small transform on those logits applied in a fixed order: repetition penalty → temperature → softmax → top-k → top-p → choose. You implement each transform and the loop that composes them, driven by a GenerationConfig, over an injected logits_fn(ids) -> logits — with sampling routed through a seeded, injected sampler so every test is reproducible.

What you build

File map

Run it

pytest test_lab.py -v                       # your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # the reference (green)
python solution.py                          # worked example

Success criteria

• Training is deterministic: same corpus + budget → identical merges and vocab, with ties broken lexicographically.
• decode(encode(text)) == text for any text over the training alphabet — including words never seen during training.
• [BOS]/[EOS] wrap encoded sequences; [PAD] fills to max_length with an attention mask of 1…10…0; truncation caps all four aligned lists.
• Offset mapping slices the source text back to each token's visible characters (a merged symbol that absorbed </w> spans only its real chars).
• Greedy decoding picks the argmax chain and temperature=0 is exactly greedy.
• top_k=2 never samples outside the two highest-probability candidates (proven over 20 seeds); top_p keeps the smallest prefix of ranked tokens reaching cumulative mass p.
• The repetition penalty measurably lowers a seen token's probability and can flip the greedy choice.
• Generation stops at both eos_token_id and max_new_tokens, and two runs with the same seed are byte-identical.
• All 30 tests pass under both lab and solution.

How this maps to the real stack

• BPE is the real algorithm behind GPT-2/GPT-4-family tokenizers and the tokenizers library's BpeTrainer — learn merges by pair frequency, apply them in order at encode time. Real HF BPE is byte-level (operates on UTF-8 bytes so any string is representable, no [UNK] ever needed) and uses a pre-tokenizer regex rather than bare whitespace; ours is char-level with an explicit </w> end-of-word marker (the original NMT-BPE formulation) so the merges are readable in tests. WordPiece (BERT) picks merges by likelihood gain instead of raw frequency and marks continuations with ##; Unigram/SentencePiece (T5, Llama) starts big and prunes rather than starting small and merging — the Warmup §5 covers all three.
• Special tokens, padding, truncation, attention mask, offset mapping are the real tokenizer(text, padding="max_length", truncation=True, return_offsets_mapping=True) surface. The attention mask really is "1 = attend, 0 = ignore," and forgetting it on padded batches is a classic silent-quality bug (the model attends to [PAD] embeddings).
• GenerationConfig is a real class (transformers.GenerationConfig), and max_new_tokens, do_sample, temperature, top_k, top_p, repetition_penalty, eos_token_id are its real field names. Our fixed transform order (penalty → temperature → top-k → top-p) mirrors HF's LogitsProcessor pipeline, where each knob is literally a LogitsProcessor/LogitsWarper object applied in sequence inside generate().
• The repetition penalty is the real CTRL-paper formulation HF implements: divide a seen token's logit when positive, multiply when negative — both push it down.
• generate() in real transformers is this exact loop plus the machinery a miniature omits: a KV cache (past attention keys/values reused so each step is O(1) in the prompt, not O(n) — the single biggest inference-speed lever), batched decoding, beam search, StoppingCriteria objects, and streaming. The Warmup covers the KV cache and beam search in prose; the loop shape you built is the part every variant shares.

Limits of the miniature. No KV cache — logits_fn receives the whole sequence each step, which is honest about the interface but not the cost (real decode reuses cached attention state). No beam search (it tracks N hypotheses with length-normalized scores; greedy is the beam=1 case). Real sampling draws on the GPU with a framework RNG; our injected seeded sampler is the same record/replay seam, made explicit. Byte-level BPE, pre-tokenizer regexes, and added-token handling (add_tokens, resize_token_embeddings) are real-world layers we skip.

Extensions (your own machine)

• Add beam search: track the top-B sequences by cumulative log-prob with length normalization; verify beam=1 reproduces greedy.
• Add a KV-cache simulation: change logits_fn to accept (new_id, cache) and return (logits, cache); prove per-step work no longer grows with sequence length.
• Make the BPE byte-level: operate on UTF-8 bytes, drop [UNK] entirely, and confirm any emoji round-trips.
• Compare your tokenizer's compression (tokens per word) against tiktoken or a real HF tokenizer on the same corpus.

Interview / resume signal

"Implemented BPE from scratch — frequency-ranked pair merges with deterministic tie-breaks, reversible encode/decode with special tokens, padding, attention masks, and offset mappings — and the full generate() decode loop: temperature, top-k, top-p nucleus filtering, CTRL-style repetition penalty, and EOS/budget stopping, composed in HF's LogitsProcessor order over an injected logit function with seeded, reproducible sampling."