Hugging Face Datasets 怎样用 map 批处理并保留列

在处理训练集时，最容易混淆的是 map 的“批处理”与“删除列”其实是两件事：batched=True 只改变回调函数收到的数据形状，默认不会把原列悄悄删掉；是否删除列，要由 remove_columns 明确表达。

官方地址：https://huggingface.co/docs/datasets

本文只围绕一个可复用的目标展开：让回调一次接收一批样本，生成 input_length 等新列，同时保留 text、label；当原始文本只作为中间输入时，再显式移除它。

先记住 map 的列规则

Dataset.map 默认以单条样本调用函数。设置 batched=True 后，函数收到的是“列名到列表”的字典，例如 {"text": [...], "label": [...]}。函数返回的新增列会并入数据集，原来的列仍会保留，除非你传入 remove_columns。

因此，批处理函数要遵守两个边界：读取的每个输入列都是列表；如果是普通等长映射，返回的新列长度要和当前批次长度一致。下面的图是字段关系说明，不是运行截图。

用 batched=True 新增一列并保留原列

假设数据集有 text 和 label 两列，我们想增加文本长度，先不删除任何字段。回调函数不需要复制整个批次，只返回要新增或更新的列：

from datasets import Dataset

# 这里的列按“整批列表”传入，而不是传入一条字符串
dataset = Dataset.from_dict({
    "text": ["map 可以接收一批样本", "保留原列需要显式理解默认行为"],
    "label": [0, 1],
})

def add_input_length(batch):
    # batched=True 时 batch["text"] 是字符串列表
    # 返回列表长度必须与本批次样本数一致
    return {
        "input_length": [len(text) for text in batch["text"]]
    }

processed = dataset.map(
    add_input_length,
    batched=True,
    batch_size=2,  # 控制一次交给回调的样本数
)

# 结果仍包含 text、label，并额外包含 input_length
print(processed.column_names)

这段写法的关键不是 batch_size=2 这个具体数字，而是返回值的形状：输入两列是列表，新增列也是列表。对于更大的数据集，可以按内存和处理函数的吞吐调整 batch_size；它只决定回调每次看到多少条样本，不代表最终数据集会变成“每批一行”。

怎样保留原列，又明确删除临时列

文本分词是另一种常见情况：text 只用于产生 input_ids 和 attention_mask，训练阶段不再需要原始字符串。这时把待删除列传给 remove_columns 更清楚。

def tokenize_batch(batch):
    # 示例 tokenizer 代表已经准备好的分词器
    encoded = tokenizer(
        batch["text"],
        truncation=True,
        padding=False,
    )
    # 返回的新列与输入批次逐项对应
    return {
        "input_ids": encoded["input_ids"],
        "attention_mask": encoded["attention_mask"],
    }

tokenized = dataset.map(
    tokenize_batch,
    batched=True,
    remove_columns=["text"],  # 回调读取 text 后，再从结果中移除它
)

# label 会保留；最终字段包含 label、input_ids、attention_mask
print(tokenized.column_names)

需要注意的是，remove_columns 不会阻止回调读取这些列。它们先作为输入提供给函数，映射结果组装时才被移除。因此可以安全地用原始 text 生成新字段，再决定最终数据集是否保留它。

变长输出时不要套用等长映射规则

批处理函数并不只能一条输入对应一条输出。例如把一批长文本切成更多片段时，返回列可以比输入批次更长，但返回结果中的各列长度必须彼此一致。此时常见做法是用 remove_columns=dataset.column_names，避免旧列长度与新片段数量冲突。

def split_batch(batch):
    chunks = []
    sources = []
    for source, text in zip(batch["source"], batch["text"]):
        # 这里只展示按空格切分的示意逻辑，实际项目应使用自己的分片规则
        words = text.split()
        for start in range(0, len(words), 32):
            chunks.append(" ".join(words[start:start + 32]))
            sources.append(source)

    # 两个输出列都按“片段”计数，而不是按原始样本计数
    return {"chunk": chunks, "source": sources}

expanded = dataset.map(
    split_batch,
    batched=True,
    remove_columns=dataset.column_names,  # 删除旧列，保留新的等长输出列
)

如果只是给每条样本增加字段，就保持等长返回；如果一个批次会展开成更多行，就同时规划所有输出列的长度，并考虑删除不再匹配的原列。drop_last_batch=True 只适用于你明确愿意丢弃最后不足一批的数据，不能用来修复返回列长度不一致。

几个容易踩到的边界

• 把批次当成单条样本：开启 batched=True 后，batch["text"] 是列表，直接把它当字符串拼接会得到错误的输入形状。
• 误以为返回字典会覆盖全部列：返回新列只表示新增或更新字段，未列出的旧字段默认仍在；需要清理时使用 remove_columns 或单独的列删除操作。
• 不同输出列长度不同：普通映射要求新增列与批次等长；变长映射则要求所有返回列彼此等长，并处理原列的去留。
• 批次过大：batch_size 越大，单次处理的列表越长，内存压力也可能越高。先按数据字段大小设置一个可控值，再根据实际任务调整。

总结

记住三句话就够了：batched=True 改变的是回调输入形状；返回新列不会自动删除原列；列的最终生命周期要用 remove_columns 明确写出来。等长特征计算适合直接新增列，拆分文本等变长处理则要统一所有输出列的长度，并重新决定旧列是否保留。

相关问题

map 的 batch_size 默认是多少？
官方 API 参考中默认值是 1000；具体任务应结合单条样本大小和回调的内存占用调整。

remove_columns 会不会让回调拿不到原列？
不会。映射函数先获得输入，指定列在结果组装阶段才被移除，所以可以用原列产生新列。

什么时候不该使用 batched=True？
当处理逻辑天然依赖单条样本、且批量化不能减少开销时，保持默认的单样本函数会更直观；如果要调用向量化 tokenizer 或批量模型接口，通常更适合批处理。