"""
TFMF 官方 Qwen3.5-4B 模型探针脚本 - 使用 enable_thinking=False 关闭 CoT
"""

import json
import time
import torch
import re
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer

# ==================== 配置区 ====================
MODEL_PATH = "/data/coding/TFMF"
DATA_PATH = "/data/coding/TFMF/语文教师_语文_高二_v1.jsonl"
MAX_NEW_TOKENS = 512
TEMPERATURE = 0.7
TOP_P = 0.9

# ==================== 加载模型 ====================
print("=" * 60)
print("开始加载官方 Qwen3.5-4B 模型...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    device_map="cuda",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)
model.eval()

print(f"模型加载完成! 显存占用: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print("=" * 60)

# ==================== 读取测试数据 ====================
test_samples = []
try:
    with open(DATA_PATH, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i >= 3:
                break
            if line.strip():
                test_samples.append(json.loads(line))
    print(f"\n成功加载 {len(test_samples)} 条测试样本")
except FileNotFoundError:
    print(f"\n警告: 未找到数据文件 {DATA_PATH}")
    test_samples = [{
        "system": "你是一位高中语文教师，教学风格直白浅近。",
        "user": "老师，什么是归谬法？",
        "assistant": "归谬法就是先假设对方的观点正确，然后推导出荒谬结论。"
    }]

print("-" * 60)

# ==================== 推理函数 ====================
def generate_response(system_prompt, user_query, stream=False):
    """
    生成回答，通过 enable_thinking=False 关闭 CoT
    """
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_query},
    ]
    
    # ===== 关键修改：在 apply_chat_template 中传入 enable_thinking=False =====
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,  # ← 官方开关，关闭 CoT
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) if stream else None
    
    print(f"\n【生成中...】")
    start_time = time.time()
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            temperature=TEMPERATURE,
            top_p=TOP_P,
            do_sample=True,
            use_cache=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1,
            streamer=streamer,
        )
    
    elapsed = time.time() - start_time
    
    full_response = tokenizer.decode(
        outputs[0][inputs['input_ids'].shape[1]:],
        skip_special_tokens=True
    )
    
    # 安全兜底：如果还有残留的 <think> 标签，用正则清理掉
    full_response = re.sub(r'<think>.*?</think>', '', full_response, flags=re.DOTALL).strip()
    
    return full_response, elapsed

# ==================== 执行测试 ====================
print("\n开始推理测试...")
print("=" * 60)

for idx, sample in enumerate(test_samples, 1):
    system = sample.get("system", "")
    user = sample.get("user", "")
    ground_truth = sample.get("assistant", "")
    
    print(f"\n【测试 {idx}】")
    print(f"用户问题: {user}")
    
    response, elapsed = generate_response(system, user, stream=False)
    
    print(f"\n【模型回答】（耗时 {elapsed:.2f}秒）")
    print(response if response else "（模型未生成有效回答）")
    
    print(f"\n【教师参考答案】")
    print(ground_truth[:300] + "..." if len(ground_truth) > 300 else ground_truth)
    print("-" * 60)

print("\n" + "=" * 60)
print("【探针完成】")
print(f"模型: Qwen3.5-4B (官方标准版)")
print(f"精度: bfloat16")
print(f"显存占用: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print("=" * 60)