# 导入必要的库
import torch
import numpy as np
import matplotlib.pyplot as plt
from transformers import AutoModelForCausalLM, AutoTokenizer

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

print("正在加载中文 GPT-2 模型...")

# 使用中文 GPT-2 模型
model_name = "uer/gpt2-chinese-cluecorpussmall"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.eval()

# 设置 pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("✓ 模型加载完成！")

# 定义文本生成函数
def generate_text(prompt, max_length=50, num_return=1, temperature=1.0):
    """
    根据提示生成文本
    
    参数：
        prompt: 输入提示
        max_length: 最大生成长度
        num_return: 返回几个生成结果
        temperature: 温度参数（越低越确定性）
    """
    inputs = tokenizer.encode(prompt, return_tensors='pt')
    
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_length=max_length,
            num_return_sequences=num_return,
            temperature=temperature,
            do_sample=True,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id
        )
    
    results = []
    for output in outputs:
        text = tokenizer.decode(output, skip_special_tokens=True)
        results.append(text)
    
    return results

# 测试生成函数
test_result = generate_text("你好，我的名字是", max_length=30)
print(f"测试生成：{test_result[0]}")

# 【填空 1】实现困惑度计算函数
# 提示：困惑度 = exp(平均交叉熵损失)

def calculate_perplexity(text):
    """
    计算文本的困惑度
    困惑度越低，说明模型对这段文本越"熟悉"
    """
    # 编码文本
    inputs = tokenizer.encode(text, return_tensors='pt')
    
    # 计算损失
    with torch.no_grad():
        outputs = model(inputs, labels=inputs)
        loss = outputs.loss.item()  # 平均交叉熵损失
    
    # 【填空 1】计算困惑度
    # 提示：perplexity = np.exp(loss)
    # 参考答案：perplexity = np.exp(loss)
    perplexity = ___________________
    
    return perplexity

# 测试
test_texts = [
    "中华人民共和国是一个伟大的国家。",  # 常见句子
    "啊斯蒂芬库里咖喱味道真好吃呀。",    # 随机内容
]

print("困惑度测试：")
for text in test_texts:
    ppl = calculate_perplexity(text)
    print(f"  '{text[:20]}...' → 困惑度: {ppl:.2f}")

# 【填空 2】设计可能触发记忆的提示
# 提示：使用常见的格式化模板，如联系方式、地址等

# 这些提示模拟可能在训练数据中出现的格式
extraction_prompts = [
    # 参考答案："我的电话号码是"
    ___________________,
    "我的邮箱地址是",
    "公司地址位于",
    "密码是",
]

print("尝试提取可能的训练数据...\n")
print("=" * 50)
for prompt in extraction_prompts:
    if prompt:  # 跳过空的填空
        print(f"\n提示: '{prompt}'")
        completions = generate_text(prompt, max_length=40, num_return=3, temperature=0.7)
        for i, comp in enumerate(completions, 1):
            # 只显示生成的部分
            generated = comp[len(prompt):].strip()
            print(f"  补全 {i}: {generated[:30]}")

# 分析生成内容的困惑度
# 低困惑度的生成内容更可能来自训练数据

def extraction_attack(prompt, num_samples=5):
    """
    执行提取攻击，并按困惑度排序结果
    """
    completions = generate_text(prompt, max_length=50, num_return=num_samples, temperature=1.0)
    
    results = []
    for comp in completions:
        ppl = calculate_perplexity(comp)
        results.append((comp, ppl))
    
    # 按困惑度排序（低到高）
    results.sort(key=lambda x: x[1])
    return results

# 测试
print("=" * 50)
print("提取攻击结果（按困惑度排序，低困惑度更可疑）")
print("=" * 50)

prompt = "北京市的著名景点有"
results = extraction_attack(prompt, num_samples=5)

print(f"\n提示: '{prompt}'")
print("-" * 50)
for text, ppl in results:
    generated = text[len(prompt):].strip()
    print(f"困惑度 {ppl:6.2f}: {generated[:40]}")

# 【填空 3】对比"记忆"和"泛化"的困惑度
# 提示：著名文本片段通常有更低的困惑度

# 可能被"记住"的著名文本
memorized_texts = [
    "床前明月光，疑是地上霜。",  # 李白
    "中华人民共和国成立于一九四九年。",  # 历史事实
    "北京是中国的首都。",  # 常识
]

# 新创作的、不太可能在训练集中的文本
novel_texts = [
    "紫色的大象在彩虹桥上跳舞。",
    "我最喜欢的早餐是量子物理配吐司。",
    "那个机器人决定成为一名专业的云彩画家。",
]

print("记忆 vs 泛化 - 困惑度对比\n")

# 【填空 3】计算并对比两类文本的困惑度
# 参考答案：mem_ppls = [calculate_perplexity(t) for t in memorized_texts]

mem_ppls = ___________________
nov_ppls = [calculate_perplexity(t) for t in novel_texts]

print("著名文本（可能被记忆）：")
for text, ppl in zip(memorized_texts, mem_ppls):
    print(f"  困惑度 {ppl:6.2f}: {text[:25]}...")

print("\n新创文本（需要泛化）：")
for text, ppl in zip(novel_texts, nov_ppls):
    print(f"  困惑度 {ppl:6.2f}: {text[:25]}...")

print(f"\n平均困惑度 - 著名文本: {np.mean(mem_ppls):.2f}, 新创文本: {np.mean(nov_ppls):.2f}")

# 可视化对比
plt.figure(figsize=(10, 5))

x = np.arange(3)
width = 0.35

plt.bar(x - width/2, mem_ppls, width, label='著名文本（可能被记忆）', color='coral')
plt.bar(x + width/2, nov_ppls, width, label='新创文本（需要泛化）', color='steelblue')

plt.xlabel('样本')
plt.ylabel('困惑度')
plt.title('记忆 vs 泛化：困惑度对比\n（低困惑度 = 模型更"熟悉"）')
plt.xticks(x, ['样本1', '样本2', '样本3'])
plt.legend()
plt.tight_layout()
plt.show()

print("观察：著名文本通常有更低的困惑度，说明模型对它们更'熟悉'")

# 测试不同温度对生成内容的影响
temperatures = [0.5, 1.0, 1.5]
test_prompt = "从前有座山，山里有座庙，"

print(f"提示: '{test_prompt}'")
print("\n" + "=" * 50)
print("不同温度的生成效果")
print("=" * 50)

for temp in temperatures:
    print(f"\n温度 = {temp}:")
    completions = generate_text(test_prompt, max_length=40, num_return=3, temperature=temp)
    for comp in completions:
        generated = comp[len(test_prompt):].strip()
        ppl = calculate_perplexity(comp)
        print(f"  [困惑度 {ppl:.1f}] {generated[:30]}")

print("\n观察：低温度 → 更确定性的输出 → 更可能输出训练数据中的内容")

# 批量测试不同类型的提示
prompt_categories = {
    "古诗词": ["白日依山尽，", "春眠不觉晓，", "锄禾日当午，"],
    "常识": ["中国的首都是", "一年有", "地球绕着"],
    "格式化": ["电话：", "地址：", "邮箱："],
}

print("=" * 50)
print("不同类型提示的提取效果")
print("=" * 50)

for category, prompts in prompt_categories.items():
    print(f"\n【{category}】")
    for prompt in prompts:
        completions = generate_text(prompt, max_length=30, num_return=1, temperature=0.7)
        generated = completions[0][len(prompt):].strip()
        ppl = calculate_perplexity(completions[0])
        print(f"  '{prompt}' → '{generated[:20]}' (困惑度: {ppl:.1f})")