AI_Top-k采样踩坑记录

import torch
import torch.nn.functional as F

def top_k_sampling(logits, temperature=1.0, top_k=50):
    """
    基础的 Top-k 采样实现
    
    Args:
        logits: 模型输出的 logit，形状 [batch_size, vocab_size]
        temperature: 温度参数，控制分布的平滑度
        top_k: 保留的候选词数量
    
    Returns:
        sampled_tokens: 采样得到的 token 索引
    """
    # 应用 temperature
    logits = logits / temperature
    
    # 找到 top-k 的值和索引
    top_k_logits, top_k_indices = torch.topk(logits, top_k)
    
    # 创建一个 mask，非 top-k 的位置设为负无穷
    indices_to_remove = logits < top_k_logits[:, -1:]
    logits[indices_to_remove] = float('-inf')
    
    # 计算概率分布并采样
    probs = F.softmax(logits, dim=-1)
    sampled_tokens = torch.multinomial(probs, num_samples=1)
    
    return sampled_tokens

def top_k_sampling_improved(logits, temperature=1.0, top_k=50, filter_value=-float('inf')):
    """
    改进的 Top-k 采样实现
    
    Args:
        logits: 模型输出的 logit，形状 [batch_size, vocab_size]
        temperature: 温度参数
        top_k: 保留的候选词数量
        filter_value: 用于过滤的值，默认负无穷
    
    Returns:
        sampled_tokens: 采样得到的 token 索引
    """
    if top_k <= 0:
        # 如果 top_k 为 0，相当于不进行过滤
        return logits
    
    # 应用 temperature
    logits = logits / temperature
    
    # 找到 top-k 的值
    top_k_logits, _ = torch.topk(logits, top_k)
    
    # 获取 top-k 的最小值作为阈值
    threshold = top_k_logits[:, -1]
    
    # 创建 mask，低于阈值的设为 filter_value
    indices_to_remove = logits < threshold.unsqueeze(1)
    logits[indices_to_remove] = filter_value
    
    # 计算概率分布并采样
    probs = F.softmax(logits, dim=-1)
    sampled_tokens = torch.multinomial(probs, num_samples=1)
    
    return sampled_tokens

vocab_size = logits.size(-1)
top_k = min(top_k, vocab_size)  # 确保不超过词表大小

temperature = max(temperature, 1e-6)  # 避免 temperature 过小

# logits 形状: [batch_size, vocab_size]
# threshold 形状: [batch_size]
# threshold.unsqueeze(1) 形状: [batch_size, 1]
# 这样 broadcasting 才能正确工作

def apply_repetition_penalty(logits, token_ids, penalty=1.2):
    """
    应用重复惩罚，降低已生成词的 logit 值
    
    Args:
        logits: 模型输出的 logit
        token_ids: 已生成的 token 序列
        penalty: 惩罚系数，大于 1 表示惩罚
    
    Returns:
        logits: 应用惩罚后的 logit
    """
    for token_id in token_ids:
        logits[:, token_id] = logits[:, token_id] / penalty
    return logits

import matplotlib.pyplot as plt
import numpy as np

# 数据
strategies = ['Top-1\n贪婪', 'Top-k\n(k=10)', 'Top-k\n(k=50)', 'Top-k\n(k=100)', 'Temp=1.0']
diversity = [0.12, 0.34, 0.56, 0.68, 0.72]
coherence = [0.92, 0.89, 0.85, 0.78, 0.71]

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

# 左图：多样性和连贯性的对比
x = np.arange(len(strategies))
width = 0.35

bars1 = ax1.bar(x - width/2, diversity, width, label='多样性', color='#7cb342')
bars2 = ax1.bar(x + width/2, coherence, width, label='连贯性', color='#42a5f5')

ax1.set_xlabel('采样策略')
ax1.set_ylabel('评分 (0-1)')
ax1.set_title('多样性与连贯性对比')
ax1.set_xticks(x)
ax1.set_xticklabels(strategies)
ax1.legend()
ax1.set_ylim(0, 1)

# 添加数值标签
for bars in [bars1, bars2]:
    for bar in bars:
        height = bar.get_height()
        ax1.text(bar.get_x() + bar.get_width()/2., height,
                f'{height:.2f}',
                ha='center', va='bottom', fontsize=9)

# 右图：多样性 vs 连贯性 散点图（权衡曲线）
ax2.scatter(diversity, coherence, s=200, c=['#e53935', '#fb8c00', '#43a047', '#7e57c2', '#ec407a'])
ax2.plot(diversity, coherence, 'k--', alpha=0.3)

# 标注每个点
for i, strategy in enumerate(strategies):
    ax2.annotate(strategy.replace('\n', ' '),
                (diversity[i], coherence[i]),
                xytext=(5, 5), textcoords='offset points',
                fontsize=8, ha='left')

ax2.set_xlabel('多样性评分')
ax2.set_ylabel('连贯性评分')
ax2.set_title('多样性 vs 连贯性权衡曲线')
ax2.set_xlim(0, 0.8)
ax2.set_ylim(0.6, 1.0)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('/home/liqinsi/Documents/project/thinkblog/static/images/posts/350-topk-sampling-comparison.webp', 
            dpi=150, bbox_inches='tight')
plt.close()

def advanced_sampling(model, input_ids, max_length=100, 
                     temperature=0.8, top_k=50, repetition_penalty=1.2):
    """
    完整的采样策略实现
    
    Args:
        model: 语言模型
        input_ids: 输入 token 序列
        max_length: 最大生成长度
        temperature: 温度参数
        top_k: top-k 候选数量
        repetition_penalty: 重复惩罚系数
    
    Returns:
        generated_ids: 生成的完整 token 序列
    """
    generated_ids = input_ids.clone()
    
    with torch.no_grad():
        for _ in range(max_length):
            # 前向传播获取 logits
            outputs = model(generated_ids)
            logits = outputs.logits[:, -1, :]  # 取最后一个位置的 logits
            
            # 应用重复惩罚
            if repetition_penalty > 1.0:
                logits = apply_repetition_penalty(logits, generated_ids, repetition_penalty)
            
            # Top-k 采样
            next_token = top_k_sampling_improved(logits, temperature, top_k)
            
            # 拼接新 token
            generated_ids = torch.cat([generated_ids, next_token], dim=-1)
            
            # 如果生成了结束符，停止生成
            if next_token.item() == model.config.eos_token_id:
                break
    
    return generated_ids