Delete scripts/inference_longemotion.py

scripts/inference_longemotion.py

@@ -1,477 +0,0 @@
-"""
-LongEmotion 测试集推理脚本
-专门处理 LongEmotion 格式的长文本多段落情感检测
-"""
-import os
-import sys
-import json
-import torch
-from pathlib import Path
-from typing import Dict, List, Any, Optional
-from tqdm import tqdm
-from collections import Counter
-
-# 添加项目根目录到路径
-sys.path.append(str(Path(__file__).parent.parent.parent))
-
-from transformers import BertTokenizer, BertModel
-
-
-class LongEmotionInference:
-    """LongEmotion 推理器"""
-    
-    # 情感标签映射 (dair数据集的6类情感)
-    EMOTION_LABELS = {
-        0: "sadness",
-        1: "joy",
-        2: "love",
-        3: "anger",
-        4: "fear",
-        5: "surprise"
-    }
-    
-    def __init__(
-        self,
-        model_path: str,
-        device: str = "cuda" if torch.cuda.is_available() else "cpu",
-        max_length: int = 512,
-        batch_size: int = 16
-    ):
-        """
-        初始化推理器
-        
-        Args:
-            model_path: 模型权重文件路径 (best_model.pt)
-            device: 设备
-            max_length: 最大序列长度
-            batch_size: 批次大小
-        """
-        # 检查设备可用性
-        if device == "cuda" and not torch.cuda.is_available():
-            print("警告: CUDA不可用，切换到CPU")
-            device = "cpu"
-        
-        self.device = device
-        self.max_length = max_length
-        self.batch_size = batch_size
-        
-        print(f"正在加载模型从 {model_path}...")
-        print(f"使用设备: {device}")
-        
-        # 加载分词器
-        try:
-            self.tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
-        except Exception as e:
-            print(f"加载分词器失败，尝试从本地加载: {e}")
-            self.tokenizer = BertTokenizer.from_pretrained("bert-base-chinese", local_files_only=False)
-        
-        # 加载模型
-        self.model = self._load_model(model_path)
-        self.model.eval()
-        
-        print("模型加载完成！")
-    
-    def _load_model(self, model_path: str):
-        """加载训练好的模型 - 使用与训练时相同的简单结构"""
-        try:
-            from transformers import AutoModel
-            import torch.nn as nn
-            
-            # 定义简单的分类器（与simple_train.py中的结构完全一致）
-            class SimpleEmotionClassifier(nn.Module):
-                """简单情感分类器 - 单层Linear"""
-                def __init__(self, model_name, num_labels=6):
-                    super().__init__()
-                    self.bert = AutoModel.from_pretrained(model_name)
-                    self.dropout = nn.Dropout(0.1)
-                    self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)
-                
-                def forward(self, input_ids, attention_mask):
-                    outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-                    pooled_output = outputs.pooler_output
-                    pooled_output = self.dropout(pooled_output)
-                    logits = self.classifier(pooled_output)
-                    return logits
-            
-            # 创建模型
-            print("创建模型结构（简单单层分类器）...")
-            model = SimpleEmotionClassifier(
-                model_name="bert-base-chinese",
-                num_labels=6  # dair数据集的6类
-            )
-            
-            # 加载权重
-            print(f"加载模型权重...")
-            checkpoint = torch.load(model_path, map_location=self.device)
-            
-            # 处理不同的保存格式
-            if isinstance(checkpoint, dict):
-                if 'model_state_dict' in checkpoint:
-                    model.load_state_dict(checkpoint['model_state_dict'])
-                elif 'state_dict' in checkpoint:
-                    model.load_state_dict(checkpoint['state_dict'])
-                else:
-                    model.load_state_dict(checkpoint)
-            else:
-                model.load_state_dict(checkpoint)
-            
-            # 移动到设备
-            print(f"移动模型到设备: {self.device}")
-            model = model.to(self.device)
-            
-            print("[OK] 模型加载成功！")
-            return model
-            
-        except Exception as e:
-            print(f"[ERROR] 加载模型时出错: {e}")
-            print(f"错误类型: {type(e).__name__}")
-            import traceback
-            traceback.print_exc()
-            raise
-    
-    def predict_segment(self, text: str) -> Dict[str, Any]:
-        """
-        预测单个段落的情感
-        
-        Args:
-            text: 段落文本
-            
-        Returns:
-            预测结果: {emotion_id, emotion_name, probabilities, confidence}
-        """
-        # 文本预处理
-        encoding = self.tokenizer(
-            text,
-            padding='max_length',
-            truncation=True,
-            max_length=self.max_length,
-            return_tensors='pt'
-        )
-        
-        input_ids = encoding['input_ids'].to(self.device)
-        attention_mask = encoding['attention_mask'].to(self.device)
-        
-        # 推理
-        with torch.no_grad():
-            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
-            logits = outputs['logits']
-            
-            # 对于分类任务，使用 softmax
-            probabilities = torch.softmax(logits, dim=-1)
-            
-            # 获取最高概率的情感
-            predicted_id = torch.argmax(probabilities, dim=-1).item()
-            confidence = probabilities[0, predicted_id].item()
-        
-        return {
-            'emotion_id': predicted_id,
-            'emotion_name': self.EMOTION_LABELS[predicted_id],
-            'probabilities': {
-                self.EMOTION_LABELS[i]: float(probabilities[0, i])
-                for i in range(len(self.EMOTION_LABELS))
-            },
-            'confidence': confidence
-        }
-    
-    def predict_segments_batch(self, texts: List[str]) -> List[Dict[str, Any]]:
-        """
-        批量预测多个段落的情感
-        
-        Args:
-            texts: 段落文本列表
-            
-        Returns:
-            预测结果列表
-        """
-        results = []
-        
-        # 分批处理
-        for i in range(0, len(texts), self.batch_size):
-            batch_texts = texts[i:i + self.batch_size]
-            
-            # 文本预处理
-            encoding = self.tokenizer(
-                batch_texts,
-                padding='max_length',
-                truncation=True,
-                max_length=self.max_length,
-                return_tensors='pt'
-            )
-            
-            input_ids = encoding['input_ids'].to(self.device)
-            attention_mask = encoding['attention_mask'].to(self.device)
-            
-            # 推理
-            with torch.no_grad():
-                logits = self.model(input_ids=input_ids, attention_mask=attention_mask)
-                
-                # 对于分类任务，使用 softmax
-                probabilities = torch.softmax(logits, dim=-1)
-                
-                # 获取最高概率的情感
-                predicted_ids = torch.argmax(probabilities, dim=-1)
-            
-            # 处理批次结果
-            for j in range(len(batch_texts)):
-                predicted_id = predicted_ids[j].item()
-                confidence = probabilities[j, predicted_id].item()
-                
-                result = {
-                    'emotion_id': predicted_id,
-                    'emotion_name': self.EMOTION_LABELS[predicted_id],
-                    'probabilities': {
-                        self.EMOTION_LABELS[k]: float(probabilities[j, k])
-                        for k in range(len(self.EMOTION_LABELS))
-                    },
-                    'confidence': confidence
-                }
-                results.append(result)
-        
-        return results
-    
-    def find_unique_emotion_segment(
-        self,
-        segment_predictions: List[Dict[str, Any]]
-    ) -> Dict[str, Any]:
-        """
-        找出表达独特情感的段落
-        
-        在 n 个段落中，n-1 个段落表达相同情感，1 个段落表达独特情感
-        
-        Args:
-            segment_predictions: 每个段落的预测结果
-            
-        Returns:
-            独特情感段落信息
-        """
-        # 统计每种情感出现的次数
-        emotion_counts = Counter([pred['emotion_name'] for pred in segment_predictions])
-        
-        # 找出只出现1次的情感 (独特情感)
-        unique_emotions = [emotion for emotion, count in emotion_counts.items() if count == 1]
-        
-        if len(unique_emotions) == 1:
-            # 找到独特情感
-            unique_emotion = unique_emotions[0]
-            
-            # 找到该情感对应的段落索引
-            for idx, pred in enumerate(segment_predictions):
-                if pred['emotion_name'] == unique_emotion:
-                    return {
-                        'unique_segment_index': idx,
-                        'unique_emotion': unique_emotion,
-                        'confidence': pred['confidence'],
-                        'emotion_distribution': dict(emotion_counts),
-                        'total_segments': len(segment_predictions),
-                        'status': 'success'
-                    }
-        
-        # 如果没有找到唯一的独特情感，使用启发式方法
-        # 方法1: 找出现次数最少且置信度最高的情感
-        min_count = min(emotion_counts.values())
-        rare_emotions = [emotion for emotion, count in emotion_counts.items() if count == min_count]
-        
-        # 在这些稀有情感中，找置信度最高的
-        best_idx = None
-        best_confidence = 0
-        best_emotion = None
-        
-        for idx, pred in enumerate(segment_predictions):
-            if pred['emotion_name'] in rare_emotions and pred['confidence'] > best_confidence:
-                best_idx = idx
-                best_confidence = pred['confidence']
-                best_emotion = pred['emotion_name']
-        
-        return {
-            'unique_segment_index': best_idx,
-            'unique_emotion': best_emotion,
-            'confidence': best_confidence,
-            'emotion_distribution': dict(emotion_counts),
-            'total_segments': len(segment_predictions),
-            'status': 'heuristic',
-            'note': f'No single unique emotion found. Used heuristic: rarest emotion ({min_count} occurrences) with highest confidence.'
-        }
-    
-    def inference_longemotion_test(
-        self,
-        test_file: str,
-        output_file: str,
-        output_detailed: Optional[str] = None
-    ):
-        """
-        对 LongEmotion 格式的测试集进行推理
-        
-        Args:
-            test_file: 测试集文件路径 (JSONL格式)
-            output_file: 输出文件路径 (提交格式)
-            output_detailed: 详细结果输出路径 (可选)
-        """
-        print(f"读取测试集: {test_file}")
-        
-        # 读取测试集
-        test_samples = []
-        with open(test_file, 'r', encoding='utf-8') as f:
-            for line in f:
-                if line.strip():
-                    test_samples.append(json.loads(line))
-        
-        print(f"测试样本数: {len(test_samples)}")
-        
-        # 推理结果
-        results = []
-        detailed_results = []
-        
-        # 处理每个样本
-        for sample_idx, sample in enumerate(tqdm(test_samples, desc="推理进度")):
-            # 提取段落
-            segments = sample['text']
-            
-            # 提取每个段落的文本
-            segment_texts = [seg['context'] for seg in segments]
-            
-            # 批量预测所有段落
-            segment_predictions = self.predict_segments_batch(segment_texts)
-            
-            # 找出独特情感段落
-            unique_result = self.find_unique_emotion_segment(segment_predictions)
-            
-            # 构建输出结果 (提交格式)
-            result = {
-                'sample_id': sample_idx,
-                'unique_segment_index': unique_result['unique_segment_index'],
-                'unique_emotion': unique_result['unique_emotion'],
-                'confidence': unique_result['confidence']
-            }
-            results.append(result)
-            
-            # 构建详细结果
-            if output_detailed:
-                detailed_result = {
-                    'sample_id': sample_idx,
-                    'total_length': sample.get('length', 0),
-                    'total_segments': len(segments),
-                    'unique_segment_index': unique_result['unique_segment_index'],
-                    'unique_emotion': unique_result['unique_emotion'],
-                    'confidence': unique_result['confidence'],
-                    'emotion_distribution': unique_result['emotion_distribution'],
-                    'status': unique_result['status'],
-                    'segment_predictions': [
-                        {
-                            'index': seg['index'],
-                            'text_preview': seg['context'][:100] + '...' if len(seg['context']) > 100 else seg['context'],
-                            'predicted_emotion': pred['emotion_name'],
-                            'confidence': pred['confidence']
-                        }
-                        for seg, pred in zip(segments, segment_predictions)
-                    ]
-                }
-                
-                if 'note' in unique_result:
-                    detailed_result['note'] = unique_result['note']
-                
-                detailed_results.append(detailed_result)
-        
-        # 保存结果
-        print(f"\n保存结果到: {output_file}")
-        os.makedirs(os.path.dirname(output_file), exist_ok=True)
-        
-        with open(output_file, 'w', encoding='utf-8') as f:
-            for result in results:
-                json.dump(result, f, ensure_ascii=False)
-                f.write('\n')
-        
-        # 保存详细结果
-        if output_detailed:
-            print(f"保存详细结果到: {output_detailed}")
-            with open(output_detailed, 'w', encoding='utf-8') as f:
-                json.dump(detailed_results, f, ensure_ascii=False, indent=2)
-        
-        # 统计信息
-        print("\n=== 推理统计 ===")
-        print(f"总样本数: {len(results)}")
-        
-        emotion_counts = Counter([r['unique_emotion'] for r in results])
-        print(f"\n独特情感分布:")
-        for emotion, count in emotion_counts.most_common():
-            print(f"  {emotion}: {count} ({count/len(results)*100:.1f}%)")
-        
-        avg_confidence = sum(r['confidence'] for r in results) / len(results)
-        print(f"\n���均置信度: {avg_confidence:.4f}")
-        
-        # 成功率统计
-        if detailed_results:
-            success_count = sum(1 for r in detailed_results if r['status'] == 'success')
-            print(f"找到唯一独特情感的样本: {success_count}/{len(results)} ({success_count/len(results)*100:.1f}%)")
-        
-        print("\n推理完成！")
-
-
-def main():
-    """主函数"""
-    import argparse
-    
-    parser = argparse.ArgumentParser(description="LongEmotion 测试集推理")
-    parser.add_argument(
-        "--model_path",
-        type=str,
-        default="../model/best_model.pt",
-        help="模型权重文件路径"
-    )
-    parser.add_argument(
-        "--test_file",
-        type=str,
-        default="../test_data/test.jsonl",
-        help="测试集文件路径"
-    )
-    parser.add_argument(
-        "--output_file",
-        type=str,
-        default="../submission/predictions.jsonl",
-        help="预测结果输出路径"
-    )
-    parser.add_argument(
-        "--output_detailed",
-        type=str,
-        default="../submission/predictions_detailed.json",
-        help="详细结果输出路径"
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cuda" if torch.cuda.is_available() else "cpu",
-        help="设备 (cuda/cpu)"
-    )
-    parser.add_argument(
-        "--max_length",
-        type=int,
-        default=512,
-        help="最大序列长度"
-    )
-    parser.add_argument(
-        "--batch_size",
-        type=int,
-        default=16,
-        help="批次大小"
-    )
-    
-    args = parser.parse_args()
-    
-    # 创建推理器
-    inference = LongEmotionInference(
-        model_path=args.model_path,
-        device=args.device,
-        max_length=args.max_length,
-        batch_size=args.batch_size
-    )
-    
-    # 执行推理
-    inference.inference_longemotion_test(
-        test_file=args.test_file,
-        output_file=args.output_file,
-        output_detailed=args.output_detailed
-    )
-
-
-if __name__ == "__main__":
-    main()
-