Several Drafts at Once

Greedy decoding commits to the single most likely token at every step and never looks back. That can walk into a dead end. A likely first token followed only by lukewarm options can be less probable overall than a less likely first token followed by a near-certain one, and greedy decoding will never find the second path.

Beam search keeps several drafts alive at once.

• model(tokens) returns a list of probabilities for the next token, for token ids 0, 1, 2, ... in order.
• A draft's score is the sum of the natural logs of the probabilities of the tokens it has generated. Adding logs is the same as multiplying the probabilities, but the numbers do not shrink toward zero as the draft grows. The prompt itself is not scored, so the starting draft has score 0.0.
• Begin with a single draft: the prompt.
• Each step: extend every current draft by every token whose probability is above zero (a token with probability 0 is never used). Pool all of these candidates together and keep the beam_width best. If there are fewer candidates than that, keep them all.
• Ranking: a higher score ranks first. If two candidates have exactly the same score, the one whose token list comes first when compared position by position ranks higher (Python's ordinary list comparison).

Task: write beam_search(model, prompt, steps, beam_width). After steps steps, return the kept drafts as a list of (tokens, score) tuples, best first. Here tokens is the full sequence (the prompt followed by the generated tokens) and score is rounded to 4 decimal places.

With beam_width = 1 this is exactly greedy decoding. In the tests, each model is a small hand-written rule rather than a trained model, so that every run is predictable.