Typical-P Sampling

Problem

Filter logits with typical-p sampling. Compute a stable softmax and its entropy H, rank tokens by | -log p - H | ascending (ties by index), keep the smallest prefix whose cumulative probability reaches p, and return the kept indices in ascending order.

The signature is typical_p_filter(logits, p).

Hint

Typical sampling keeps tokens whose surprise is close to the distribution's entropy.

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.5]

expected: [0,1]

case 5

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

expected: [0,1,2]

Starter codePython

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

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