import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

PRUNED_MODEL_PATH = "./QiMing-Polaris-Qwen3-4B-HardPruned"

print("正在加载剪枝后的模型到 GPU...")
tokenizer = AutoTokenizer.from_pretrained(PRUNED_MODEL_PATH, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    PRUNED_MODEL_PATH,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    trust_remote_code=True,
)

prompt = "What is the 'Burden-based Training' method?"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

start_time = time.time()
with torch.no_grad():
    outputs = model.generate(
        **inputs, max_new_tokens=256, temperature=0.7, do_sample=True
    )
end_time = time.time()

response = tokenizer.decode(
    outputs[0][inputs.input_ids.shape[1] :], skip_special_tokens=True
)

print("\n==== 剪枝模型回答 ====")
print(response)
print(f"\n耗时: {end_time - start_time:.4f} 秒")