Top-p-nucleus sampling

# High-confidence situation: model strongly prefers 'the'
high_confidence_logits = np.array([5.0, 1.0, 0.5, 0.1, -0.5])
hc_probs = softmax(high_confidence_logits)
print('High confidence probs:', np.round(hc_probs, 3))
# [0.974, 0.018, 0.011, 0.007, 0.004]
# Top-p=0.9 nucleus: just 1 token (cumulative after token 0 = 97.4% > 90%)

# Low-confidence situation: model unsure
low_confidence_logits = np.array([1.1, 1.0, 0.9, 0.8, 0.7])
lc_probs = softmax(low_confidence_logits)
print('Low confidence probs:', np.round(lc_probs, 3))
# [0.218, 0.208, 0.199, 0.190, 0.181]
# Top-p=0.9 nucleus: all 5 tokens (need all to reach 90%)

import openai
client = openai.OpenAI(api_key='sk-...')

# Nucleus sampling: top 90% of probability mass
resp = client.chat.completions.create(
    model='gpt-4o',
    messages=[{'role': 'user', 'content': 'Tell me an interesting fact about the ocean.'}],
    temperature=1.0,   # leave temperature at default
    top_p=0.9          # sample from top 90% nucleus
)
print(resp.choices[0].message.content)

# p=1.0: all tokens in nucleus
resp_full = client.chat.completions.create(
    model='gpt-4o',
    messages=[{'role': 'user', 'content': 'Generate a creative story opening.'}],
    temperature=1.0,
    top_p=1.0  # no nucleus restriction
)

# p=0.5: very focused nucleus
resp_focused = client.chat.completions.create(
    model='gpt-4o',
    messages=[{'role': 'user', 'content': 'Generate a creative story opening.'}],
    temperature=1.0,
    top_p=0.5  # only top 50% probability mass
)

# The tail problem with temperature alone
high_temp_logits = np.array([3.0, 2.0, 1.0, 0.0, -1.0, -5.0, -10.0])
probs_high_temp = softmax(high_temp_logits, temperature=2.0)
print('High temp probs:', np.round(probs_high_temp, 4))
# The last token (logit=-10) still has a small probability
# With many tokens in a real vocab, these rare tokens accumulate
# and occasionally get sampled, producing incoherent output

# Top-p=0.9 cuts these off entirely
# ensuring only tokens contributing to the top 90% are considered

def combined_sample(logits, temperature=1.0, top_p=0.9):
    # Step 1: apply temperature
    probs = softmax(logits, temperature=temperature)

    # Step 2: apply top-p nucleus
    sorted_idx = np.argsort(probs)[::-1]
    sorted_probs = probs[sorted_idx]
    cumulative = np.cumsum(sorted_probs)
    nucleus_size = np.searchsorted(cumulative, top_p) + 1
    nucleus_idx = sorted_idx[:nucleus_size]
    nucleus_probs = probs[nucleus_idx]
    nucleus_probs = nucleus_probs / nucleus_probs.sum()

    return np.random.choice(nucleus_idx, p=nucleus_probs)

def detect_repetition(text, window=20):
    words = text.split()
    if len(words) < window * 2:
        return False
    # Check if any window of words repeats within the text
    for i in range(len(words) - window):
        phrase = ' '.join(words[i:i + window])
        rest = ' '.join(words[i + window:])
        if phrase in rest:
            return True
    return False

response = call_llm(prompt)
if detect_repetition(response):
    print('Warning: repetition detected — consider increasing top_p or temperature')

# Top-k equivalent for comparison
def top_k_sample(logits, k=50):
    probs = softmax(logits)
    # Keep only top-k tokens
    top_k_indices = np.argsort(probs)[::-1][:k]
    top_k_probs = probs[top_k_indices]
    top_k_probs = top_k_probs / top_k_probs.sum()
    return np.random.choice(top_k_indices, p=top_k_probs)

# Key difference: k is always 50, regardless of model confidence
# Top-p nucleus size varies from 1 to thousands depending on confidence

# Guidance: what to adjust based on symptoms
TROUBLESHOOTING = {
    'output is incoherent or contains random words': {
        'fix': 'lower top_p to 0.9 or 0.85',
        'or': 'lower temperature'
    },
    'output is repetitive and looping': {
        'fix': 'increase top_p or temperature',
        'also': 'try adding frequency_penalty or presence_penalty'
    },
    'output is too predictable and boring': {
        'fix': 'increase temperature to 0.9-1.2',
        'keep': 'top_p at 0.95'
    },
    'output needs to be deterministic': {
        'fix': 'set temperature=0, top_p=1.0'
    }
}

import anthropic
claude = anthropic.Anthropic(api_key='sk-ant-...')

# Creative writing with nucleus sampling
message = claude.messages.create(
    model='claude-opus-4-5',
    max_tokens=256,
    temperature=1.0,
    top_p=0.95,
    messages=[{
        'role': 'user',
        'content': 'Write a short poem about the ocean.'
    }]
)
print(message.content[0].text)