evaluate.py

"""
评估脚本 - 在测试集上评估训练好的模型
"""

import os
import json
import argparse
from typing import List, Dict, Tuple
from collections import Counter

import torch
from tqdm import tqdm
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    confusion_matrix,
    classification_report,
)
from transformers import AutoModelForCausalLM, AutoTokenizer

from config import (
    DataConfig,
    build_prompt_for_inference,
    LABEL2ID,
    ID2LABEL,
)


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="评估意图识别模型")

    parser.add_argument(
        "--model_path",
        type=str,
        required=True,
        help="训练好的模型路径",
    )
    parser.add_argument(
        "--test_file",
        type=str,
        default=None,
        help="测试集文件路径（默认使用 config 中的路径）",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=32,
        help="评估时的 batch size",
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=5,
        help="生成的最大 token 数",
    )
    parser.add_argument(
        "--output_file",
        type=str,
        default=None,
        help="输出详细结果的文件路径",
    )
    parser.add_argument(
        "--show_errors",
        action="store_true",
        help="显示错误案例",
    )
    parser.add_argument(
        "--max_errors",
        type=int,
        default=20,
        help="最多显示的错误案例数",
    )

    return parser.parse_args()


def load_test_data(test_file: str) -> List[Dict]:
    """加载测试数据"""
    data = []
    with open(test_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                item = json.loads(line)
                data.append(item)
    return data


def extract_label_from_output(output: str) -> int:
    """
    从模型输出中提取标签

    Args:
        output: 模型生成的文本

    Returns:
        提取的标签（0 或 1），如果无法提取则返回 -1
    """
    # 清理输出
    output = output.strip()

    # 尝试直接提取数字
    if output.startswith("1"):
        return 1
    elif output.startswith("0"):
        return 0

    # 尝试在输出中找到 0 或 1
    for char in output:
        if char == "1":
            return 1
        elif char == "0":
            return 0

    # 无法提取
    return -1


def evaluate_model(
    model,
    tokenizer,
    test_data: List[Dict],
    batch_size: int = 32,
    max_new_tokens: int = 5,
) -> Tuple[List[int], List[int], List[str], List[str]]:
    """
    对模型进行评估

    Args:
        model: 模型
        tokenizer: 分词器
        test_data: 测试数据
        batch_size: batch size
        max_new_tokens: 最大生成 token 数

    Returns:
        (预测标签列表, 真实标签列表, 原始文本列表, 模型输出列表)
    """
    model.eval()
    device = next(model.parameters()).device

    predictions = []
    references = []
    texts = []
    outputs = []

    # 分批处理
    for i in tqdm(range(0, len(test_data), batch_size), desc="评估中"):
        batch = test_data[i : i + batch_size]

        # 构建 prompts
        prompts = [build_prompt_for_inference(item["text"]) for item in batch]
        true_labels = [item["label"] for item in batch]
        batch_texts = [item["text"] for item in batch]

        # 分词
        inputs = tokenizer(
            prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=256,
        ).to(device)

        # 生成
        with torch.no_grad():
            generated = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,  # 使用贪心解码
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        # 解码输出（只取新生成的部分）
        input_length = inputs["input_ids"].shape[1]
        for j, gen in enumerate(generated):
            # 只取生成的新 token
            new_tokens = gen[input_length:]
            output_text = tokenizer.decode(new_tokens, skip_special_tokens=True)

            # 提取标签
            pred_label = extract_label_from_output(output_text)

            predictions.append(pred_label)
            references.append(true_labels[j])
            texts.append(batch_texts[j])
            outputs.append(output_text)

    return predictions, references, texts, outputs


def print_evaluation_results(
    predictions: List[int],
    references: List[int],
    texts: List[str],
    outputs: List[str],
    show_errors: bool = False,
    max_errors: int = 20,
):
    """打印评估结果"""
    # 过滤掉无法解析的样本
    valid_indices = [i for i, p in enumerate(predictions) if p != -1]
    invalid_count = len(predictions) - len(valid_indices)

    valid_preds = [predictions[i] for i in valid_indices]
    valid_refs = [references[i] for i in valid_indices]

    print("\n" + "=" * 60)
    print("评估结果")
    print("=" * 60)

    print(f"\n总样本数: {len(predictions)}")
    print(f"有效样本数: {len(valid_preds)}")
    if invalid_count > 0:
        print(f"无法解析的样本数: {invalid_count} ({invalid_count/len(predictions)*100:.2f}%)")

    if len(valid_preds) == 0:
        print("错误：没有有效的预测结果！")
        return

    # 计算指标
    accuracy = accuracy_score(valid_refs, valid_preds)
    precision = precision_score(valid_refs, valid_preds, pos_label=1, zero_division=0)
    recall = recall_score(valid_refs, valid_preds, pos_label=1, zero_division=0)
    f1 = f1_score(valid_refs, valid_preds, pos_label=1, zero_division=0)

    print("\n--- 主要指标 ---")
    print(f"准确率 (Accuracy):  {accuracy:.4f} ({accuracy*100:.2f}%)")
    print(f"精确率 (Precision): {precision:.4f} ({precision*100:.2f}%)")
    print(f"召回率 (Recall):    {recall:.4f} ({recall*100:.2f}%)")
    print(f"F1 Score:           {f1:.4f} ({f1*100:.2f}%)")

    # 目标检查
    print("\n--- 目标达成情况 ---")
    targets = {"准确率": (accuracy, 0.95), "精确率": (precision, 0.93), "召回率": (recall, 0.93), "F1": (f1, 0.93)}
    for name, (value, target) in targets.items():
        status = "✓ 达标" if value >= target else "✗ 未达标"
        print(f"{name}: {value*100:.2f}% (目标: {target*100:.0f}%) {status}")

    # 混淆矩阵
    cm = confusion_matrix(valid_refs, valid_preds)
    print("\n--- 混淆矩阵 ---")
    print("           预测=0  预测=1")
    print(f"实际=0     {cm[0][0]:5d}   {cm[0][1]:5d}")
    print(f"实际=1     {cm[1][0]:5d}   {cm[1][1]:5d}")

    # 分类报告
    print("\n--- 详细分类报告 ---")
    print(classification_report(valid_refs, valid_preds, target_names=["不生成游戏(0)", "生成游戏(1)"]))

    # 统计预测分布
    pred_counter = Counter(valid_preds)
    ref_counter = Counter(valid_refs)
    print("--- 预测分布 ---")
    print(f"预测为 0: {pred_counter.get(0, 0)} ({pred_counter.get(0, 0)/len(valid_preds)*100:.2f}%)")
    print(f"预测为 1: {pred_counter.get(1, 0)} ({pred_counter.get(1, 0)/len(valid_preds)*100:.2f}%)")
    print(f"实际为 0: {ref_counter.get(0, 0)} ({ref_counter.get(0, 0)/len(valid_refs)*100:.2f}%)")
    print(f"实际为 1: {ref_counter.get(1, 0)} ({ref_counter.get(1, 0)/len(valid_refs)*100:.2f}%)")

    # 显示错误案例
    if show_errors:
        print("\n--- 错误案例 ---")
        error_count = 0
        for i in range(len(predictions)):
            if predictions[i] != references[i]:
                if error_count >= max_errors:
                    remaining = sum(1 for j in range(i, len(predictions)) if predictions[j] != references[j])
                    print(f"\n... 还有 {remaining} 个错误案例未显示")
                    break

                print(f"\n[错误 {error_count + 1}]")
                print(f"  输入: {texts[i]}")
                print(f"  预测: {predictions[i]} (输出: '{outputs[i]}')")
                print(f"  实际: {references[i]}")
                error_count += 1

    print("\n" + "=" * 60)


def save_detailed_results(
    predictions: List[int],
    references: List[int],
    texts: List[str],
    outputs: List[str],
    output_file: str,
):
    """保存详细结果到文件"""
    results = []
    for i in range(len(predictions)):
        results.append({
            "text": texts[i],
            "prediction": predictions[i],
            "reference": references[i],
            "model_output": outputs[i],
            "correct": predictions[i] == references[i],
        })

    with open(output_file, "w", encoding="utf-8") as f:
        for item in results:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    print(f"\n详细结果已保存到: {output_file}")


def main():
    """主函数"""
    args = parse_args()

    # 确定测试文件路径
    if args.test_file:
        test_file = args.test_file
    else:
        data_config = DataConfig()
        test_file = data_config.test_path

    print(f"测试文件: {test_file}")
    print(f"模型路径: {args.model_path}")

    # 加载测试数据
    print("\n加载测试数据...")
    test_data = load_test_data(test_file)
    print(f"测试样本数: {len(test_data)}")

    # 加载模型和分词器
    print("\n加载模型...")
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        padding_side="left",  # 生成时使用左 padding
    )

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        torch_dtype=torch.float16,
        device_map="auto",
    )

    print(f"模型参数量: {model.num_parameters():,}")

    # 评估模型
    predictions, references, texts, outputs = evaluate_model(
        model,
        tokenizer,
        test_data,
        batch_size=args.batch_size,
        max_new_tokens=args.max_new_tokens,
    )

    # 打印结果
    print_evaluation_results(
        predictions,
        references,
        texts,
        outputs,
        show_errors=args.show_errors,
        max_errors=args.max_errors,
    )

    # 保存详细结果
    if args.output_file:
        save_detailed_results(predictions, references, texts, outputs, args.output_file)


if __name__ == "__main__":
    main()