基于HuggingFace Transformer的句子级序列到序列分类技术问询

核心思路 #

基于HuggingFace预训练Transformer模型改造，适配句子级seq2seq分类任务，摒弃滑动窗口拼接的冗余方式，直接利用Transformer自注意力机制捕捉全文上下文关联。

具体实现步骤 #

• 输入格式处理：将整段文档的句子序列按顺序拼接，给每个句子添加专属分隔符（如自定义<sent_sep>），让模型明确区分句子边界；无需额外拼接窗口，直接输入完整文档。
• 模型改造：
  • 选用bert-base-uncased/roberta-base等预训练模型，在顶部新增句子级分类头。
  • 提取每个句子起始位置token的隐藏状态作为该句子的语义表征，输入分类头完成标签预测。
• 训练策略：
  • 损失函数采用交叉熵损失，对每个句子的预测结果独立计算损失后求和。
  • 采用整文档批量训练，配合梯度累积平衡显存占用与训练效率。

核心代码片段 #

from transformers import BertTokenizer, BertModel
import torch
import torch.nn as nn

# 定义带句子分类头的Transformer模型
class SentenceLevelClassifier(nn.Module):
    def __init__(self, model_name, num_labels):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)
    
    def forward(self, input_ids, attention_mask, sent_start_positions):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        # 提取每个句子起始位置的隐藏状态
        sent_reps = outputs.last_hidden_state[torch.arange(input_ids.shape[0])[:, None], sent_start_positions]
        logits = self.classifier(sent_reps)
        return logits

# 初始化模型与tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = SentenceLevelClassifier('bert-base-uncased', num_labels=5)

# 构造输入示例
sentences = ["First sentence of the doc.", "Second sentence here.", "Third one follows."]
# 拼接带分隔符的文本
input_text = "<sent_sep>".join(sentences)
tokenized_input = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)

# 定位每个句子的起始token位置
sent_start_pos = []
current_idx = 1  # 跳过[CLS]
for sent in sentences:
    sent_start_pos.append(current_idx)
    current_idx += len(tokenizer.tokenize(sent)) + 1  # +1对应<sent_sep>的token数
sent_start_pos = torch.tensor([sent_start_pos])

# 前向预测
logits = model(**tokenized_input, sent_start_positions=sent_start_pos)
# 计算损失（labels为句子对应的标签张量）
labels = torch.tensor([[0, 1, 2]])
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits.view(-1, 5), labels.view(-1))

性能优化方案 #

• 取消滑动窗口：直接输入完整文档，利用Transformer自注意力捕捉长距离上下文，避免重复拼接计算的冗余开销。
• 动态Padding：按批次内最长文档的长度做padding，而非固定窗口长度，提升显存利用率。
• 量化训练：借助HuggingFace的bitsandbytes库实现4/8位量化，降低显存占用，加快训练与推理速度。

内容的提问来源于stack exchange，提问作者Keren L