如何在HuggingFace中创建含多选与序列分类双任务头的语言模型？

要实现共享基础模型、同时带序列分类和多选任务头的模型，核心思路是复用预训练的DistilBERT主体权重，仅新增两个任务头的可训练参数，这样就能保证和单头模型的差异只在新增的任务头上。具体实现步骤如下：

1. 核心实现逻辑 #

自定义模型类整合DistilBERT基础编码器、序列分类头、多选头，前向传播时共享编码器输出，分别传入两个任务头得到结果。

2. 完整代码示例 #

from transformers import DistilBertModel, DistilBertPreTrainedModel, AutoTokenizer
import torch
import torch.nn as nn

# 自定义双头模型类，继承自DistilBertPreTrainedModel以复用预训练权重加载逻辑
class DistilBertDualHeadModel(DistilBertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        # 加载预训练的DistilBERT基础编码器（无任务头）
        self.distilbert = DistilBertModel(config)
        
        # 序列分类头：对应AutoModelForSequenceClassification的结构
        self.classifier = nn.Sequential(
            nn.Dropout(config.seq_classif_dropout),
            nn.Linear(config.dim, config.num_labels)
        )
        
        # 多选任务头：对应AutoModelForMultipleChoice的结构
        self.multiple_choice_classifier = nn.Sequential(
            nn.Dropout(config.seq_classif_dropout),
            nn.Linear(config.dim, 1)
        )
        
        # 初始化新增的任务头参数
        self.init_weights()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,  # 格式为[sequence_label, multiple_choice_labels]
    ):
        # 共享基础编码器的输出
        outputs = self.distilbert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )
        # 获取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token的输出（序列分类用）
        pooled_output = outputs[0][:, 0]
        # 获取所有token的输出（多选任务用，会在后续做池化）
        sequence_output = outputs[0]

        # 序列分类任务的logits
        cls_logits = self.classifier(pooled_output)
        # 多选任务的logits：对每个选项的序列做池化后传入分类头
        # 假设输入的input_ids是[batch_size, num_choices, seq_length]格式
        mc_logits = self.multiple_choice_classifier(sequence_output).squeeze(-1)
        mc_logits = mc_logits.mean(dim=-1)  # 对每个选项的序列做平均池化，可根据需求调整

        # 准备输出结果
        loss = None
        if labels is not None:
            seq_label, mc_labels = labels
            # 计算序列分类损失
            cls_loss_fct = nn.CrossEntropyLoss()
            cls_loss = cls_loss_fct(cls_logits, seq_label)
            # 计算多选任务损失
            mc_loss_fct = nn.CrossEntropyLoss()
            mc_loss = mc_loss_fct(mc_logits, mc_labels)
            # 总损失：可根据任务权重调整比例
            loss = cls_loss + mc_loss

        return (loss, cls_logits, mc_logits) if loss is not None else (cls_logits, mc_logits)

# 加载预训练权重并初始化模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = DistilBertDualHeadModel.from_pretrained("distilbert-base-cased", num_labels=2).to(device)
# 示例：加载tokenizer
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased")

3. 关键说明 #

• 参数复用：基础的DistilBERT编码器直接加载预训练权重，仅新增的classifier和multiple_choice_classifier两个任务头会随机初始化，保证和单头模型的差异仅在新增参数上。
• 输入格式适配：
  • 文本输入需处理为[batch_size, num_choices, seq_length]格式，使用tokenizer时指定padding=True和return_tensors="pt"即可实现。
  • 标签输入需为[sequence_label, multiple_choice_labels]，其中sequence_label是形状为[batch_size]的张量，multiple_choice_labels是形状为[batch_size]的张量（每个样本对应正确选项的索引）。
• 损失计算：可根据两个任务的重要性调整损失权重，比如设置loss = 0.3*cls_loss + 0.7*mc_loss。

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