Beam Search One Step (Beam 2)

Problem

Expand a beam by combining each beam's cumulative score with its per-token log-probabilities: candidate score = beam_scores[b] + token_logprobs[b][t]. Return the top two candidates sorted by score descending, breaking ties by smaller beam index and then smaller token index.

The signature is beam_search_step(beam_scores, token_logprobs). Return a list of [beam_index, token_index, score].

Hint

Enumerate beam-token pairs, add the log-probabilities to the beam scores, then sort with a deterministic tie-break.

Test cases

case 1

in: [[0,-1],[[-0.5,-2],[-0.5,-3]]]

expected: [[0,0,-0.5],[1,0,-1.5]]

case 2

in: [[0,0],[[-1,-1],[-1,-1]]]

expected: [[0,0,-1],[0,1,-1]]

case 3

in: [[-0.1,-0.2],[[-0.1,-5],[-0.05,-5]]]

expected: [[0,0,-0.2],[1,0,-0.25]]

case 4

in: [[-1,-2,-3],[[-1],[-0.5],[-0.1]]]

expected: [[0,0,-2],[1,0,-2.5]]

Starter codePython

def beam_search_step(beam_scores, token_logprobs):
    # Your code here
    pass

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