466_文本分类任务实战¶

# 使用 HuggingFace Transformers 做文本分类微调

from transformers import BertTokenizer, BertForSequenceClassification
from torch.utils.data import DataLoader, Dataset
import torch

# ---- 1. 自定义数据集 ----
class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        # 分词并转换为模型输入格式
        enc = self.tokenizer(
            self.texts[idx],
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': enc['input_ids'].squeeze(0),          # token id序列
            'attention_mask': enc['attention_mask'].squeeze(0), # 注意力掩码
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }

# ---- 2. 加载预训练BERT ----
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
# num_labels 设置为分类类别数
model = BertForSequenceClassification.from_pretrained(
    'bert-base-chinese', num_labels=2
)

# ---- 3. 简单训练循环 ----
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)  # BERT标准学习率

# 模拟数据
texts = ["今天天气真好", "这个产品太差了"]
labels = [1, 0]  # 1=正面, 0=负面

dataset = TextDataset(texts, labels, tokenizer)
loader = DataLoader(dataset, batch_size=2)

model.train()
for batch in loader:
    outputs = model(
        input_ids=batch['input_ids'],
        attention_mask=batch['attention_mask'],
        labels=batch['label']  # 传入labels自动计算交叉熵损失
    )
    loss = outputs.loss  # 模型内部已计算loss
    loss.backward()      # 反向传播
    optimizer.step()     # 更新参数
    optimizer.zero_grad()  # 清零梯度
    print(f"Loss: {loss.item():.4f}")

# ---- 4. 推理预测 ----
model.eval()
with torch.no_grad():
    enc = tokenizer("这个电影很精彩", return_tensors='pt',
                    max_length=128, truncation=True, padding='max_length')
    logits = model(**enc).logits          # 原始输出分数
    pred = logits.argmax(dim=-1).item()   # 取最大值对应类别
    print(f"预测类别: {pred}")