Implement Top-p (Nucleus) Sampling in NumPy

Nucleus (Top-p) Sampling

In Large Language Model decoding, nucleus sampling selects the smallest set of tokens whose cumulative probability exceeds a threshold p. This technique helps prevent the model from picking low-probability tokens while maintaining diversity.

Task

Implement a function nucleus_sampling(logits, p) that performs the following steps:

1. Compute the probability distribution using the softmax function.
2. Sort the probabilities in descending order.
3. Compute the cumulative sum of the sorted probabilities.
4. Create a mask to identify tokens where the cumulative sum is less than or equal to p (include the first token that exceeds p).
5. Normalize the filtered probabilities and sample a single index based on the resulting distribution.

Function Signature

import numpy as np

def nucleus_sampling(logits: np.ndarray, p: float) -> int:
    # Returns the sampled index
    pass

Constraints

• logits is a 1D array of shape (vocab_size, ).
• 0 < p <= 1.0.
• Use np.exp(logits - np.max(logits)) for numerical stability when calculating softmax.

Examples

# Example 1
logits = np.array([2.0, 1.0, 0.1])
p = 0.9
# Probabilities: [0.659, 0.242, 0.098]
# Sorted: [0.659, 0.242, 0.098]
# Cumulative: [0.659, 0.901, 1.0]
# Tokens kept: indices corresponding to [0.659, 0.242]

Edge Cases

• What happens if p=1.0?
• How does the function handle a single token with a very high logit value?