Hard End-to-End

Beam Search

Practice the expand, score and prune steps of beam search using a fixed table of token scores. Keep several partial sequences during the search, then return only the best final sequence.

Your task

Implement beam_search(log_probs, beam_width).

log_probs is a floating-point matrix of shape (T, V), with V > 0. Row t supplies the finite additive scores for token IDs 0 through V-1 at step t. Use these scores directly; do not apply softmax or take logs again.
beam_width is a positive integer limiting how many partial sequences survive each step. Keep fewer if fewer candidates exist.
Return a Python dictionary with "sequences" and "score". Despite its plural name, "sequences" contains one list of T integer token IDs. "score" is that sequence's total score as a Python float, rounded to seven decimal places. Leave log_probs unchanged.

The parameter name reflects log-probability scoring, but the supplied rows need not be normalized probability distributions. The algorithm only adds and compares their scores.

Expand, score, prune

Start with one empty sequence of score 0. For each row of scores:

Visit the surviving beams in their current order. For each beam, try token IDs in ascending order from 0 to V-1.
Append each token to a copy of the sequence. Add its score from the current row to the beam's score.
Sort candidates by descending total score, keeping their existing order when scores tie. Retain the first beam_width candidates.

After T steps, return the first surviving sequence and its score. Round only that final returned score; do not round partial scores while ranking candidates. If T=0, return {"sequences": [], "score": 0.0}. There is no prompt, EOS token, length penalty or early stopping in this exercise.

For rows [[-0.1, -1.0], [-0.2, -0.5]] and beam_width=2, the first step keeps [0] with score -0.1 and [1] with score -1.0. Expanding them gives:

[0, 0]: -0.3
[0, 1]: -0.6
[1, 0]: -1.2
[1, 1]: -1.5

Keep the first two candidates. Return {"sequences": [0, 0], "score": -0.3}. Higher scores are better: -0.3 beats -1.2.

What this simplified setup leaves out

Every beam uses the same score row at a given step. Scores do not depend on the tokens already chosen. As a result, independently choosing the largest score in each row also finds a best sequence here, for any positive beam width. Use beam expansion to practice the requested algorithm, but do not expect a wider beam to improve this particular table's best score.

A language model instead computes next-token scores from each candidate's prefix. In that setting, beam search is an approximate search and does not guarantee the globally best sequence.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Choose the next token using the prefix so far

An autoregressive model predicts a distribution for the next token given the existing sequence. Greedy decoding picks the most likely token and repeats. Beam search keeps several candidate prefixes, expands each with possible next tokens, and retains the best candidates under its scoring rule.

Beam search considers more alternatives but can still discard a prefix that would later lead to the best complete sequence. Keep each candidate's state, score, and stopping status separate. A real model's next-token distribution depends on its prefix; a fixed table of scores by time step is only a simplified exercise.

Sequence log probability is a sum of token log probabilities. Since these terms are nonpositive, shorter sequences can be favored. Length normalization changes candidate rankings, so use the exact rule in the task.

Sample from a restricted distribution

Top-k sampling retains the k highest-scoring tokens. Nucleus, or top-p, sampling sorts tokens by probability and keeps the smallest prefix reaching the requested cumulative probability. Include the token that crosses the threshold.

For probabilities [0.6, 0.3, 0.1] and threshold 0.8, nucleus sampling retains the first two tokens because their total is 0.9. Renormalizing gives retained probabilities [2/3, 1/3]. Map any sampled sorted index back to its original vocabulary ID.

Temperature divides logits before softmax. A positive temperature changes how concentrated the distribution is. Check the task's end-of-sequence handling, maximum length, whether returned tokens include the prompt, and random-key usage.

Evaluate overlap without confusing it with meaning

BLEU-1 measures clipped unigram overlap with a reference and applies a brevity penalty. Repeating a reference word cannot earn unlimited credit. This unigram score ignores word order and does not establish semantic correctness. Follow the specified tokenization and empty-output conventions when computing it.

Detecting runtime…