Hugging Face Datasets select_columns 怎样只保留训练所需字段

Hugging Face Datasets 里，select_columns 解决的是“数据集已经加载完成，但训练只需要其中几列”的问题。它会返回只保留指定字段的数据集，并同步保留这些字段对应的 features；它不负责选择 train 或 test，也不会替代分词和标签转换。实际使用时先确认字段名，再做列投影，最后检查每个 split 的结构。

最小可用写法是 trimmed = dataset.select_columns(["text", "label"])。如果对象是 DatasetDict，这个选择会应用到其中的各个 split；如果某个 split 没有指定列，应先修正 schema 或单独处理 split。

先分清 Dataset、DatasetDict 与字段集合

Dataset 表示一份数据，DatasetDict 通常把 train、validation、test 等 split 放在同一个字典式对象中。两者都能做列投影，但查看字段的方式略有区别。不要直接凭感觉写列名，先从对象读取真实 schema。

from datasets import DatasetDict, load_dataset

# 读取公开数据集；示例只用于说明字段检查，不代表本机运行结果
ds = load_dataset("rotten_tomatoes")

# DatasetDict 的字段要按 split 查看，避免把 split 名当成列名
for split_name, split_data in ds.items():
    print(split_name, split_data.column_names)
    print(split_data.features)

# 单个 Dataset 则直接查看这一份数据的字段
train_data = ds["train"]
print(train_data.column_names)

这里通常会看到 text 和 label。column_names 适合确认字段是否存在，features 适合确认字段类型，例如标签是否仍是整数、文本是否仍是字符串。两者都比凭经验猜字段可靠。

用 select_columns 保留训练真正需要的字段

列名确认后，把“训练输入”和“训练标签”组成显式清单。对于单个 Dataset，返回值仍是 Dataset；对于 DatasetDict，返回值仍保留原来的 split 结构。这样可以在进入分词或训练器前去掉来源、原始编号等不参与训练的元数据。

from datasets import load_dataset

dataset = load_dataset("rotten_tomatoes")

# 只保留后续训练需要的输入和监督标签
keep_columns = ["text", "label"]
trimmed = dataset.select_columns(keep_columns)

# DatasetDict 的每个 split 都只留下这两列
for split_name, split_data in trimmed.items():
    print(split_name, split_data.column_names)

# 如果只处理训练集，先取 split，再对单个 Dataset 做投影
train_only = dataset["train"].select_columns(["text", "label"])

select_columns 与“先选 split”是两个不同动作：dataset["train"] 选择数据分片，select_columns 选择字段。把它们写成两行，后续排查更清楚。

字段缺失或 split 不一致时怎样定位

如果传入不存在的字段，问题不在训练器，而在列清单与当前 schema 不一致。DatasetDict 还要额外注意：不同 split 可能来自不同文件或不同预处理阶段，字段不一定完全相同。先逐个 split 检查，再决定统一补字段、改清单，还是只处理字段完整的 split。

def check_columns(dataset, required):
    """检查每个 split 是否都包含训练所需字段。"""
    required = set(required)
    for split_name, split_data in dataset.items():
        actual = set(split_data.column_names)
        missing = required - actual
        if missing:
            # 让错误带上 split 名，避免只看到模糊的 KeyError
            raise ValueError(f"{split_name} 缺少字段: {sorted(missing)}")

check_columns(dataset, ["text", "label"])
trimmed = dataset.select_columns(["text", "label"])

如果训练数据的输入列叫 sentence，就应该先明确重命名或修改清单，不能把 text 当作通用保留名。列投影不会自动推断语义，也不会把 sentence 映射成模型需要的 input_ids。

交给 Trainer 前做一次最小核对

完成列投影后，建议只做三项检查：字段集合是否准确、标签类型是否符合损失函数、样本是否仍包含后续预处理需要的原始文本。若下一步要调用 tokenizer，就必须保留文本列；若已经完成 tokenization，则应保留训练器读取的 input_ids、attention_mask 和标签列。

# 投影后先看字段和第一条样本，再连接后续训练流程
print(trimmed["train"].column_names)
print(trimmed["train"].features)
print(trimmed["train"][0])  # 用一条样本确认键名没有被误删

这一步的重点不是验证训练效果，而是确认数据接口没有断裂。select_columns 只做字段层面的收敛，分词、标签编码、批处理和张量格式仍应由各自的步骤负责。

常见问题

select_columns 会改变原始 Dataset 吗？

应把返回值当作新的数据集对象使用，并显式赋给变量，例如 trimmed = dataset.select_columns(...)。不要依赖原变量已经自动变成筛选后的结构。

DatasetDict 的每个 split 必须有完全相同的字段吗？

如果一次对整个 DatasetDict 使用同一组列名，实践上应先确认所有 split 都包含这些列。字段不一致时，可以修正上游 schema，或先取出字段完整的 split 分别处理。

select_columns 能代替 remove_columns 吗？

两者目标相反：select_columns 给出要保留的白名单，适合训练字段明确的场景；remove_columns 给出要删除的黑名单，适合只排除少量已知元数据的场景。训练入口更建议白名单思路。

筛选后为什么仍然没有 input_ids？

因为列投影不会执行 tokenizer。保留 text 后，还需要单独调用分词函数；如果已经完成分词，则在投影清单中保留 input_ids 等训练字段。