Greedy and beam search decoding for a language model

You are writing the decoding step for a language model. Decoding turns the model's per-step scores into a sequence of token ids.

Part 1: greedy decoding

Implement greedy_decode(scores: list[list[float]]) -> list[int].

• scores[i][j] is the score of token j at time step i. Each step is independent of the others.
• Return one token index per step: the token with the highest score at that step.
• If two or more tokens tie for the highest score, choose the smaller index.
• Scores can be negative. An empty scores returns [].

Example: scores = [[2.0, 5.0, 1.0], [3.0, 3.0, -1.0]] returns [1, 0] (step 1 is a tie, so index 0 wins).

Part 2: beam search

In a real model the scores at a step depend on the tokens already chosen, so the best token now can lead to a worse sequence. Implement:

beam_search(step_fn, beam_width: int, max_len: int, eos: int | None = None) -> tuple[list[int], float]

• step_fn(prefix: tuple[int, ...]) returns a list of V log-probabilities (all <= 0) for the next token, given the tokens chosen so far.
• The score of a sequence is the sum of the log-probabilities of its tokens.
• Keep at most beam_width partial sequences at each step. Return the highest-scoring sequence found and its score.
• A sequence that emits eos is complete (eos is included in it) and is not extended further. Stop after max_len tokens or when no partial sequence remains.
• Break score ties by choosing the lexicographically smaller sequence.
• Raise ValueError if beam_width < 1 or max_len < 0.

Example (probabilities shown; step_fn returns their logs, V = 2, max_len = 2):

With beam_width = 1 (greedy) the result is [0, 0] with probability 0.30. With beam_width = 2 the result is [1, 0] with probability 0.36, score log(0.36).

Constraints

• 1 <= V <= 50,000, 1 <= beam_width <= 32, 0 <= max_len <= 512.
• State the time and space complexity of each function. Calls to step_fn are expensive; say how many your solution makes.

Reconstructed from the reported question; details may differ from what interviewers ask.