HardNLPnlp-287

Beam Top-k After Two Steps

Problem

Run beam search over probability rows and return the surviving beams. vocab aligns with each row of step_probs, beams start as one empty hypothesis with log score 0.0, and every beam is extended by every token with positive probability.

After each step keep the top beam_width hypotheses ranked by score descending and then by token sequence lexicographically. Return all surviving beams as [[tokens], score] pairs in that order.

Hint

Scores add in log space, and sequences break ties lexicographically.

Test cases

case 1

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

expected: [[["a","a"],-0.6161861394238171],[["b","a"],-1.0216512475319812]]

case 2

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

expected: [[["a","a"],-0.6161861394238171]]

case 3

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

expected: [[["a","a"],-0.6161861394238171],[["b","a"],-1.0216512475319812],[["a","b"],-2.8134107167600364],[["b","b"],-3.2188758248682006]]

case 4

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

expected: [[["a","a"],0]]

case 5

in: [["a","b"],[[0,0]],2]

expected: []

Starter codePython

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

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