Nucleus Sampling
Implement top-p (nucleus) sampling from a probability distribution.
Nucleus sampling selects the smallest set of tokens whose cumulative probability exceeds a threshold p, then renormalizes.
Algorithm:
- Sort tokens by probability (descending)
- Compute cumulative probabilities
- Find the smallest set of tokens where cumulative probability >= p
- Renormalize the selected probabilities to sum to 1
Since we need deterministic output, return the renormalized probability distribution (with non-selected tokens set to 0).
Input format:
- Line 1: p threshold (float, 0 < p <= 1)
- Line 2: Token names (space-separated)
- Line 3: Corresponding probabilities (space-separated floats, sum to 1)
Output: A dictionary of token: renormalized_probability for selected tokens (prob > 0), sorted by probability descending. Round to 4 decimal places.
Example:
0.8 the a an this that 0.35 0.25 0.20 0.15 0.05
{'the': 0.4375, 'a': 0.3125, 'an': 0.25}Step 1: Sort by probability (already sorted) the: 0.35, a: 0.25, an: 0.20, this: 0.15, that: 0.05
Step 2: Cumulative probabilities the: 0.35 the + a: 0.60 the + a + an: 0.80 >= 0.8 => STOP
Step 3: Selected tokens: {the, a, an} with probs [0.35, 0.25, 0.20]
Step 4: Renormalize Sum = 0.80 the: 0.35/0.80 = 0.4375 a: 0.25/0.80 = 0.3125 an: 0.20/0.80 = 0.25
Constraints:
- Include the minimum number of tokens to reach cumulative prob >= p
- Renormalize selected probabilities to sum to 1
- Round to 4 decimal places
- Output only tokens with non-zero probability
Background Knowledge
Text Generation is a crucial aspect of Natural Language Processing (NLP), involving the creation of human-like text based on a given context or prompt. One key technique in text generation is sampling from a probability distribution, where the goal is to select the next token (word or character) based on its probability of occurrence. Top-p (nucleus) sampling is a specific method that aims to balance the trade-off between exploration and exploitation by selecting a subset of tokens with the highest cumulative probability.
In NLP, probability distributions are often used to model the likelihood of different tokens appearing in a sequence. These distributions can be obtained through various means, such as training a language model on a large corpus of text. The cumulative probability of a set of tokens is the sum of their individual probabilities, and it represents the likelihood of any token in the set being selected. Renormalization is the process of adjusting the probabilities of a subset of tokens so that they sum to 1, which is essential for ensuring that the resulting distribution is valid.
The concept of top-p sampling is closely related to information theory and decision theory, as it involves making choices based on uncertain outcomes. The threshold p determines the minimum cumulative probability required for a set of tokens to be considered, and it controls the trade-off between exploring less likely tokens and exploiting the most likely ones. Understanding these concepts is essential for implementing top-p sampling and other text generation techniques.
Algorithm/Approach
The general approach to solving this problem involves implementing the top-p sampling algorithm, which consists of sorting, cumulative probability computation, threshold-based selection, and renormalization. This algorithm can be viewed as a filtering process, where the goal is to select a subset of tokens that meet a certain criterion (cumulative probability exceeding the threshold p). The renormalization step ensures that the resulting distribution is valid and can be used for further processing or decision-making.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.