Top-p Sampling

The problem, from first principles

Sampling from the full distribution occasionally picks a token from the long tail of unlikely, incoherent options, and top-k sampling uses a fixed number of candidates whether the model is certain or unsure. Nucleus (top-p) sampling adapts: sort tokens by probability and keep the smallest set whose cumulative probability reaches p. When the model is confident the nucleus is a token or two, when it is unsure the nucleus is wide. The kept probabilities are renormalized so they sum to one and sampling proceeds from this truncated distribution.

From theory to code

Implement top_p_filter.

Constraints

• probs is a 1-D probability vector and 0 < p <= 1. Sort descending (stable, ties by lower index first).
• Keep the smallest prefix of the sorted order whose cumulative probability is >= p. At least one token is always kept.
• Return a vector of the same length with the kept tokens renormalized to sum to 1 and every other entry exactly 0.
• Do not modify probs.

Hints