text_classification/test.py

# BERT文本分类测试脚本
import torch
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from tqdm import tqdm
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, precision_score, recall_score, f1_score
import os
import seaborn as sns
import matplotlib.pyplot as plt
import gc
import time


class TextClassificationTester:
    """BERT文本分类测试器"""
    
    def __init__(self, model_path, is_english=False, num_classes=2, max_length=256):
        """
        初始化测试器
        
        Args:
            model_path: 模型权重文件路径
            is_english: 是否为英文模型
            num_classes: 分类类别数
            max_length: 最大序列长度
        """
        # 选择模型
        if is_english:
            self.model_name = "bert-base-uncased"
        else:
            self.model_name = "bert-base-chinese"
        self.is_english = is_english
        
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.max_length = max_length
        self.num_classes = num_classes
        
        print(f"\n{'='*60}")
        print("初始化测试器")
        print(f"{'='*60}")
        print(f"使用设备: {self.device}")
        print(f"模型名称: {self.model_name}")
        print(f"模型路径: {model_path}")
        print(f"类别数: {num_classes}")
        print(f"最大序列长度: {max_length}")
        
        # 加载模型和分词器
        print("\n加载模型和分词器...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.model_name, 
            num_labels=num_classes
        ).to(self.device)
        
        # 加载训练好的权重
        if os.path.exists(model_path):
            print(f"加载模型权重: {model_path}")
            self.model.load_state_dict(torch.load(model_path, map_location=self.device))
            print("✓ 模型权重加载成功")
        else:
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        
        # 设置为评估模式
        self.model.eval()
        
    def load_test_data(self, test_csv_file):
        """
        加载测试数据
        
        Args:
            test_csv_file: 测试数据CSV文件路径
        """
        print(f"\n加载测试数据: {test_csv_file}")
        self.test_data = pd.read_csv(test_csv_file)
        self.test_data = self.test_data.dropna()
        
        print(f"测试样本数: {len(self.test_data)}")
        print(f"列名: {self.test_data.columns.tolist()}")
        
        # 检查必要的列
        if 'record' not in self.test_data.columns:
            raise ValueError("数据文件中缺少 'record' 列")
        if 'label' not in self.test_data.columns:
            raise ValueError("数据文件中缺少 'label' 列")
        
        # 统计标签分布
        label_counts = self.test_data['label'].value_counts().sort_index()
        print(f"\n标签分布:")
        for label, count in label_counts.items():
            print(f"  类别 {label}: {count} ({count/len(self.test_data)*100:.1f}%)")
    
    def predict(self, batch_size=32):
        """
        对测试数据进行预测
        
        Args:
            batch_size: 批次大小
        """
        print(f"\n开始预测 (batch_size={batch_size})...")
        
        self.predictions = []
        self.labels = []
        self.texts = []
        
        with torch.no_grad():
            for i in tqdm(range(0, len(self.test_data), batch_size), desc="预测中"):
                batch_texts = self.test_data['record'].iloc[i:i+batch_size].tolist()
                batch_labels = self.test_data['label'].iloc[i:i+batch_size].tolist()
                
                # 分词
                inputs = self.tokenizer(
                    batch_texts,
                    max_length=self.max_length,
                    padding='max_length',
                    truncation=True,
                    return_tensors='pt'
                )
                
                # 移动到设备
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                
                # 预测
                outputs = self.model(**inputs)
                preds = torch.argmax(outputs.logits, dim=-1).cpu().numpy()
                
                self.predictions.extend(preds)
                self.labels.extend(batch_labels)
                self.texts.extend(batch_texts)
        
        self.predictions = np.array(self.predictions)
        self.labels = np.array(self.labels)
        
        print(f"✓ 预测完成，共 {len(self.predictions)} 个样本")
    
    def evaluate(self):
        """
        评估模型性能并输出详细指标
        """
        print(f"\n{'='*60}")
        print("评估结果")
        print(f"{'='*60}")
        
        # 计算各项指标
        accuracy = accuracy_score(self.labels, self.predictions)
        precision_macro = precision_score(self.labels, self.predictions, average='macro')
        recall_macro = recall_score(self.labels, self.predictions, average='macro')
        f1_macro = f1_score(self.labels, self.predictions, average='macro')
        
        precision_weighted = precision_score(self.labels, self.predictions, average='weighted')
        recall_weighted = recall_score(self.labels, self.predictions, average='weighted')
        f1_weighted = f1_score(self.labels, self.predictions, average='weighted')
        
        # 输出总体指标
        print("\n【总体指标】")
        print(f"  准确率 (Accuracy): {accuracy:.4f} ({accuracy*100:.2f}%)")
        print(f"\n  Macro平均:")
        print(f"    精确率 (Precision): {precision_macro:.4f}")
        print(f"    召回率 (Recall): {recall_macro:.4f}")
        print(f"    F1分数: {f1_macro:.4f}")
        print(f"\n  Weighted平均:")
        print(f"    精确率 (Precision): {precision_weighted:.4f}")
        print(f"    召回率 (Recall): {recall_weighted:.4f}")
        print(f"    F1分数: {f1_weighted:.4f}")
        
        # 输出每个类别的详细指标
        print("\n【各类别详细指标】")
        print(classification_report(self.labels, self.predictions, digits=4))
        
        # 混淆矩阵
        print("\n【混淆矩阵】")
        cm = confusion_matrix(self.labels, self.predictions)
        print(cm)
        
        # 绘制混淆矩阵热力图
        self._plot_confusion_matrix(cm)
        
        return {
            'accuracy': accuracy,
            'precision_macro': precision_macro,
            'recall_macro': recall_macro,
            'f1_macro': f1_macro,
            'precision_weighted': precision_weighted,
            'recall_weighted': recall_weighted,
            'f1_weighted': f1_weighted
        }
    
    def _plot_confusion_matrix(self, cm):
        """
        绘制混淆矩阵热力图
        """
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=True)
        plt.title('Confusion Matrix')
        plt.ylabel('True Label')
        plt.xlabel('Predicted Label')
        
        # 保存图片
        if self.is_english:
            output_path = './confusion_matrix_en.png'
        else:
            output_path = './confusion_matrix_zh.png'
        plt.savefig(output_path, dpi=300, bbox_inches='tight')
        print(f"\n✓ 混淆矩阵已保存至: {output_path}")
        plt.close()
    
    def save_errors(self, output_csv_file='errors.csv'):
        """
        保存预测错误的样本到CSV文件
        
        Args:
            output_csv_file: 输出CSV文件路径
        """
        print(f"\n保存错误预测样本...")
        
        # 找出预测错误的样本
        error_indices = np.where(self.predictions != self.labels)[0]
        
        if len(error_indices) == 0:
            print("✓ 没有预测错误的样本")
            return
        
        # 创建错误样本DataFrame
        errors_df = pd.DataFrame({
            'record': [self.texts[i] for i in error_indices],
            'predict': self.predictions[error_indices],
            'label': self.labels[error_indices]
        })
        
        # 保存到CSV
        errors_df.to_csv(output_csv_file, index=False, encoding='utf-8-sig')
        
        print(f"✓ 错误样本已保存至: {output_csv_file}")
        print(f"  错误样本数: {len(errors_df)} / {len(self.labels)} ({len(errors_df)/len(self.labels)*100:.2f}%)")
        
        # 统计错误类型
        print(f"\n【错误类型统计】")
        error_types = errors_df.groupby(['label', 'predict']).size().reset_index(name='count')
        for _, row in error_types.iterrows():
            print(f"  真实标签 {row['label']} → 预测为 {row['predict']}: {row['count']} 个")
    
    def run_test(self, test_csv_file, output_csv_file='errors.csv', batch_size=32):
        """
        运行完整的测试流程
        
        Args:
            test_csv_file: 测试数据文件路径
            output_csv_file: 错误样本输出文件路径
            batch_size: 批次大小
        """
        # 加载测试数据
        self.load_test_data(test_csv_file)
        
        # 预测
        self.predict(batch_size)
        
        # 评估
        metrics = self.evaluate()
        
        # 保存错误样本
        self.save_errors(output_csv_file)
        
        print(f"\n{'='*60}")
        print("测试完成!")
        print(f"{'='*60}")
        
        return metrics


if __name__ == '__main__':
    # 示例用法
    tester = TextClassificationTester(
        model_path='./checkpoints/bert-base-chinese_best_acc.pth',  # 模型路径
        is_english=False,  # 中文模型
        num_classes=2,     # 二分类
        max_length=256     # 最大序列长度
    )
    
    # 运行测试
    metrics = tester.run_test(
        test_csv_file='data_zh.csv',  # 测试数据文件
        output_csv_file='errors_zh.csv',  # 错误样本输出文件
        batch_size=32                  # 批次大小
    )

    del tester, metrics
    gc.collect()  # 强制 Python 垃圾回收
    torch.cuda.empty_cache()  # 清空 CUDA 缓存
    time.sleep(2)  # 短暂等待即可，不需要太久

    tester = TextClassificationTester(
        model_path='./checkpoints/bert-base-uncased_best_acc.pth',  # 模型路径
        is_english=True,  # 英文模型
        num_classes=2,  # 二分类
        max_length=256  # 最大序列长度
    )

    # 运行测试
    metrics = tester.run_test(
        test_csv_file='data_en.csv',  # 测试数据文件
        output_csv_file='errors_en.csv',  # 错误样本输出文件
        batch_size=32  # 批次大小
    )