HardNLPnlp-078

Beam Search (Two Steps)

Problem

Run beam search over step_probs, a list of probability rows aligned with vocab. Start from one empty hypothesis with score 0, extend every beam by every token with positive probability, and keep the top beam_width hypotheses.

Rank hypotheses by score descending and then by sequence lexicographically, and return the tokens of the best final sequence.

Hint

Scores are log probabilities, so they add as sequences grow.

Test cases

case 1

in: [["a","b"],[[0.6,0.4],[0.9,0.1]],2]

expected: ["a","a"]

case 2

in: [["a","b"],[[0.5,0.5],[0.02,0.98]],1]

expected: ["a","b"]

case 3

in: [["a","b","c"],[[0.5,0.3,0.2],[0.1,0.2,0.7]],2]

expected: ["a","c"]

case 4

in: [["a"],[[1],[1]],1]

expected: ["a","a"]

Starter codePython

import math
def beam_search_two_steps(vocab, step_probs, beam_width):
    # Your code here
    pass

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