Top-P Nucleus Filter

Problem

Nucleus (top-p) sampling filter. Compute the stable softmax over the logits, sort indices by descending probability (ties by smaller index), and keep the smallest prefix whose cumulative probability reaches p. Return the kept indices sorted in ascending order.

The signature is top_p_filter(logits, p).

Hint

Always keep at least the most likely token, even when its probability already exceeds p.

Test cases

case 1

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

expected: [1,2]

case 2

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

expected: [2]

case 3

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

expected: [0,1,2]

case 4

in: [[0,0,0],0.7]

expected: [0,1,2]

case 5

in: [[0,0,0],0.5]

expected: [0,1]

Starter codePython

import math
def top_p_filter(logits, p):
    # Your code here
    pass

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