Nucleus Sampling Cutoff

Problem

Return the number of top tokens kept by nucleus sampling.

Sort the probabilities in descending order and count how many are needed for their cumulative sum to reach top_p, keeping at least one token and at most all of them. An empty distribution returns 0.

Hint

Top-p adapts the candidate set size to how concentrated the distribution is.

Test cases

case 1

in: [[0.5,0.3,0.2],0.6]

expected: 2

case 2

in: [[0.5,0.3,0.2],0.5]

expected: 1

case 3

in: [[0.5,0.3,0.2],1]

expected: 3

case 4

in: [[],0.9]

expected: 0

Starter codePython

def nucleus_cutoff(probs, top_p):
    # Your code here
    pass

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