Exercise 09 — Beam Search

Timebox: 45 minutes

Goal

Implement beam search decoding for a toy autoregressive scorer. Common at companies that hire from NLP roots.

Work in

• beam_search.py

Tasks

1. Take a score_fn(prefix) -> dict[token, log_prob] so you can swap in a real model later.
2. Maintain beam_width partial sequences with cumulative log-probs.
3. Expand each beam with all candidate tokens, keep the top beam_width by score.
4. Stop when all beams emit EOS or reach max_len.
5. Return all completed sequences sorted by normalized score (length-normalize to avoid short-sequence bias).

Done when

• Hand-built score_fn covers a small vocab and you can verify the top beam by inspection
• Length normalization is implemented (mention α in your README writeup)
• You can articulate beam search vs greedy vs sampling tradeoffs

Stretch

• Diverse beam search (penalize beams that share tokens)
• Coverage penalty for sequence-to-sequence