基于自有数据集微调BERT下一句预测模型的数据集与Hugging Face Trainer实现咨询

适配NSP任务的数据集结构要求 #

• 基础语料拆分规则：先将所有原始文档拆分为语义完整的独立句子，保留句子在原文档中的先后顺序，不得拆分出残缺的半句话作为训练单元。
• 训练样本构造逻辑：每条训练样本必须由成对文本+二分类标签组成：
  • 正样本（标签=1）：取同一文档内语义连贯的相邻两个句子
  • 负样本（标签=0）：从无关联的其他语料中随机抽取句子，与当前句子配对，不得使用原文档中相邻的下一句，避免语义连贯性干扰标签正确性
• 内置数据集类的格式要求：如果使用Hugging Face提供的TextDatasetForNextSentencePrediction工具类，无需手动构造标签和文本对，只需要将按序排列的完整句子以每行一句的格式存入纯文本文件即可，类内部会自动完成相邻句正样本、随机抽句负样本的构造与标签标注。

原代码的核心错误点 #

• 数据集传参错误：Trainer的train_dataset参数需要传入PyTorch Dataset类实例，不能直接传入本地文件路径字符串
• 数据收集器选型错误：DataCollatorForLanguageModeling是掩码语言模型（MLM）任务的专用收集器，会对输入token做动态掩码，和NSP任务的输入预期完全不匹配
• 分词器传参错误：Trainer初始化时的tokenizer参数需要传入已加载完成的分词器实例，不能直接传入BertTokenizer类本身

基于Trainer接口的正确微调实现 #

修正后的完整代码如下：

import torch
from transformers import (
    BertTokenizer,
    BertForNextSentencePrediction,
    Trainer,
    TrainingArguments,
    TextDatasetForNextSentencePrediction,
    DataCollatorWithPadding
)

def train(bert_model, bert_tokenizer, data_path, eval_path=None):
    out_dir = "/content/drive/My Drive/next_sentence/"
    training_args = TrainingArguments(
        output_dir=out_dir,
        overwrite_output_dir=True,
        num_train_epochs=1,
        per_device_train_batch_size=30,
        save_steps=100,
        save_total_limit=5,
        prediction_loss_only=True
    )
    # 加载NSP专用数据集，自动完成文本配对与标签构造
    train_dataset = TextDatasetForNextSentencePrediction(
        tokenizer=bert_tokenizer,
        file_path=data_path,
        block_size=128
    )
    # 替换错用的MLM收集器，使用通用Padding收集器适配NSP输入
    data_collator = DataCollatorWithPadding(tokenizer=bert_tokenizer)
    trainer = Trainer(
        model=bert_model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        eval_dataset=None if eval_path is None else TextDatasetForNextSentencePrediction(
            tokenizer=bert_tokenizer,
            file_path=eval_path,
            block_size=128
        ),
        tokenizer=bert_tokenizer
    )
    trainer.train()
    trainer.save_model(out_dir)

if __name__ == "__main__":
    bert_tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
    bert_model = BertForNextSentencePrediction.from_pretrained("bert-base-cased")
    train_data_set_path = "c:/data.txt"
    train(bert_model, bert_tokenizer, train_data_set_path)

额外注意事项 #

• 手动构造自定义Dataset时，需要保证每条样本的编码结果包含input_ids、attention_mask、token_type_ids、labels（NSP二分类标签0/1）四个字段，不需要对输入做动态掩码处理
• 构造负样本时，尽量从跨文档的语料中随机抽取，避免引入同文档内的语义关联，导致模型学习偏差

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