Greedy Decoding

Task

Implement greedy_decoding(step_logits: list[list[float]], eos_token_id: int, max_new_tokens: int) -> list[int]. Choose the highest-logit token at each generation step, stopping after EOS, the generation limit, or the available step logits are exhausted.

Requirements

• Each row in step_logits contains the vocabulary logits for one autoregressive generation step.
• At every step, select the token ID with the largest logit.
• Break equal-logit ties by choosing the smaller token ID.
• Append the selected token before checking whether it equals eos_token_id.
• Stop immediately after selecting EOS.
• Generate at most max_new_tokens tokens.
• If fewer logit rows are available, stop after the final available row.
• Return only newly generated token IDs and do not mutate step_logits.
• Return an empty list when max_new_tokens is zero or no step logits are available.

Example

step_logits = [
    [0.1, 0.9, 0.0],
    [0.2, 0.1, 0.8],
    [0.9, 0.0, 0.1],
]

greedy_decoding(step_logits, eos_token_id=2, max_new_tokens=3)
# [1, 2]