如何使用字符串格式自定义数据集微调掩码语言模型

算力阿峰

2026-8-30

问题

参考Hugging Face教程做掩码语言模型（随机掩码部分词元并完成预测的任务）微调时，教程默认使用load_dataset直接加载的平台内置数据集，但我手头只有长字符串格式的自定义数据，示例如下：

text = "This is an attempt of a great example. "
dataset = text * 3000

我按教程流程先写了分词代码：

from transformers import AutoTokenizer
from transformers import AutoModelForMaskedLM
import torch
from transformers import DataCollatorForLanguageModeling

model_checkpoint = "distilbert-base-uncased"
model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

def tokenize_long_text(tokenizer, long_text):
    individual_sentences = long_text.split('.')
    tokenized_sentences_list = tokenizer(individual_sentences)['input_ids']
    tokenized_sequence = [x for xs in tokenized_sentences_list for x in xs]
    return tokenized_sequence

    tokenized_sequence = tokenize_long_text(tokenizer, long_text)

之后写了固定长度切分文本块的代码：

def chunk_long_tokenized_text(tokenizer_text, chunk_size):
    # 计算长分词文本总长度
    total_length = len(tokenizer_text)
    # 丢弃长度不足chunk_size的最后一个文本块
    total_length = (total_length // chunk_size) * chunk_size

    return [tokenizer_text[i : i + chunk_size] for i in range(0, total_length, chunk_size)]

chunked_sequence = chunk_long_tokenized_text(tokenized_sequence, 30)

随后构建了掩码概率0.15的DataCollatorForLanguageModeling：

# 该组件接收字典列表，每个字典对应一段连续文本块
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)

测试数据整理器时得到如下输出，能看到连续重复的特殊标记：

d = {}
d['input_ids'] = chunked_sequence[0]
d

>>>{'input_ids': [101,
  2023,
  2003,
  1037,
  2307,
  103,...

for chunk in data_collator([ d ])["input_ids"]:
    print(f"'>>> {tokenizer.decode(chunk)}'")
    >>>'>>> [CLS] this is a great [MASK] [SEP] [CLS] this is a great [MASK] [SEP] [CLS] this is a great [MASK] [SEP] [CLS] this is a great [MASK] [SEP] [CLS] this'

现在的问题是Trainer仅支持平台规定格式的数据集，如何适配自定义字符串数据集，完成掩码语言模型全流程微调。

解决方案

核心问题有两个：一是你没有把处理好的分块数据封装为datasets库支持的Dataset格式，二是现有分词逻辑存在缩进bug和冗余设计，会生成大量重复的特殊标记影响训练效果。按以下步骤调整即可：

先修复基础代码错误：你写的tokenize_long_text函数中，调用函数生成tokenized_sequence的代码缩进在函数内部，永远不会执行，需要把这行挪到函数外。
弃用手动分句再拼接的逻辑：你按.拆分句子后逐句分词，会给每个句子都加上<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>和[SEP]标记，拼接后就会出现测试输出里连续重复特殊标记的问题，直接对整段文本分词后再切分块即可。
用Dataset.from_dict把内存里的字符串转为标准数据集格式，通过map方法批量完成分词、切分，输出结果可以直接传入Trainer。

完整可运行代码如下：

from datasets import Dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForMaskedLM, 
    DataCollatorForLanguageModeling, 
    TrainingArguments, 
    Trainer
)

# 初始化基础组件
model_checkpoint = "distilbert-base-uncased"
model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
chunk_size = 128  # 可按需调整，不要超过模型支持的最大上下文长度（distilbert为512）

# 把长字符串转为标准Dataset格式
raw_text = "This is an attempt of a great example. " * 3000
dataset = Dataset.from_dict({"text": [raw_text]})
# 如果是多个文本样本，把所有文本放到list里即可，例如{"text": [text1, text2, text3]}

# 批量分词+分块处理函数
def tokenize_and_chunk(examples):
    # 批量分词，关闭截断保留全部分词结果
    tokenized_res = tokenizer(examples["text"], truncation=False)
    chunks = []
    for input_ids in tokenized_res["input_ids"]:
        # 按固定步长切分，丢弃最后一个长度不足的块
        for i in range(0, len(input_ids), chunk_size):
            chunk = input_ids[i:i+chunk_size]
            if len(chunk) == chunk_size:
                chunks.append({"input_ids": chunk})
    # 整理为map方法要求的字典格式
    return {k: [dic[k] for dic in chunks] for k in chunks[0]}

# 处理数据集，开启批量处理提速，处理完删除原始文本列
processed_dataset = dataset.map(
    tokenize_and_chunk, 
    batched=True, 
    remove_columns=dataset.column_names
)
# 拆分训练集和验证集
processed_dataset = processed_dataset.train_test_split(test_size=0.1)

# 初始化数据整理器、训练参数
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)
training_args = TrainingArguments(
    output_dir="./mlm_finetune_result",
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_steps=10,
    report_to="none"
)

# 初始化Trainer启动训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=processed_dataset["train"],
    eval_dataset=processed_dataset["test"],
    data_collator=data_collator,
    tokenizer=tokenizer
)

trainer.train()

补充说明：

如果你的数据存储在本地文本文件中，可以直接用Dataset.from_text("your_file.txt")读取生成数据集，不需要手动加载字符串。
训练过程中数据整理器会自动完成掩码操作，不需要提前对数据集做掩码处理。
如果需要调整掩码策略，只需要修改DataCollatorForLanguageModeling的初始化参数即可。

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