Greedy Decoding with EOS

Problem

Consume a fixed sequence of next-token logit rows. At each row choose the largest logit, breaking ties by smallest token ID. Stop immediately when EOS is chosen and exclude EOS from the returned generated IDs. Return [] when no rows are supplied.

def greedy_decode(logit_rows, eos_id):

Constraints

• 0..32 rows; each nonempty row has the same vocabulary size V in 1..32.
• eos_id is in [0,V) when rows exist; all logits are finite.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      0,
      2,
      1
    ],
    [
      0,
      0,
      4
    ],
    [
      5,
      0,
      0
    ]
  ],
  2
]
Expected result: [
  1
]

Arguments: [
  [
    [
      2,
      2,
      0
    ]
  ],
  2
]
Expected result: [
  0
]