Top-K Sampling

Sample the next token using top-$kk$ sampling. Keep only the k largest logits, set all others to a very large negative value, turn the result into probabilities with a softmax and draw one token with rng.choice(len(logits), p=p) where rng is a numpy.random.Generator.

Implement solve(logits, k, rng).

Returns. Return the sampled token index (an integer in 0 .. len(logits)-1). With k=1 the largest logit is always returned.

Examples

Example 1

Input:

solve([1.0, 2.0, 3.0], 1, np.random.default_rng(0))

Output:

2

Example 2

Input:

solve([1.0, 2.0, 3.0, 4.0], 2, np.random.default_rng(0))

Output:

3