Top-P Sampling

Sample the next token with top-$pp$ (nucleus) sampling. Turn the logits into probabilities with a softmax, sort them from largest to smallest, keep every token whose cumulative probability is at most p_cut plus the first token that pushes the total past p_cut, renormalise those probabilities and draw one index with rng.choice(len(logits), p=q) where rng is a numpy.random.Generator.

Implement solve(logits, p_cut, rng).

Returns. Return the sampled token index (an integer). With a small p_cut only the most likely token survives.

Examples

Example 1

Input:

solve([2.0, 1.0, 0.0], 0.5, np.random.default_rng(0))

Output:

0

Example 2

Input:

solve([0.0, 0.0, 0.0, 0.0], 1.0, np.random.default_rng(1))

Output:

2