Nucleus Sampling

hardPROText Generation & AdvancedNLP 3: Advanced

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

Input

0.8
the a an this that
0.35 0.25 0.20 0.15 0.05

Output

{'the': 0.4375, 'a': 0.3125, 'an': 0.25}

Reasoning

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

Learn this in

Probability & StatisticsFoundations Text Generation & SummarizationNLP Study Plan Sequence ModelsNLP Study Plan

More from NLP 3: Advanced

Temperature Samplereasy ROUGE-1 Score Calculatoreasy Context Window Findereasy Keyword Answer Extractoreasy Levenshtein Distanceeasy

Concept primerPREVIEW

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.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

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.

solution.py

Python 3.13

Test results