Added a base model class and training scripts for various sentiment analysis models, including Naive Bayes, SVM, XGBoost, LSTM, and BERT. Also, improved prediction functionality and the model loading mechanism.

2025-08-04 22:07:30 +08:00
parent bd60e2ed1b
commit 43525c5ca8
23 changed files with 1940 additions and 2362 deletions
@@ -0,0 +1,413 @@
+# -*- coding: utf-8 -*-
+"""
+BERT情感分析模型训练脚本
+"""
+import argparse
+import os
+import pandas as pd
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import BertTokenizer, BertModel
+from sklearn.metrics import accuracy_score, f1_score, classification_report, roc_auc_score
+from typing import List, Tuple
+import warnings
+import requests
+from pathlib import Path
+
+from base_model import BaseModel
+from utils import load_corpus_bert
+
+# 忽略transformers的警告
+warnings.filterwarnings("ignore")
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+
+class BertDataset(Dataset):
+    """BERT数据集"""
+    
+    def __init__(self, data: List[Tuple[str, int]]):
+        self.data = [item[0] for item in data]
+        self.labels = [item[1] for item in data]
+    
+    def __getitem__(self, index):
+        return self.data[index], self.labels[index]
+    
+    def __len__(self):
+        return len(self.labels)
+
+
+class BertClassifier(nn.Module):
+    """BERT分类器网络"""
+    
+    def __init__(self, input_size):
+        super(BertClassifier, self).__init__()
+        self.fc = nn.Linear(input_size, 1)
+        self.sigmoid = nn.Sigmoid()
+    
+    def forward(self, x):
+        out = self.fc(x)
+        out = self.sigmoid(out)
+        return out
+
+
+class BertModel_Custom(BaseModel):
+    """BERT情感分析模型"""
+    
+    def __init__(self, model_path: str = "./model/chinese_wwm_pytorch"):
+        super().__init__("BERT")
+        self.model_path = model_path
+        self.tokenizer = None
+        self.bert = None
+        self.classifier = None
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+    def _download_bert_model(self):
+        """自动下载BERT预训练模型"""
+        print(f"BERT模型不存在，正在下载中文BERT预训练模型...")
+        print("下载来源: bert-base-chinese (Hugging Face)")
+        
+        try:
+            # 创建模型目录
+            os.makedirs(self.model_path, exist_ok=True)
+            
+            # 使用Hugging Face的中文BERT模型
+            model_name = "bert-base-chinese"
+            print(f"正在从Hugging Face下载 {model_name}...")
+            
+            # 下载tokenizer
+            print("下载分词器...")
+            tokenizer = BertTokenizer.from_pretrained(model_name)
+            tokenizer.save_pretrained(self.model_path)
+            
+            # 下载模型
+            print("下载BERT模型...")
+            bert_model = BertModel.from_pretrained(model_name)
+            bert_model.save_pretrained(self.model_path)
+            
+            print(f"✅ BERT模型下载完成，保存在: {self.model_path}")
+            return True
+            
+        except Exception as e:
+            print(f"❌ BERT模型下载失败: {e}")
+            print("\n💡 您可以手动下载BERT模型:")
+            print("1. 访问 https://huggingface.co/bert-base-chinese")
+            print("2. 或使用哈工大中文BERT: https://github.com/ymcui/Chinese-BERT-wwm")
+            print(f"3. 将模型文件解压到: {self.model_path}")
+            return False
+    
+    def _load_bert(self):
+        """加载BERT模型和分词器"""
+        print(f"加载BERT模型: {self.model_path}")
+        
+        # 如果模型不存在，尝试自动下载
+        if not os.path.exists(self.model_path) or not any(os.scandir(self.model_path)):
+            print("BERT模型不存在，尝试自动下载...")
+            if not self._download_bert_model():
+                raise FileNotFoundError(f"BERT模型下载失败，请手动下载到: {self.model_path}")
+        
+        try:
+            self.tokenizer = BertTokenizer.from_pretrained(self.model_path)
+            self.bert = BertModel.from_pretrained(self.model_path).to(self.device)
+            
+            # 冻结BERT参数
+            for param in self.bert.parameters():
+                param.requires_grad = False
+                
+            print("✅ BERT模型加载完成")
+            
+        except Exception as e:
+            print(f"❌ BERT模型加载失败: {e}")
+            print("尝试使用在线模型...")
+            
+            # 如果本地加载失败，尝试直接使用在线模型
+            try:
+                model_name = "bert-base-chinese"
+                self.tokenizer = BertTokenizer.from_pretrained(model_name)
+                self.bert = BertModel.from_pretrained(model_name).to(self.device)
+                
+                # 冻结BERT参数
+                for param in self.bert.parameters():
+                    param.requires_grad = False
+                    
+                print("✅ 在线BERT模型加载完成")
+                
+            except Exception as e2:
+                print(f"❌ 在线模型也加载失败: {e2}")
+                raise FileNotFoundError(f"无法加载BERT模型，请检查网络连接或手动下载模型到: {self.model_path}")
+    
+    def train(self, train_data: List[Tuple[str, int]], **kwargs) -> None:
+        """训练BERT模型"""
+        print(f"开始训练 {self.model_name} 模型...")
+        
+        # 加载BERT
+        self._load_bert()
+        
+        # 超参数
+        learning_rate = kwargs.get('learning_rate', 1e-3)
+        num_epochs = kwargs.get('num_epochs', 10)
+        batch_size = kwargs.get('batch_size', 100)
+        input_size = kwargs.get('input_size', 768)
+        decay_rate = kwargs.get('decay_rate', 0.9)
+        
+        print(f"BERT超参数: lr={learning_rate}, epochs={num_epochs}, "
+              f"batch_size={batch_size}, input_size={input_size}")
+        
+        # 创建数据集
+        train_dataset = BertDataset(train_data)
+        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+        
+        # 创建分类器
+        self.classifier = BertClassifier(input_size).to(self.device)
+        
+        # 损失函数和优化器
+        criterion = nn.BCELoss()
+        optimizer = torch.optim.Adam(self.classifier.parameters(), lr=learning_rate)
+        scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=decay_rate)
+        
+        # 训练循环
+        self.bert.eval()  # BERT始终保持评估模式
+        self.classifier.train()
+        
+        for epoch in range(num_epochs):
+            total_loss = 0
+            num_batches = 0
+            
+            for i, (words, labels) in enumerate(train_loader):
+                # 分词和编码
+                tokens = self.tokenizer(words, padding=True, truncation=True, 
+                                      max_length=512, return_tensors='pt')
+                input_ids = tokens["input_ids"].to(self.device)
+                attention_mask = tokens["attention_mask"].to(self.device)
+                labels = torch.tensor(labels, dtype=torch.float32).to(self.device)
+                
+                # 获取BERT输出（冻结参数）
+                with torch.no_grad():
+                    bert_outputs = self.bert(input_ids, attention_mask=attention_mask)
+                    bert_output = bert_outputs[0][:, 0]  # [CLS] token的输出
+                
+                # 分类器前向传播
+                optimizer.zero_grad()
+                outputs = self.classifier(bert_output)
+                logits = outputs.view(-1)
+                loss = criterion(logits, labels)
+                
+                # 反向传播
+                loss.backward()
+                optimizer.step()
+                
+                total_loss += loss.item()
+                num_batches += 1
+                
+                if (i + 1) % 10 == 0:
+                    avg_loss = total_loss / num_batches
+                    print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}], Loss: {avg_loss:.4f}")
+                    total_loss = 0
+                    num_batches = 0
+            
+            # 学习率衰减
+            scheduler.step()
+            
+            # 保存每个epoch的模型
+            if kwargs.get('save_each_epoch', False):
+                epoch_model_path = f"./model/bert_epoch_{epoch+1}.pth"
+                os.makedirs(os.path.dirname(epoch_model_path), exist_ok=True)
+                torch.save(self.classifier.state_dict(), epoch_model_path)
+                print(f"已保存模型: {epoch_model_path}")
+        
+        self.is_trained = True
+        print(f"{self.model_name} 模型训练完成！")
+    
+    def predict(self, texts: List[str]) -> List[int]:
+        """预测文本情感"""
+        if not self.is_trained:
+            raise ValueError(f"模型 {self.model_name} 尚未训练，请先调用train方法")
+        
+        predictions = []
+        batch_size = 32
+        
+        self.bert.eval()
+        self.classifier.eval()
+        
+        with torch.no_grad():
+            for i in range(0, len(texts), batch_size):
+                batch_texts = texts[i:i+batch_size]
+                
+                # 分词和编码
+                tokens = self.tokenizer(batch_texts, padding=True, truncation=True,
+                                      max_length=512, return_tensors='pt')
+                input_ids = tokens["input_ids"].to(self.device)
+                attention_mask = tokens["attention_mask"].to(self.device)
+                
+                # 获取BERT输出
+                bert_outputs = self.bert(input_ids, attention_mask=attention_mask)
+                bert_output = bert_outputs[0][:, 0]
+                
+                # 分类器预测
+                outputs = self.classifier(bert_output)
+                outputs = outputs.view(-1)
+                
+                # 转换为类别标签
+                preds = (outputs > 0.5).cpu().numpy()
+                predictions.extend(preds.astype(int).tolist())
+        
+        return predictions
+    
+    def predict_single(self, text: str) -> Tuple[int, float]:
+        """预测单条文本的情感"""
+        if not self.is_trained:
+            raise ValueError(f"模型 {self.model_name} 尚未训练，请先调用train方法")
+        
+        self.bert.eval()
+        self.classifier.eval()
+        
+        with torch.no_grad():
+            # 分词和编码
+            tokens = self.tokenizer([text], padding=True, truncation=True,
+                                  max_length=512, return_tensors='pt')
+            input_ids = tokens["input_ids"].to(self.device)
+            attention_mask = tokens["attention_mask"].to(self.device)
+            
+            # 获取BERT输出
+            bert_outputs = self.bert(input_ids, attention_mask=attention_mask)
+            bert_output = bert_outputs[0][:, 0]
+            
+            # 分类器预测
+            output = self.classifier(bert_output)
+            prob = output.item()
+            
+            prediction = int(prob > 0.5)
+            confidence = prob if prediction == 1 else 1 - prob
+        
+        return prediction, confidence
+    
+    def save_model(self, model_path: str = None) -> None:
+        """保存模型"""
+        if not self.is_trained:
+            raise ValueError(f"模型 {self.model_name} 尚未训练，无法保存")
+        
+        if model_path is None:
+            model_path = f"./model/{self.model_name.lower()}_model.pth"
+        
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+        
+        # 保存分类器和相关信息
+        model_data = {
+            'classifier_state_dict': self.classifier.state_dict(),
+            'model_path': self.model_path,
+            'input_size': 768,
+            'device': str(self.device)
+        }
+        
+        torch.save(model_data, model_path)
+        print(f"模型已保存到: {model_path}")
+    
+    def load_model(self, model_path: str) -> None:
+        """加载模型"""
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"模型文件不存在: {model_path}")
+        
+        model_data = torch.load(model_path, map_location=self.device)
+        
+        # 设置BERT模型路径
+        self.model_path = model_data['model_path']
+        
+        # 加载BERT
+        self._load_bert()
+        
+        # 重建分类器
+        input_size = model_data['input_size']
+        self.classifier = BertClassifier(input_size).to(self.device)
+        
+        # 加载分类器权重
+        self.classifier.load_state_dict(model_data['classifier_state_dict'])
+        
+        self.is_trained = True
+        print(f"已加载模型: {model_path}")
+    
+    @staticmethod
+    def load_data(train_path: str, test_path: str) -> Tuple[List[Tuple[str, int]], List[Tuple[str, int]]]:
+        """加载BERT格式的数据"""
+        print("加载训练数据...")
+        train_data = load_corpus_bert(train_path)
+        print(f"训练数据量: {len(train_data)}")
+        
+        print("加载测试数据...")
+        test_data = load_corpus_bert(test_path)
+        print(f"测试数据量: {len(test_data)}")
+        
+        return train_data, test_data
+
+
+def main():
+    """主函数"""
+    parser = argparse.ArgumentParser(description='BERT情感分析模型训练')
+    parser.add_argument('--train_path', type=str, default='./data/weibo2018/train.txt',
+                        help='训练数据路径')
+    parser.add_argument('--test_path', type=str, default='./data/weibo2018/test.txt',
+                        help='测试数据路径')
+    parser.add_argument('--model_path', type=str, default='./model/bert_model.pth',
+                        help='模型保存路径')
+    parser.add_argument('--bert_path', type=str, default='./model/chinese_wwm_pytorch',
+                        help='BERT预训练模型路径')
+    parser.add_argument('--epochs', type=int, default=10,
+                        help='训练轮数')
+    parser.add_argument('--batch_size', type=int, default=100,
+                        help='批大小')
+    parser.add_argument('--learning_rate', type=float, default=1e-3,
+                        help='学习率')
+    parser.add_argument('--eval_only', action='store_true',
+                        help='仅评估已有模型，不进行训练')
+    
+    args = parser.parse_args()
+    
+    # 创建模型
+    model = BertModel_Custom(args.bert_path)
+    
+    if args.eval_only:
+        # 仅评估模式
+        print("评估模式：加载已有模型进行评估")
+        model.load_model(args.model_path)
+        
+        # 加载测试数据
+        _, test_data = model.load_data(args.train_path, args.test_path)
+        
+        # 评估模型
+        model.evaluate(test_data)
+    else:
+        # 训练模式
+        # 加载数据
+        train_data, test_data = model.load_data(args.train_path, args.test_path)
+        
+        # 训练模型
+        model.train(
+            train_data,
+            num_epochs=args.epochs,
+            batch_size=args.batch_size,
+            learning_rate=args.learning_rate
+        )
+        
+        # 评估模型
+        model.evaluate(test_data)
+        
+        # 保存模型
+        model.save_model(args.model_path)
+        
+        # 示例预测
+        print("\n示例预测:")
+        test_texts = [
+            "今天天气真好，心情很棒",
+            "这部电影太无聊了，浪费时间",
+            "哈哈哈，太有趣了"
+        ]
+        
+        for text in test_texts:
+            pred, conf = model.predict_single(text)
+            sentiment = "正面" if pred == 1 else "负面"
+            print(f"文本: {text}")
+            print(f"预测: {sentiment} (置信度: {conf:.4f})")
+            print()
+
+
+if __name__ == "__main__":
+    main()