Implement beam search decoding

Maintain k live hypotheses; each step, expand every hypothesis by the top candidates, then keep the global top k by cumulative log-prob. Move finished (EOS) hypotheses to a completed pool instead of expanding them.

import heapq

def beam_search(model, bos, eos, k=4, max_len=64, alpha=0.6):
    beams = [(0.0, [bos])]           # (cum_logprob, tokens)
    done = []
    for _ in range(max_len):
        candidates = []
        for score, seq in beams:
            log_probs = model(seq)   # dict/array: token -> logprob
            for tok, lp in topk(log_probs, k):
                candidates.append((score + lp, seq + [tok]))
        beams = []
        for score, seq in heapq.nlargest(2 * k, candidates):
            if seq[-1] == eos:
                done.append((score / (len(seq) ** alpha), seq))  # length norm
            else:
                beams.append((score, seq))
            if len(beams) == k:
                break
        if not beams:
            break
    done += [(s / (len(q) ** alpha), q) for s, q in beams]
    return max(done)[1]

Explain length normalization unprompted: raw log-prob sums penalize long sequences (every token adds negative log-prob), so you divide by len^alpha — otherwise beam search degenerates to short outputs.

Follow-ups: Complexity vs greedy? (k× model calls per step; batch the k hypotheses into one forward pass.) Why do chat LLMs use sampling instead of beam search? (Beam search finds high-likelihood but degenerate/repetitive text for open-ended generation; it wins for closed tasks like translation.) How does beam search interact with a KV cache? (Cache must fork/reorder with the beams — a real serving cost.)