final_report.py

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

# --- ⚙️ 配置区 ---
base_model_path = "./models/gemma-3-270m-it"
lora_path = "./tmodels/gemma-3-270m-it-FT-lora"
test_prompt = "you are fox,give say a ..."
# -----------------

def inject_chat_template(tokenizer):
    tokenizer.chat_template = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/model/user/model/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '<start_of_turn>user\n' + message['content'] | trim + '<end_of_turn>\n' }}{% elif message['role'] == 'model' %}{{ '<start_of_turn>model\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<start_of_turn>model\n' }}{% endif %}"

def print_top_k(probs, tokenizer, k=5):
    top_probs, top_indices = torch.topk(probs, k)
    for i in range(k):
        token = tokenizer.decode([top_indices[i]]).replace('\n', '\\n')
        print(f"    - Rank {i+1}: [{token}] \t 概率: {top_probs[i].item()*100:.2f}%")

# 💥【修正点】函数名已从 generate_decision_chain_report 改为 full_audit_analysis
def full_audit_analysis(model_name, model, tokenizer, prompt):
    print("\n" + "="*80)
    print(f"📄 开始对模型 [{model_name}] 进行终极决策链审计")
    print("="*80)

    # --- 准备工作 ---
    messages = [{"role": "user", "content": prompt}]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=True)

    hidden_states = outputs.hidden_states

    if isinstance(model, PeftModel):
        base_model_inner = model.base_model.model
    else:
        base_model_inner = model
    final_norm_layer = base_model_inner.model.norm
    lm_head_layer = base_model_inner.lm_head

    # --- 开始逐帧回放 ---

    # 阶段 1 & 2: 从 Input 到 Layer 18 Raw
    print("\n[阶段 1 & 2] 从输入到 Layer 18 Raw (部门主管的最终提案形成过程)")
    print("-" * 80)
    print("这是每一层计算完毕后，未经任何修正的“原始念头”：")
    for i, layer_hidden in enumerate(hidden_states):
        raw_vec = layer_hidden[0, -1, :].to(base_model_inner.dtype)
        raw_logits = lm_head_layer(raw_vec)
        raw_probs = F.softmax(raw_logits, dim=-1)
        top_prob, top_id = torch.topk(raw_probs, 1)
        top_word = tokenizer.decode([top_id[0]]).replace('\n', '\\n')

        layer_name = "Embed (Raw)" if i == 0 else f"L-{i} (RAW)"
        print(f"  - {layer_name:<12}: 最可能的词是 [{top_word}] ({top_prob[0].item()*100:.1f}%)")
    print("-" * 80)

    # 阶段 3: Layer 18 Raw -> Final Norm -> Normalized Vector
    print("\n[阶段 3] Layer 18 Raw -> Final Norm (技术总监审查并修改提案)")
    print("-" * 80)
    raw_last_layer_vec = hidden_states[-1][0, -1, :].to(base_model_inner.dtype)
    normalized_vec = final_norm_layer(raw_last_layer_vec)
    cos_sim = F.cosine_similarity(raw_last_layer_vec.unsqueeze(0), normalized_vec.unsqueeze(0))
    print("1. 部门主管 (L-18 Raw) 提交的原始提案翻译如下:")
    print_top_k(F.softmax(lm_head_layer(raw_last_layer_vec), dim=-1), tokenizer)
    print(f"\n2. 技术总监 (Final Norm) 对提案向量进行了修正。")
    print(f"   (向量方向偏移度: {cos_sim.item():.4f}, 1.0 表示未修正)")
    print("-" * 80)

    # 阶段 4: Normalized Vector -> LM Head -> Logits
    print("\n[阶段 4] Normalized Vector -> LM Head (秘书处将修改后的提案翻译成具体方案)")
    print("-" * 80)
    print("技术总监修正后的提案，经秘书处翻译，内容变为:")
    normed_logits = lm_head_layer(normalized_vec)
    normed_probs = F.softmax(normed_logits, dim=-1)
    print_top_k(normed_probs, tokenizer)
    print("-" * 80)

    # 阶段 5: Logits -> Decoding Strategy -> Final Token
    print("\n[阶段 5] CEO (Decoding Strategy) 结合所有信息做出最终裁决")
    print("-" * 80)
    print("1. CEO 在做决定前，参考的最终概率分布 (outputs.logits) 是:")
    final_logits_for_generation = outputs.logits[0, -1, :]
    final_probs_for_generation = F.softmax(final_logits_for_generation, dim=-1)
    print_top_k(final_probs_for_generation, tokenizer)

    print("\n2. 经过对上下文、风险和连贯性的最终权衡，CEO 发表了公开声明:")
    gen_output = model.generate(
        **inputs,
        max_new_tokens=100,
        do_sample=False,
        repetition_penalty=1.1
    )
    response = tokenizer.decode(gen_output[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)
    print("   >>> " + response.strip())
    print("-" * 80)
    print(f"✅ 模型 [{model_name}] 决策链审计完成。")


# --- 主程序 ---
print("🚀 启动终极决策链全景报告生成器...")
print(f"📝 测试 Prompt: '{test_prompt}'")

tokenizer = AutoTokenizer.from_pretrained(base_model_path)
inject_chat_template(tokenizer)

# --- 审计 Base-IT 模型 ---
base_model = AutoModelForCausalLM.from_pretrained(base_model_path, device_map="auto")
full_audit_analysis("Base-IT (老黄牛)", base_model, tokenizer, test_prompt)

# --- 审计 FT 模型 ---
base_model_for_ft = AutoModelForCausalLM.from_pretrained(base_model_path, device_map="auto")
ft_model = PeftModel.from_pretrained(base_model_for_ft, lora_path)
full_audit_analysis("FT (监工介入)", ft_model, tokenizer, test_prompt)

print("\n\n" + "="*80)
print("🎉 所有审计工作已完成。")
print("="*80)