AI核采样：Top-k不够用了之后

import torch
import torch.nn.functional as F

def apply_temperature(logits, temperature):
    """
    应用温度参数
    logits: [vocab_size] 或 [batch_size, vocab_size]
    temperature: float
    """
    # 防止除以0
    if temperature == 0:
        temperature = 1e-5
    return logits / temperature

def nucleus_sampling(logits, top_p=0.9, temperature=1.0):
    """
    实现 Top-p (Nucleus) 采样
    logits: 模型输出的原始 logits [vocab_size]
    top_p: 核采样的阈值，通常 0.9 或 0.95
    temperature: 温度参数
    """
    # 1. 应用 Temperature
    scaled_logits = apply_temperature(logits, temperature)

    # 2. 计算 Softmax 得到概率分布
    probs = F.softmax(scaled_logits, dim=-1)

    # 3. 按概率从大到小排序
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)

    # 4. 计算累积概率
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    # 5. 移除累积概率超过 top_p 的 Token
    # 这里的减法是为了确保在临界点时保留该 Token (例如 0.9 - 0.9 = 0)
    sorted_indices_to_remove = cumulative_probs > top_p
    # 第一个超过 top_p 的 Token 也要保留
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0

    # 6. 将被移除的 Token 概率置为 0
    sorted_probs[sorted_indices_to_remove] = 0.0

    # 7. 重新归一化
    sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True)

    # 8. 根据排序后的索引，把概率放回原来的位置（这一步为了后续从原始采样）
    # 但实际上采样可以直接在 sorted_probs 上进行，然后映射回原始索引
    # 这里为了简化展示，直接采样
    next_token_idx = torch.multinomial(sorted_probs, num_samples=1)

    # 9. 将排序后的索引映射回原始词汇表索引
    original_next_token_idx = sorted_indices.gather(dim=-1, index=next_token_idx)

    return original_next_token_idx.item()

sorted_indices_to_remove = cumulative_probs > top_p

sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()

import matplotlib.pyplot as plt
import numpy as np

# 模拟一个词表大小为 1000 的 Logits 分布
# 假设前几个词概率很高，后面是长尾
vocab_size = 1000
logits = torch.randn(vocab_size)
# 手动放大前几个词的 Logits，模拟真实的高概率预测
logits[0] += 5
logits[1] += 3
logits[2] += 2
logits[3] += 1

# 设置参数
temperature = 0.8
top_p = 0.9

# 1. 原始 Softmax
probs_orig = F.softmax(logits, dim=-1).numpy()

# 2. Top-p 采样后的分布
# 这里复用上面的 nucleus_sampling 逻辑来获取处理后的概率矩阵用于绘图
def get_filtered_probs(logits, top_p, temperature):
    scaled_logits = logits / temperature
    probs = F.softmax(scaled_logits, dim=-1)
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
    sorted_indices_to_remove = cumulative_probs > top_p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0
    sorted_probs[sorted_indices_to_remove] = 0.0
    sorted_probs /= sorted_probs.sum()
    # 恢复原始顺序
    original_probs = torch.zeros_like(probs)
    original_probs.scatter_(0, sorted_indices, sorted_probs)
    return original_probs.numpy()

probs_filtered = get_filtered_probs(logits, top_p, temperature)

# 绘图
plt.figure(figsize=(12, 6))
indices = np.arange(vocab_size)

plt.bar(indices, probs_orig, alpha=0.5, label='Original Probabilities', width=1.0)
plt.bar(indices, probs_filtered, alpha=0.8, label=f'Nucleus Sampled (Top-p={top_p})', width=1.0)

plt.yscale('log') # 使用对数坐标以便观察长尾
plt.xlabel('Token Index')
plt.ylabel('Probability (Log Scale)')
plt.title('Impact of Nucleus Sampling on Probability Distribution')
plt.legend()
plt.grid(True, which="both", ls="-", alpha=0.2)
plt.show()