import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

# 探针1: 使用小模型（GPT-2）
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 探针2: 简单的prompt（便于分析）
tracer_text = "The meaning of life is"

print(f"\n[实验设置]")
print(f"输入文本：\"{tracer_text}\"")

# ============================================================
# 步骤1：获取下一个token的logits
# ============================================================
inputs = tokenizer(tracer_text, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits[0, -1, :]  # 最后一个位置的logits

print(f"Logits shape: {logits.shape}")  # [50257]

# 找出top 10
top_values, top_indices = torch.topk(logits, 10)
print(f"\nTop 10 Logits（原始分数）：")
for i, (idx, val) in enumerate(zip(top_indices, top_values)):
    token = tokenizer.decode([idx])
    print(f"  {i+1}. {token:<20} logit={val.item():>8.3f}")

# ============================================================
# 步骤2：Softmax转为概率分布
# ============================================================
probs = F.softmax(logits, dim=-1)

print(f"\nTop 10 概率：")
top_probs, top_indices = torch.topk(probs, 10)
for i, (idx, prob) in enumerate(zip(top_indices, top_probs)):
    token = tokenizer.decode([idx])
    print(f"  {i+1}. {token:<20} prob={prob.item():>.6f}")

# ============================================================
# 步骤3：对比不同采样策略
# ============================================================
def greedy_sample(probs):
    """贪婪采样：选择概率最大的"""
    return torch.argmax(probs).item()

def temperature_sample(probs, temperature=1.0):
    """温度采样：调整概率分布的平滑度"""
    scaled_logits = torch.log(probs) / temperature
    scaled_probs = F.softmax(scaled_logits, dim=-1)
    return torch.multinomial(scaled_probs, 1).item()

def top_k_sample(probs, k=50):
    """Top-K采样：只从前k个候选中采样"""
    top_k_probs, top_k_indices = torch.topk(probs, k)
    top_k_probs = top_k_probs / top_k_probs.sum()
    sampled_idx = torch.multinomial(top_k_probs, 1).item()
    return top_k_indices[sampled_idx].item()

# ============================================================
# 步骤4：多次采样观察分布
# ============================================================
import random
random.seed(42)
torch.manual_seed(42)

print(f"\n[多次采样观察（各10次）]")
strategies = [
    ("Greedy", lambda: greedy_sample(probs)),
    ("Temperature=0.7", lambda: temperature_sample(probs, 0.7)),
    ("Temperature=1.5", lambda: temperature_sample(probs, 1.5)),
    ("Top-K=50", lambda: top_k_sample(probs, 50)),
]

for strategy_name, strategy_fn in strategies:
    samples = [strategy_fn() for _ in range(10)]
    from collections import Counter
    counter = Counter(samples)

    print(f"\n{strategy_name}:")
    print(f"  采样了{len(counter)}个不同的token")
    most_common = counter.most_common(1)[0]
    token = tokenizer.decode([most_common[0]])
    print(f"  最常见：{token} ({most_common[1]}次)")

# ============================================================
# 步骤5：可视化Temperature的影响
# ============================================================
print(f"\n[Temperature对概率分布的影响]")

top_probs_orig, top_indices_orig = torch.topk(probs, 5)
print(f"\n原始概率分布（Top 5）：")
for i, (idx, prob) in enumerate(zip(top_indices_orig, top_probs_orig)):
    token = tokenizer.decode([idx])
    print(f"  {i+1}. {token:<20} {prob.item():.6f}")

for temp in [0.5, 0.7, 1.0, 1.5, 2.0]:
    scaled_probs = F.softmax(torch.log(probs) / temp, dim=-1)
    top_probs_temp, _ = torch.topk(scaled_probs, 5)

    print(f"\nTemperature={temp}:")
    for i, (idx, prob) in enumerate(zip(top_indices_orig, top_probs_temp)):
        token = tokenizer.decode([idx])
        print(f"  {i+1}. {token:<20} {prob.item():.6f}")

# ============================================================
# 步骤6：完整生成对比
# ============================================================
print(f"\n[完整生成对比（生成20个token）]")

generation_configs = [
    {"do_sample": False, "name": "Greedy"},
    {"do_sample": True, "temperature": 0.7, "name": "Temp=0.7"},
    {"do_sample": True, "temperature": 1.5, "name": "Temp=1.5"},
    {"do_sample": True, "top_k": 50, "name": "Top-K=50"},
    {"do_sample": True, "top_p": 0.9, "name": "Top-P=0.9"},
]

for config in generation_configs:
    name = config.pop("name")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=20, **config)
    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"\n{name}:")
    print(f"  {generated}")

# 步骤1: 运行探针
python sampling_tracer.py

# 步骤2: 测试确定性输出
python -c "
import torch
from transformers import set_seed, pipeline

set_seed(42)
generator = pipeline('text-generation', model='gpt2')

# 两次生成应该相同
for i in range(2):
    result = generator('The future of AI is', max_new_tokens=20,
                      do_sample=False)
    print(f'Run {i+1}: {result[0][\"generated_text\"]}\n')
"