如何向字典结构的HuggingFace数据集新增labels列解决KeyError报错

现有分词后的DatasetDict类型数据集tokenized_datasets，结构如下：

需要新增labels列，值为input_ids列的副本。直接对DatasetDict调用自定义map函数时报KeyError: 'input_ids'错误，错误代码如下：

def new_column(example):
    example["labels"] = example["input_ids"].copy()
    return example

dataset_new = tokenized_datasets.map(new_column)

DatasetDict是按训练/验证/测试切分（split）存储子数据集的字典结构，抛出KeyError: 'input_ids'的常见原因有两个：

• 至少一个split下不存在input_ids列：绝大多数情况是分词步骤只处理了部分split（比如只处理了训练集，漏了验证/测试集），map遍历到无input_ids的split时触发键错误
• 版本兼容问题：少数旧版本🤗 Datasets库对DatasetDict执行map时的列自动识别逻辑存在bug

方案1：修正map逻辑（推荐） #

先校验所有split的列完整性，再执行map操作，代码可直接运行：

# 先检查所有split是否存在input_ids列，提前定位漏分词的split
for split, ds in tokenized_datasets.items():
    if "input_ids" not in ds.column_names:
        raise ValueError(f"{split}切分集缺少input_ids列，请先完成全部分集的分词操作")

def add_labels(example):
    example["labels"] = example["input_ids"].copy()
    return example

# 显式指定处理参数，避免兼容问题
dataset_new = tokenized_datasets.map(
    add_labels,
    batched=False,
    desc="添加labels列"
)

处理完成后所有split下都会自动生成和input_ids完全一致的labels列，符合大模型微调的数据集格式要求。

方案2：逐split调用add_column #

如果偏好使用add_column接口，遍历每个split单独加列后重新封装为DatasetDict即可：

from datasets import DatasetDict

processed_splits = {}
for split_name, split_ds in tokenized_datasets.items():
    labels_col = split_ds["input_ids"].copy()
    processed_splits[split_name] = split_ds.add_column("labels", labels_col)

dataset_new = DatasetDict(processed_splits)

处理完成后可打印数据集结构确认：

print(dataset_new)

输出内容中每个split的features列表下存在labels列，且形状与input_ids一致即为操作成功。

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