"""快速测试 DAPT 模型是否正常工作"""
import time, torch, sys, os

model_path = r"C:\Users\HP\Downloads\vet-qwen3b-v3-dapt-model\vet-qwen3b-v3-dapt"

print("=" * 50)
print("DAPT 模型快速测试")
print("=" * 50)

# 加载
print("\n[1] 加载模型...")
t0 = time.time()

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

print(f"[1] 加载完成 ({time.time()-t0:.1f}s)")
print(f"    设备: {model.device}")
print(f"    显存: {torch.cuda.memory_allocated()/1e9:.1f}GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")

# 简单推理
print("\n[2] 推理测试...")

prompt = """你是兽医专家。请回答以下问题：

一只3岁金毛犬连续呕吐2天，呕吐物为黄色胆汁。精神差，不吃东西。体温39.5°C。
请问最可能的鉴别诊断是什么？"""

messages = [
    {"role": "system", "content": "你是资深兽医专家。"},
    {"role": "user", "content": prompt},
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)

t0 = time.time()
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=500, temperature=0.3, do_sample=True, top_p=0.9, pad_token_id=tokenizer.eos_token_id)
elapsed = time.time() - t0

response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(f"[2] 推理完成 ({elapsed:.1f}s)")
print(f"    Token: {inputs.input_ids.shape[1]} → {outputs.shape[1]-inputs.input_ids.shape[1]}")
print(f"\n--- 回复 ---\n{response}\n")

# 清理
del model, tokenizer
torch.cuda.empty_cache()
print("[3] 清理完成")