Hugging Face自定义data_collator触发input_ids的KeyError问题求助

核心原因 #

你的数据集大概率被设置了仅输出单个字段的格式，导致传入data_collator的样本只保留了一个字段，而非完整的5个键。常见场景是预处理时调用了set_format()方法，指定仅输出pixel_values这类单字段。

解决步骤 #

1. 检查数据集格式设置 #

在定义train_dataset后，添加代码查看当前数据集的格式配置：

print(train_dataset.format)

如果输出显示仅包含单个字段（比如{"type": "torch", "columns": ["pixel_values"]}），说明格式被限制，这就是问题根源。

2. 重置或重新设置数据集格式 #

执行以下操作恢复全字段访问：

# 方式1：重置为默认格式，保留所有字段
train_dataset.reset_format()
val_dataset.reset_format()

# 方式2：显式指定需要保留的所有字段（替换为你的数据集实际字段）
train_dataset.set_format(type="torch", columns=["pixel_values", "input_ids", "attention_mask", "xxx", "xxx"])
val_dataset.set_format(type="torch", columns=["pixel_values", "input_ids", "attention_mask", "xxx", "xxx"])

3. 优化自定义data_collator（可选） #

可以简化实现逻辑，避免手动遍历的冗余：

def data_collator(data):
    pixel_values = torch.stack([torch.tensor(item["pixel_values"]) for item in data])
    labels = torch.stack([torch.tensor(item["input_ids"]) for item in data])
    return {"pixel_values": pixel_values, "labels": labels}

或者借助官方默认工具简化：

from transformers import default_data_collator

def custom_data_collator(data):
    batch = default_data_collator(data)
    # 将input_ids映射为Seq2Seq任务需要的labels字段
    batch["labels"] = batch.pop("input_ids")
    return batch

4. 验证修复效果 #

在data_collator开头添加打印语句，确认样本字段完整：

def data_collator(data):
    print(data[0].keys())  # 打印第一个样本的所有键，确认input_ids存在
    # 后续原有逻辑...

额外注意点 #

• 如果数据集经过map()预处理，检查映射函数是否不小心丢弃了input_ids字段；
• Seq2SeqTrainer默认会处理部分字段映射，但自定义data_collator时需确保模型所需字段都被正确收集。

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