文本生成：语言模型与解码策略

# N-gram 语言模型：计数、平滑与困惑度
from collections import Counter, defaultdict
import math

class NgramLM:
    """N-gram 语言模型（带加 K 平滑）"""
    
    def __init__(self, n=3, k=0.1):
        self.n = n
        self.k = k  # 拉普拉斯平滑参数
        self.ngram_counts = Counter()
        self.context_counts = Counter()
        self.vocab_size = 0
    
    def fit(self, sentences):
        """训练 N-gram 模型"""
        for sent in sentences:
            tokens = ["<s>"] * (self.n - 1) + sent + ["</s>"]
            for i in range(len(tokens) - self.n + 1):
                ngram = tuple(tokens[i:i + self.n])
                context = ngram[:-1]
                self.ngram_counts[ngram] += 1
                self.context_counts[context] += 1
        self.vocab_size = len(set(t for s in sentences for t in s)) + 2  # +2 for <s> </s>
    
    def prob(self, word, context):
        """P(word | context) 带加 K 平滑"""
        count = self.ngram_counts[context + (word,)]
        total = self.context_counts[context]
        return (count + self.k) / (total + self.k * self.vocab_size)
    
    def perplexity(self, sentences):
        """计算模型在测试集上的困惑度"""
        log_prob_sum = 0
        total_tokens = 0
        for sent in sentences:
            tokens = ["<s>"] * (self.n - 1) + sent + ["</s>"]
            for i in range(self.n - 1, len(tokens)):
                context = tuple(tokens[i - self.n + 1:i])
                p = self.prob(tokens[i], context)
                log_prob_sum += math.log2(max(p, 1e-10))
                total_tokens += 1
        return 2 ** (-log_prob_sum / total_tokens)

# 训练与评估
corpus = [
    ["i", "love", "machine", "learning"],
    ["machine", "learning", "is", "amazing"],
    ["i", "love", "natural", "language", "processing"],
]
model = NgramLM(n=3, k=0.5)
model.fit(corpus)
print(f"Trigram 困惑度: {model.perplexity(corpus):.2f}")

# Transformer 因果语言模型（简化版）
import torch
import torch.nn as nn
import math

class CausalTransformerLM(nn.Module):
    """仅解码器的因果语言模型"""
    
    def __init__(self, vocab_size=50257, d_model=768, n_layers=12,
                 n_heads=12, d_ff=3072, max_len=1024, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(max_len, d_model)
        self.dropout = nn.Dropout(dropout)
        
        # Transformer 解码器层
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=n_heads,
            dim_feedforward=d_ff, dropout=dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=n_layers)
        self.lm_head = nn.Linear(d_model, vocab_size)
    
    def forward(self, input_ids):
        batch_size, seq_len = input_ids.shape
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
        
        # 词嵌入 + 位置嵌入
        x = self.token_embed(input_ids) + self.pos_embed(positions)
        x = self.dropout(x)
        
        # 因果掩码：禁止关注未来位置
        causal_mask = torch.triu(
            torch.ones(seq_len, seq_len, device=input_ids.device), diagonal=1
        ).bool()
        causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq, seq)
        
        # Transformer 解码
        memory = torch.zeros_like(x)  # 编码器-解码器架构需要，这里用零填充
        output = self.transformer(x, memory, tgt_mask=causal_mask)
        
        # 输出词表概率分布
        logits = self.lm_head(output)  # (batch, seq, vocab)
        return logits
    
    def generate_probs(self, input_ids):
        """获取下一个词的概率分布"""
        logits = self.forward(input_ids)
        next_token_logits = logits[:, -1, :]  # 取最后一个位置
        probs = torch.softmax(next_token_logits, dim=-1)
        return probs  # (batch, vocab_size)

# 贪心搜索实现
import torch
import torch.nn.functional as F

def greedy_decode(model, prompt_ids, max_length=50, eos_token_id=2):
    """
    贪心搜索解码
    model: 语言模型，接受 input_ids 返回 logits
    prompt_ids: (1, prompt_len) 提示词 token IDs
    返回: 生成的完整 token IDs
    """
    input_ids = prompt_ids.clone()
    
    for _ in range(max_length):
        # 获取模型输出
        with torch.no_grad():
            logits = model(input_ids)  # (1, seq_len, vocab_size)
        
        # 取最后一个位置的 logits
        next_token_logits = logits[:, -1, :]  # (1, vocab_size)
        
        # 贪心选择概率最高的词
        next_token_id = torch.argmax(next_token_logits, dim=-1)  # (1,)
        
        # 追加到序列
        input_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=1)
        
        # 遇到 EOS 停止
        if next_token_id.item() == eos_token_id:
            break
    
    return input_ids

# 贪心搜索的问题演示
print("=== 贪心搜索的问题 ===")
print("贪心搜索在每一步做局部最优选择:")
print("  Step 1: P('the') = 0.40 → 选择 'the' ✓")
print("  Step 2: P('cat') = 0.35 → 选择 'cat' ✓")
print("  Step 3: P('is')  = 0.60 → 选择 'is'  ✓")
print("  但全局最优可能是:")
print("  Step 1: P('a')   = 0.30 (非最优)")
print("  Step 2: P('beautiful') = 0.25 (非最优)")
print("  Step 3: P('sunset') = 0.45 → 整体概率可能更高!")
print("贪心搜索会错过这个全局更优的路径。")

# Beam Search 完整实现
import torch
import torch.nn.functional as F
import heapq

def beam_search_decode(model, prompt_ids, beam_size=4, max_length=50,
                       eos_token_id=2, length_penalty=0.6):
    """
    Beam Search 解码
    model: 语言模型
    prompt_ids: (1, prompt_len)
    beam_size: 束宽度
    length_penalty: 长度归一化指数
    """
    # 初始化 beam: [(score, sequence, is_done)]
    beams = [(0.0, prompt_ids.clone(), False)]
    completed = []
    
    for step in range(max_length):
        candidates = []
        
        for score, seq, is_done in beams:
            if is_done:
                candidates.append((score, seq, True))
                continue
            
            with torch.no_grad():
                logits = model(seq)
                next_logits = logits[:, -1, :]  # (1, vocab)
                log_probs = F.log_softmax(next_logits, dim=-1).squeeze(0)
            
            # 取 top-k 候选
            top_k_probs, top_k_ids = torch.topk(log_probs, beam_size)
            
            for prob, token_id in zip(top_k_probs, top_k_ids):
                new_score = score + prob.item()
                new_seq = torch.cat([seq, token_id.unsqueeze(0).unsqueeze(0)], dim=1)
                is_finished = token_id.item() == eos_token_id
                
                if is_finished:
                    # 长度归一化
                    norm_score = new_score / (new_seq.shape[1] ** length_penalty)
                    completed.append((norm_score, new_seq))
                else:
                    candidates.append((new_score, new_seq, False))
        
        # 保留 top-k 候选
        beams = heapq.nlargest(beam_size, candidates, key=lambda x: x[0])
        
        if not beams:
            break
    
    # 合并已完成和未完成的序列
    all_seqs = completed + [(s / (seq.shape[1] ** length_penalty), seq)
                            for s, seq, _ in beams]
    all_seqs.sort(key=lambda x: x[0], reverse=True)
    
    return all_seqs  # [(score, seq), ...] 按分数排序

# 对比演示
print("=== Beam Search vs 贪心搜索 ===")
print("Beam Size = 4 时的搜索空间:")
print("  Step 1: 评估整个词表 → 保留 top-4")
print("  Step 2: 每个候选扩展 top-4 → 4×4=16 个候选 → 保留 top-4")
print("  Step 3: 同上 → 16 个候选 → 保留 top-4")
print("  ...")
print("  相比贪心搜索，Beam Search 探索了 4^N 倍的路径空间")

# 温度采样详解
import torch
import torch.nn.functional as F

def temperature_sampling(logits, temperature=1.0):
    """温度采样：控制生成的随机性程度"""
    scaled_logits = logits / temperature
    probs = F.softmax(scaled_logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)

# 温度对概率分布的影响
logits = torch.tensor([5.0, 3.0, 1.0, 0.5, -1.0, -2.0])
words = ["the", "a", "one", "this", "that", "some"]

print("=== 温度对概率分布的影响 ===")
print(f"{'词':<8} {'Logits':>8}", end="")
for t in [0.1, 0.5, 1.0, 2.0, 5.0]:
    print(f" {f'T={t:.1f}':>10}", end="")
print()

for word, logit in zip(words, logits):
    print(f"{word:<8} {logit:>8.1f}", end="")
    for t in [0.1, 0.5, 1.0, 2.0, 5.0]:
        p = F.softmax(logit.unsqueeze(0) / t, dim=-1).item()
        print(f" {p:>10.4f}", end="")
    print()

print("\n观察:")
print("  T=0.1: 几乎 100% 选择 'the' (贪心搜索)")
print("  T=0.5: 'the' 占主导，但 'a' 也有机会")
print("  T=1.0: 原始分布，概率自然衰减")
print("  T=2.0: 分布变平，低概率词机会增加")
print("  T=5.0: 接近均匀分布，几乎随机")

# Top-k 与 Top-p 采样对比
import torch
import torch.nn.functional as F

def top_k_sampling(logits, k=50):
    """Top-k 采样：只从概率最高的 k 个词中采样"""
    top_k_values, top_k_indices = torch.topk(logits, k)
    # 将非 top-k 的 logits 设为负无穷
    filtered_logits = torch.full_like(logits, float('-inf'))
    filtered_logits.scatter_(0, top_k_indices, top_k_values)
    probs = F.softmax(filtered_logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)

def top_p_sampling(logits, p=0.9):
    """Top-p (Nucleus) 采样：按累积概率截断"""
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
    
    # 找到累积概率首次 >= p 的位置
    remove_mask = cumulative_probs > p
    # 保留第一个超过 p 的词（确保概率和 >= p）
    remove_mask[1:] = remove_mask[:-1].clone()
    remove_mask[0] = False
    
    # 移除低概率词
    filtered_logits = sorted_logits.masked_fill(remove_mask, float('-inf'))
    # 恢复原始顺序
    original_order = torch.argsort(sorted_indices)
    restored_logits = filtered_logits[original_order]
    
    probs = F.softmax(restored_logits, dim=-1)
    return torch.multinomial(probs, num_samples=1)

# 对比两种策略
logits = torch.randn(1000)  # 模拟 1000 个候选词的 logits

# Top-k: 固定保留 50 个词
top_k_candidates = torch.topk(logits, 50)
print(f"Top-k=50: 始终保留 {len(top_k_candidates.values)} 个候选词")

# Top-p: 自适应保留
sorted_probs = F.softmax(torch.sort(logits, descending=True).values, dim=-1)
cum_probs = torch.cumsum(sorted_probs, dim=-1)
top_p_count = (cum_probs <= 0.9).sum().item() + 1
print(f"Top-p=0.9: 保留 {top_p_count} 个候选词 (自适应)")

print("\nTop-k vs Top-p 的关键区别:")
print("  Top-k: 固定数量，不考虑概率分布形状")
print("    → 确定上下文 vs 模糊上下文都保留 50 个词")
print("  Top-p: 按累积概率，自适应调整候选数量")
print("    → 确定上下文可能只保留 3 个词")
print("    → 模糊上下文可能保留 100 个词")

# 核采样的可视化分析
import torch
import torch.nn.functional as F

def nucleus_sampling_analysis(logits, p_values=[0.5, 0.7, 0.9, 0.95, 0.99]):
    """分析不同 top-p 值对候选集大小的影响"""
    probs = F.softmax(logits, dim=-1)
    sorted_probs, _ = torch.sort(probs, descending=True)
    cum_probs = torch.cumsum(sorted_probs, dim=-1)
    
    print("=== 核采样候选集分析 ===")
    print(f"{'top-p':>6} {'候选词数':>8} {'累积概率':>8} {'最高词概率':>10}")
    print("-" * 40)
    
    for p in p_values:
        # 找到最小 k 使得累积概率 >= p
        k = (cum_probs < p).sum().item() + 1
        cum_at_k = cum_probs[k - 1].item()
        top_prob = sorted_probs[0].item()
        print(f"{p:>6.2f} {k:>8} {cum_at_k:>8.4f} {top_prob:>10.4f}")
    
    return k

# 场景 1: 高确定性上下文
print("场景 1: 「The capital of France is ___」")
logits_certain = torch.tensor([8.0, 2.0, 1.0, 0.5, 0.0, -1.0, -1.5, -2.0])
nucleus_sampling_analysis(logits_certain)

# 场景 2: 低确定性上下文
print("\n场景 2: 「The movie was very ___」")
logits_uncertain = torch.tensor([3.0, 2.8, 2.7, 2.5, 2.4, 2.3, 2.2, 2.1])
nucleus_sampling_analysis(logits_uncertain)

print("\n关键发现:")
print("  高确定性: top-p=0.9 可能只保留 1-2 个词")
print("  低确定性: top-p=0.9 可能保留全部 8 个词")
print("  → 核采样根据上下文不确定性自动调节候选集大小")

# 核采样与 Beam Search 生成质量对比
import torch
import torch.nn.functional as F

def compare_decoding_strategies():
    """对比不同解码策略的生成结果特征"""
    strategies = {
        "贪心搜索": {
            "优点": ["确定性", "计算高效", "适合结构化输出"],
            "缺点": ["容易重复", "缺乏多样性", "局部最优陷阱"],
            "典型输出": "这是一个非常好的产品。这是一个非常好的产品。",
        },
        "Beam Search (k=4)": {
            "优点": ["质量稳定", "翻译效果好", "长度可控"],
            "缺点": ["保守乏味", "缺乏惊喜", "可能重复"],
            "典型输出": "该产品具有良好的性能和出色的用户体验。",
        },
        "Top-k (k=50)": {
            "优点": ["避免尾部噪声", "适度多样性", "实现简单"],
            "缺点": ["候选集大小固定", "不考虑概率分布形状"],
            "典型输出": "这个产品真的很棒，使用体验超出预期。",
        },
        "Top-p (p=0.9)": {
            "优点": ["自适应候选集", "自然多样", "论文推荐"],
            "缺点": ["计算略复杂", "需要调 p 值"],
            "典型输出": "用了这款产品之后，感觉生活品质都提升了。",
        },
        "Temperature (T=0.7)": {
            "优点": ["简单直观", "连续可调", "与 top-p 兼容"],
            "缺点": ["全局影响分布", "不区分上下文确定性"],
            "典型输出": "这款产品的表现确实让人印象深刻。",
        },
    }
    
    for name, info in strategies.items():
        print(f"\n{'='*40}")
        print(f"策略: {name}")
        print(f"{'='*40}")
        print(f"  优点: {', '.join(info['优点'])}")
        print(f"  缺点: {', '.join(info['缺点'])}")
        print(f"  示例: {info['典型输出']}")

compare_decoding_strategies()

# 推荐组合
print("\n\n=== 推荐解码组合 ===")
print("代码生成: temperature=0.2 + top_p=1.0 (确定性优先)")
print("技术写作: temperature=0.7 + top_p=0.9 (平衡质量与多样性)")
print("创意写作: temperature=1.0 + top_p=0.95 (多样性优先)")
print("对话系统: temperature=0.8 + top_p=0.9 (自然流畅)")
print("翻译任务: beam_size=4-6 (确定性最佳) 或 top_p=0.95 (多样性) ")

# 重复惩罚（Repetition Penalty）实现
import torch
import torch.nn.functional as F

def apply_repetition_penalty(logits, generated_ids, penalty=1.2):
    """
    应用重复惩罚到 logits
    logits: (1, vocab_size) 下一步的 logits
    generated_ids: (seq_len,) 已经生成的词 IDs
    penalty: 惩罚因子 γ (γ > 1 表示惩罚)
    """
    for token_id in generated_ids:
        if logits[0, token_id] > 0:
            # 正 logits → 除以惩罚因子 (减小)
            logits[0, token_id] /= penalty
        else:
            # 负 logits → 乘以惩罚因子 (更负)
            logits[0, token_id] *= penalty
    
    return logits

# 演示重复惩罚的效果
vocab_words = ["the", "cat", "sat", "on", "mat", "and", "slept", "peacefully"]
logits = torch.tensor([5.0, 4.0, 3.5, 3.0, 2.5, 1.0, 0.5, 0.0])
generated = torch.tensor([0, 1, 2, 0, 3, 4])  # 已生成: "the cat sat the on mat"

print("=== 重复惩罚效果 ===")
print("原始概率分布:")
probs_before = F.softmax(logits, dim=-1)
for word, prob in zip(vocab_words, probs_before):
    print(f"  {word:<12} {prob:.4f}")

# 应用重复惩罚
penalized_logits = apply_repetition_penalty(logits.clone(), generated, penalty=1.5)
probs_after = F.softmax(penalized_logits, dim=-1)

print("\n应用重复惩罚后 (γ=1.5):")
for word, prob in zip(vocab_words, probs_after):
    marker = " ⚠️ 已出现" if int(vocab_words.index(word)) in generated.tolist() else ""
    print(f"  {word:<12} {prob:.4f}{marker}")

print("\n观察: 'the'(0)、'cat'(1)、'sat'(2)、'on'(3)、'mat'(4) 的概率都被降低")
print("模型更可能选择 'and'、'slept' 等新词，打破重复循环")

# N-gram 阻塞与生成多样性增强
import torch
import torch.nn.functional as F
from collections import Counter

def get_ngrams(token_ids, n):
    """提取序列中的所有 n-gram"""
    return [tuple(token_ids[i:i+n]) for i in range(len(token_ids) - n + 1)]

def apply_ngram_blocking(logits, generated_ids, no_repeat_ngram_size=3):
    """
    N-gram 阻塞: 禁止生成会导致重复 n-gram 的词
    """
    if len(generated_ids) < no_repeat_ngram_size - 1:
        return logits  # 生成的词太少，无法构成 n-gram
    
    # 提取最近的 (n-1)-gram
    prefix = tuple(generated_ids[-(no_repeat_ngram_size - 1):])
    
    # 提取所有已出现的 n-gram
    all_ngrams = get_ngrams(generated_ids.tolist(), no_repeat_ngram_size)
    
    # 找到不允许接在 prefix 后面的词
    forbidden = set()
    for ngram in all_ngrams:
        if ngram[:-1] == prefix:
            forbidden.add(ngram[-1])
    
    # 将禁止词的 logits 设为负无穷
    for token_id in forbidden:
        logits[0, token_id] = float('-inf')
    
    return logits

# 多样性增强策略汇总
print("=== 生成多样性增强策略 ===\n")

strategies = {
    "重复惩罚 (Repetition Penalty)": {
        "参数": "penalty=1.0-2.0",
        "原理": "对已出现词的 logits 除以惩罚因子",
        "效果": "词级去重，温和",
        "推荐": "通用场景，γ=1.2"
    },
    "N-gram 阻塞": {
        "参数": "no_repeat_ngram_size=2-4",
        "原理": "禁止生成已出现的 n-gram 序列",
        "效果": "n-gram 级去重，强硬",
        "推荐": "需要严格避免重复时"
    },
    "多样性惩罚 (Diverse Beam Search)": {
        "参数": "num_beam_groups, diversity_penalty",
        "原理": "Beam Search 分多组，组间施加多样性惩罚",
        "效果": "一组内相似，组间差异大",
        "推荐": "需要多个不同候选时"
    },
    "对比搜索 (Contrastive Search)": {
        "参数": "alpha (惩罚权重), k (候选数)",
        "原理": "在概率高且与历史差异大的词中选择",
        "效果": "平衡流畅度和多样性",
        "推荐": "2022 年论文推荐，优于 top-p"
    },
}

for name, info in strategies.items():
    print(f"策略: {name}")
    for key, val in info.items():
        print(f"  {key}: {val}")
    print()

# 困惑度计算与解读
import torch
import torch.nn.functional as F
import math

def compute_perplexity(model, data_loader, device='cpu'):
    """计算语言模型在数据集上的困惑度"""
    model.eval()
    total_loss = 0
    total_tokens = 0
    
    with torch.no_grad():
        for batch in data_loader:
            input_ids = batch['input_ids'].to(device)
            labels = batch['labels'].to(device)
            
            outputs = model(input_ids, labels=labels)
            loss = outputs.loss  # 交叉熵损失
            
            total_loss += loss.item() * labels.numel()
            total_tokens += labels.numel()
    
    avg_loss = total_loss / total_tokens
    perplexity = math.exp(avg_loss)
    return perplexity

# 困惑度的直观理解
print("=== 困惑度的直观理解 ===\n")
ppl_values = {
    10: "极好 - 模型几乎确定了每个词（接近完美）",
    20: "很好 - 每步约 20 个候选（高质量语言模型）",
    50: "不错 - 每步约 50 个候选（可用水平）",
    100: "一般 - 每步约 100 个候选（基线水平）",
    200: "较差 - 模型困惑度较高（需要改进）",
    500: "很差 - 模型几乎在随机猜测",
    1000: "失败 - 模型未学到有效语言模式",
}

for ppl, desc in sorted(ppl_values.items()):
    marker = "★" if ppl <= 50 else "  "
    print(f"  PPL = {ppl:<5} {marker} {desc}")

print("\n注意: 困惑度的绝对值取决于词表大小!")
print("  词表=10K: PPL=50 可能已经很好")
print("  词表=50K: PPL=50 可能还不够好")
print("  比较困惑度时，必须在相同词表和数据集上比较")

# LLM-as-a-Judge: 用大模型评估生成质量
import json

def llm_judge_evaluation():
    """模拟 LLM-as-a-Judge 评估流程"""
    evaluation_prompt = """
请评估以下生成文本的质量，从 1-10 分打分：

【参考文本】
{reference}

【生成文本】
{generated}

请从以下维度评分：
1. 流畅性 (Fluency): 文本是否自然流畅，语法是否正确？
2. 一致性 (Coherence): 逻辑是否连贯，前后是否矛盾？
3. 信息完整性 (Informativeness): 是否包含了参考文本的关键信息？
4. 简洁性 (Conciseness): 是否简洁明了，没有冗余？

请以 JSON 格式输出评分：
{{"fluency": X, "coherence": X, "informativeness": X, "conciseness": X, "overall": X}}
"""

    # 评估示例
    examples = [
        {
            "reference": "Transformer 架构完全基于注意力机制，摒弃了传统的循环结构。",
            "generated": "Transformer 模型的核心是注意力机制，它不再使用循环神经网络。",
            "expected_scores": {"fluency": 9, "coherence": 9, "informativeness": 8, "conciseness": 9, "overall": 8.75}
        },
        {
            "reference": "深度学习需要大量标注数据和计算资源。",
            "generated": "深度学习需要数据和计算。深度学习需要数据和计算。",
            "expected_scores": {"fluency": 5, "coherence": 4, "informativeness": 6, "conciseness": 3, "overall": 4.5}
        },
    ]

    print("=== LLM-as-a-Judge 评估示例 ===\n")
    for i, ex in enumerate(examples):
        print(f"示例 {i + 1}:")
        print(f"  参考: {ex['reference']}")
        print(f"  生成: {ex['generated']}")
        print(f"  预期评分: {ex['expected_scores']}")
        print()

    print("LLM-as-a-Judge 的优势:")
    print("  1. 自动化，可扩展")
    print("  2. 与人类评估相关性高 (0.7-0.9)")
    print("  3. 可评估语义质量，不仅 n-gram 匹配")
    print("  4. 支持多维度细粒度评估")
    print("\n局限性:")
    print("  1. 评估模型自身可能存在偏见")
    print("  2. 不同评估模型之间结果可能不一致")
    print("  3. 计算成本较高（每次评估需要推理）")

llm_judge_evaluation()

# HuggingFace 文本生成：从入门到精通
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import torch

# === 方式 1: Pipeline (最简单) ===
print("=== 方式 1: Pipeline ===")
generator = pipeline("text-generation", model="gpt2")
result = generator(
    "人工智能的未来发展",
    max_new_tokens=50,
    do_sample=True,
    temperature=0.7,
    top_p=0.9,
    num_return_sequences=1,
)
print(f"生成结果: {result[0]['generated_text']}")

# === 方式 2: model.generate() (推荐) ===
print("\n=== 方式 2: model.generate() ===")
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

prompt = "人工智能的未来发展"
inputs = tokenizer(prompt, return_tensors="pt")

# 贪心搜索
greedy_output = model.generate(
    **inputs,
    max_new_tokens=50,
    do_sample=False,  # 贪心搜索
)
print(f"贪心搜索: {tokenizer.decode(greedy_output[0], skip_special_tokens=True)}")

# Beam Search
beam_output = model.generate(
    **inputs,
    max_new_tokens=50,
    num_beams=5,
    early_stopping=True,
)
print(f"Beam Search: {tokenizer.decode(beam_output[0], skip_special_tokens=True)}")

# 核采样 + 温度
sampled_output = model.generate(
    **inputs,
    max_new_tokens=50,
    do_sample=True,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.2,
)
print(f"核采样: {tokenizer.decode(sampled_output[0], skip_special_tokens=True)}")

# 高级生成控制：自定义 LogitsProcessor
from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    LogitsProcessor, LogitsProcessorList,
    GenerationConfig
)
import torch

class KeywordForcingProcessor(LogitsProcessor):
    """强制生成包含特定关键词的文本"""
    
    def __init__(self, keyword_ids, min_step=10):
        self.keyword_ids = set(keyword_ids)
        self.min_step = min_step  # 至少生成 min_step 步后才开始强制
    
    def __call__(self, input_ids, scores):
        step = input_ids.shape[1]
        if step < self.min_step:
            # 降低非关键词的概率
            mask = torch.ones_like(scores, dtype=torch.bool)
            mask[:, list(self.keyword_ids)] = False
            scores[mask] -= 10.0  # 大幅降低非关键词概率
        return scores

class LengthControlProcessor(LogitsProcessor):
    """控制生成长度：接近目标长度时提高 EOS 概率"""
    
    def __init__(self, target_length, eos_token_id, scale=5.0):
        self.target_length = target_length
        self.eos_token_id = eos_token_id
        self.scale = scale
    
    def __call__(self, input_ids, scores):
        current_length = input_ids.shape[1]
        if current_length >= self.target_length * 0.8:
            # 接近目标长度时，逐渐提高 EOS 概率
            progress = (current_length - self.target_length * 0.8) / (self.target_length * 0.2)
            scores[:, self.eos_token_id] += self.scale * progress
        return scores

# === 使用 GenerationConfig 管理参数 ===
print("=== GenerationConfig 最佳实践 ===\n")

config = GenerationConfig(
    max_new_tokens=100,
    do_sample=True,
    temperature=0.7,
    top_p=0.9,
    top_k=50,
    repetition_penalty=1.2,
    no_repeat_ngram_size=3,
    num_return_sequences=3,
    return_dict_in_generate=True,
    output_scores=True,
)

print(f"生成配置:")
for key, value in vars(config).items():
    if not key.startswith('_') and value is not None:
        print(f"  {key}: {value}")

# === 解码策略速查表 ===
print("\n=== 解码策略参数速查 ===")
strategies = {
    "贪心搜索": {"do_sample": False},
    "Beam Search": {"num_beams": 5, "early_stopping": True},
    "Top-k 采样": {"do_sample": True, "top_k": 50, "temperature": 0.7},
    "Top-p 采样": {"do_sample": True, "top_p": 0.9, "temperature": 0.7},
    "对比搜索": {"penalty_alpha": 0.6, "top_k": 4},
    "多样 Beam Search": {"num_beams": 5, "num_beam_groups": 5, "diversity_penalty": 1.0},
}

for name, params in strategies.items():
    params_str = ", ".join(f"{k}={v}" for k, v in params.items())
    print(f"  {name}: {params_str}")