Delete scripts/detection_model.py

scripts/detection_model.py

@@ -1,238 +0,0 @@
-"""
-情感检测模型
-Emotion Detection Task
-多标签分类任务
-"""
-import torch
-import torch.nn as nn
-from transformers import (
-    BertModel,
-    BertTokenizer,
-    AutoModel,
-    AutoTokenizer
-)
-from typing import Dict, List, Optional
-
-
-class EmotionDetectionModel(nn.Module):
-    """情感检测模型（多标签分类）"""
-    
-    def __init__(
-        self,
-        model_name: str = "bert-base-chinese",
-        num_emotions: int = 7,
-        dropout: float = 0.1
-    ):
-        """
-        初始化检测模型
-        
-        Args:
-            model_name: 预训练模型名称
-            num_emotions: 情感类别数量
-            dropout: Dropout 比率
-        """
-        super().__init__()
-        
-        self.bert = AutoModel.from_pretrained(model_name)
-        hidden_size = self.bert.config.hidden_size
-        
-        # 多标签分类头
-        self.classifier = nn.Sequential(
-            nn.Dropout(dropout),
-            nn.Linear(hidden_size, hidden_size // 2),
-            nn.ReLU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_size // 2, num_emotions)
-        )
-        
-        self.num_emotions = num_emotions
-    
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        labels: Optional[torch.Tensor] = None
-    ):
-        """
-        前向传播
-        
-        Args:
-            input_ids: 输入ID
-            attention_mask: 注意力掩码
-            labels: 标签（多标签，形状为 [batch_size, num_emotions]）
-            
-        Returns:
-            模型输出
-        """
-        outputs = self.bert(
-            input_ids=input_ids,
-            attention_mask=attention_mask
-        )
-        
-        # 使用 [CLS] token 的表示
-        pooled_output = outputs.last_hidden_state[:, 0, :]
-        
-        # 分类
-        logits = self.classifier(pooled_output)
-        
-        loss = None
-        if labels is not None:
-            # 使用 BCE Loss 进行多标签分类
-            loss_fct = nn.BCEWithLogitsLoss()
-            loss = loss_fct(logits, labels.float())
-        
-        return {
-            'loss': loss,
-            'logits': logits
-        }
-
-
-class EmotionDetectionModelWrapper:
-    """情感检测模型封装器"""
-    
-    def __init__(
-        self,
-        model_name: str = "bert-base-chinese",
-        num_emotions: int = 7,
-        max_length: int = 512,
-        device: str = "cuda" if torch.cuda.is_available() else "cpu",
-        threshold: float = 0.5
-    ):
-        """
-        初始化
-        
-        Args:
-            model_name: 预训练模型名称
-            num_emotions: 情感类别数量
-            max_length: 最大序列长度
-            device: 设备
-            threshold: 分类阈值
-        """
-        self.model_name = model_name
-        self.num_emotions = num_emotions
-        self.max_length = max_length
-        self.device = device
-        self.threshold = threshold
-        
-        # 加载分词器
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        
-        # 加载模型
-        self.model = EmotionDetectionModel(
-            model_name=model_name,
-            num_emotions=num_emotions
-        ).to(device)
-        
-        # 情感标签名称
-        self.emotion_names = [
-            "happiness",
-            "sadness",
-            "anger",
-            "fear",
-            "surprise",
-            "disgust",
-            "neutral"
-        ]
-    
-    def preprocess(self, texts: List[str]) -> Dict[str, torch.Tensor]:
-        """预处理文本"""
-        encoding = self.tokenizer(
-            texts,
-            padding=True,
-            truncation=True,
-            max_length=self.max_length,
-            return_tensors="pt"
-        )
-        
-        return {k: v.to(self.device) for k, v in encoding.items()}
-    
-    def predict(self, texts: List[str]) -> List[Dict[str, any]]:
-        """
-        批量预测
-        
-        Args:
-            texts: 文本列表
-            
-        Returns:
-            预测结果列表
-        """
-        self.model.eval()
-        
-        inputs = self.preprocess(texts)
-        
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-            logits = outputs['logits']
-            probabilities = torch.sigmoid(logits)  # 多标签使用 sigmoid
-        
-        results = []
-        for i, probs in enumerate(probabilities):
-            # 获取超过阈值的情感
-            detected_emotions = []
-            emotion_scores = {}
-            
-            for j, prob in enumerate(probs):
-                emotion_scores[self.emotion_names[j]] = float(prob)
-                if prob >= self.threshold:
-                    detected_emotions.append({
-                        'emotion': self.emotion_names[j],
-                        'score': float(prob)
-                    })
-            
-            results.append({
-                'text': texts[i],
-                'emotions': detected_emotions,
-                'all_scores': emotion_scores
-            })
-        
-        return results
-    
-    def save(self, save_path: str):
-        """保存模型"""
-        import os
-        os.makedirs(save_path, exist_ok=True)
-        
-        # 保存模型权重
-        torch.save(self.model.state_dict(), f"{save_path}/model.pt")
-        
-        # 保存分词器
-        self.tokenizer.save_pretrained(save_path)
-        
-        # 保存配置
-        import json
-        config = {
-            'model_name': self.model_name,
-            'num_emotions': self.num_emotions,
-            'max_length': self.max_length,
-            'threshold': self.threshold
-        }
-        with open(f"{save_path}/config.json", 'w') as f:
-            json.dump(config, f, indent=2)
-        
-        print(f"检测模型已保存到 {save_path}")
-    
-    @classmethod
-    def load(cls, model_path: str, device: str = "cuda"):
-        """加载模型"""
-        import json
-        
-        # 加载配置
-        with open(f"{model_path}/config.json", 'r') as f:
-            config = json.load(f)
-        
-        instance = cls(
-            model_name=config['model_name'],
-            num_emotions=config['num_emotions'],
-            max_length=config['max_length'],
-            device=device,
-            threshold=config.get('threshold', 0.5)
-        )
-        
-        # 加载权重
-        instance.model.load_state_dict(
-            torch.load(f"{model_path}/model.pt", map_location=device)
-        )
-        
-        print(f"检测模型已从 {model_path} 加载")
-        return instance
-