BGE-M3 文本情感分类实战：预训练模型微调，导出ONNX并测试

基于BGE-M3开源模型（由智谱 AI 开发的中英双语通用向量表征模型），实现一个简单的二分类情感分析任务（积极 / 消极文本判断）。通过微调预训练模型，展示 NLP 中经典的 “预训练 + 微调” 范式。微调采用基于 [CLS] token 的单特征提取或者基于池化策略的多特征融合。并导出ONNX模型，进行测试，以适应生产环境。

微调步骤

1. 加载预训练模型与分词器；
2. 构建分类头并整合到模型中；
3. 数据预处理与 DataLoader 构建；
4. 定义优化器、损失函数与训练循环；
5. 验证与推理。
   实际应用中，需根据数据规模和任务复杂度调整模型结构与训练策略。BGE-M3 在语义表征任务中表现优异，适合作为文本分类、语义搜索等下游任务的基础模型。

微调代码

from transformers import AutoModel, AutoTokenizer,  get_linear_schedule_with_warmup
import torch
from torch.optim import AdamW  
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from sklearn.model_selection import train_test_split
from tqdm import tqdm

# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-m3")
model = AutoModel.from_pretrained("BAAI/bge-m3")

# 定义分类模型
class TextClassifier(nn.Module):
    def __init__(self, model, hidden_size=1024, num_classes=2):
        super().__init__()
        self.model = model
        self.classifier = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, num_classes)
        )
    
    def forward(self, inputs):
        # 获取模型输出
        outputs = self.model(**inputs)
        # 使用 [CLS] token
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        # 通过分类器
        logits = self.classifier(cls_embedding)
        return logits

# 模拟数据集
class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
        
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        print(torch.tensor(label, dtype=torch.long))
        return {
   
   
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'label': torch.tensor(label, dtype=torch.long)
        }

# 初始化分类器
classifier = TextClassifier(model)

# 模拟数据
texts = [
    "这是一个积极的句子，充满了正能量。",
    "这是一个消极的句子，感觉非常糟糕。",
    "今天天气真好，阳光明媚。",
    "这个电影太无聊了，浪费时间。",
    "我喜欢这个产品，质量非常好。",
    "这个服务太差劲了，非常不满意。",
    "大模型对程序员来说是一个很好的工具。",
    "大模型对初级开发者来说是一个坏消息。"
]
#  PyTorch 的交叉熵损失函数 nn.CrossEntropyLoss 要求标签必须是从 0 开始的连续整数（如 0、1、2...）
labels = [1, 0, 1, 0, 1, 0,1,0]  # 1表示积极，0表示消极

# 划分训练集和验证集
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)

# 创建数据加载器
train_dataset = TextDataset(train_texts, train_labels, tokenizer)
val_dataset = TextDataset(val_texts, val_labels, tokenizer)

train_dataloader = DataLoader(train_dataset, batch_size=2, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=2)

# 训练配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
classifier.to(device)

optimizer = AdamW(classifier.parameters(), lr=2e-5)
epochs = 3
total_steps = len(train_dataloader) * epochs
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_steps)

# 训练循环
def train_epoch(model, dataloader, optimizer,