import pandas as pd
import numpy as np
import torch
from sklearn.metrics import accuracy_score, f1_score, classification_report
from transformers import AutoModelForCausalLM, AutoTokenizer

def evaluate(model_path: str, template: str = "qwen2"):
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True
    )

    df = pd.read_csv("sst2_validation.csv")
    texts = df["text"].tolist()
    labels = df["label"].tolist()

    preds = []
    for text in texts:
        prompt = (
            "<|im_start|>user\n"
            "请判断以下电影评论的情感倾向，只回答'正面'或'负面'：\n\n"
            f"评论：{text}\n情感：<|im_end|>\n"
            "<|im_start|>assistant\n"
        )
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=6, do_sample=False,
                                 pad_token_id=tokenizer.eos_token_id)
        resp = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
        if "正面" in resp:
            preds.append(1)
        elif "负面" in resp:
            preds.append(0)
        else:
            preds.append(np.random.randint(0, 2))

    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds)
    rep = classification_report(labels, preds, target_names=["负面","正面"], output_dict=True)
    print(f"准确率: {acc:.3f}  F1: {f1:.3f}  精确率(正面): {rep['正面']['precision']:.3f}  召回率(正面): {rep['正面']['recall']:.3f}")
    return {"accuracy": acc, "f1": f1, "precision": rep['正面']['precision'], "recall": rep['正面']['recall']}

if __name__ == "__main__":
    print("\n[LoRA 合并模型]")
    evaluate("./output/sft_qwen_lora_merged")

    print("\n[全参数模型]")
    evaluate("./output/sft_qwen_full/checkpoint-3000")