使用HuggingFace库微调BLOOM模型时遭遇分词KeyError问题求助

错误原因 #

触发KeyError: 'dialog'的核心原因是代码中访问的字段名与实际数据集字段不匹配，或数据集加载异常：

1. 代码假设conv_ai_2数据集包含dialog字段，但实际可能因数据集版本、加载方式问题，字段名不符；
2. 多进程映射（num_proc=4）可能导致子进程中数据集加载异常，无法正确读取字段。

另外代码还存在两个隐藏问题：

• BLOOM的Tokenizer默认无pad_token，直接分词会触发padding相关错误；
• 对话数据是多轮列表格式，直接传给Tokenizer无法被因果语言模型正确处理，需先拼接为连续文本。

修复步骤 #

步骤1：确认数据集实际字段 #

在加载数据集后添加打印代码，明确字段名：

# 加载数据集后添加
print("数据集字段：", dataset['train'].features.keys())
print("样本示例：", dataset['train'][0])

运行后即可看到实际的对话字段名（conv_ai_2正常情况下字段为dialog，若不符则替换为实际字段名）。

步骤2：修正分词函数 #

针对对话数据集特性，将多轮对话拼接为模型可识别的连续文本，同时设置Tokenizer的pad_token：

# 设置pad_token（BLOOM默认无pad_token，用eos_token替代）
tokenizer.pad_token = tokenizer.eos_token

# 修正后的分词函数
def tokenize_function(examples):
    # 将多轮对话拼接为统一格式文本，用User/Bot区分轮次
    formatted_texts = []
    for dialog in examples["dialog"]:
        dialog_str = ""
        for idx, utterance in enumerate(dialog):
            if idx % 2 == 0:
                dialog_str += f"User: {utterance}\n"
            else:
                dialog_str += f"Bot: {utterance}\n"
        formatted_texts.append(dialog_str.strip())
    
    # 分词并设置padding/truncation，适配模型输入长度
    tokenized = tokenizer(
        formatted_texts,
        truncation=True,
        padding="max_length",
        max_length=512,
        return_tensors="pt"
    )
    # 因果语言模型用input_ids作为labels
    tokenized["labels"] = tokenized["input_ids"].clone()
    return tokenized

步骤3：调整数据集映射参数 #

如果多进程导致字段访问异常，暂时去掉num_proc=4，改用单进程处理：

tokenized_dataset = dataset.map(tokenize_function, batched=True)

后续确认无问题后，可尝试恢复多进程。

步骤4：完整修正后的代码 #

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
import datasets

# 加载模型与Tokenizer
model_name = "bigscience/bloom-560m"
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 设置pad_token
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_name)

# 加载数据集并确认字段
dataset = datasets.load_dataset('conv_ai_2')
print("数据集字段：", dataset['train'].features.keys())
print("样本示例：", dataset['train'][0])

# 修正后的分词函数
def tokenize_function(examples):
    formatted_texts = []
    for dialog in examples["dialog"]:
        dialog_str = ""
        for idx, utterance in enumerate(dialog):
            if idx % 2 == 0:
                dialog_str += f"User: {utterance}\n"
            else:
                dialog_str += f"Bot: {utterance}\n"
        formatted_texts.append(dialog_str.strip())
    
    tokenized = tokenizer(
        formatted_texts,
        truncation=True,
        padding="max_length",
        max_length=512,
        return_tensors="pt"
    )
    tokenized["labels"] = tokenized["input_ids"].clone()
    return tokenized

# 映射分词函数（单进程避免字段访问异常）
tokenized_dataset = dataset.map(tokenize_function, batched=True)

# 拆分数据集
train_dataset = tokenized_dataset['train']
val_dataset = tokenized_dataset['valid']

# 训练参数
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy = "epoch",
    num_train_epochs=1,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    logging_steps=500,
    save_steps=500,
    seed=42,
    learning_rate=5e-5,
    report_to="none"
)

# 初始化Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)

# 开始微调
trainer.train()

# 生成回复函数
def generate_response(input_text, model, tokenizer):
    # 拼接成模型训练时的格式
    input_prompt = f"User: {input_text}\nBot:"
    input_ids = tokenizer.encode(input_prompt, return_tensors='pt')
    chat_history_ids = model.generate(
        input_ids=input_ids,
        max_length=1000,
        do_sample=True,
        top_p=0.9,
        top_k=50,
        pad_token_id=tokenizer.pad_token_id  # 指定pad_token_id避免警告
    )
    return tokenizer.decode(chat_history_ids[0], skip_special_tokens=True).split("Bot:")[-1].strip()

# 测试对话机器人
while True:
    user_input = input("You: ")
    if user_input.lower() in ["quit", "exit"]:
        break
    response = generate_response(user_input, model, tokenizer)
    print("Bot: " + response)

关键说明 #

• 因果语言模型需要连续文本输入，必须将多轮对话拼接为统一格式的字符串；
• BLOOM的Tokenizer默认没有pad_token，必须手动设置，否则训练时会触发padding错误；
• 训练时需要设置labels字段，因果语言模型通常直接用input_ids作为标签（克隆一份避免修改原数据）。

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