Beam Search Top-K

Run beam search over precomputed per-step scores. step_scores[t][v] is the score (for example a log-probability) of choosing token v at step t, independent of earlier choices. Keep the k best partial sequences by cumulative score after each step; ties are broken in favour of the lexicographically smaller token sequence. k must be positive.

Implement solve(step_scores, k).

Returns. Return a list of at most k tuples (tokens, score) sorted from best to worst, where tokens is a list of token ids.

Examples

Example 1

Input:

solve([[0.0, -1.0], [-2.0, 0.0]], 2)

Output:

[([0, 1], 0.0), ([1, 1], -1.0)]

Example 2

Input:

solve([[-1.0, 0.5, 0.0]], 1)

Output:

[([1], 0.5)]