Greedy Token Generation

Problem

Generate tokens greedily from precomputed logits.

logits_seq is a list where each element is a vector of next-token logits for one generation step. At each step pick the index of the largest value, breaking ties by choosing the lowest index, and stop after max_tokens tokens or when the logits run out.

Hint

argmax with strict greater-than keeps the earliest index on ties.

Test cases

case 1

in: [[[0.1,0.9],[0.7,0.2]],2]

expected: [1,0]

case 2

in: [[[1,2,3]],1]

expected: [2]

case 3

in: [[[5,1]],3]

expected: [0]

case 4

in: [[[0.4,0.4,0.1]],1]

expected: [0]

Starter codePython

def generate_greedy(logits_seq, max_tokens):
    # Your code here
    pass

4 test cases · run your code against them in the editor