Hugging Face Datasets map batched 返回长度不一致时怎么定位

遇到 Hugging Face Datasets 的 map(batched=True) 报“返回列长度不一致”时，先别改 batch_size。真正要查的是：同一次调用返回的每个列值是否拥有同一个行数。batched=True 允许函数保持原批次大小，也允许像切分长文本那样扩展成更多行，但不能让 input_ids、attention_mask 和新标签各自返回不同长度。原列需要被新列替换时，再用 remove_columns 明确旧列的去留。

官方文档：https://huggingface.co/docs/datasets/en/process

先看清 batched map 的长度契约

官方文档把 batched=True 定义为“每次接收一批样本”。在一次调用里，输入是类似 {"text": [...], "label": [...]} 的字典；返回值也应是字典。返回字典中的列共同决定要写入的数据行，因而它们必须能组成同一张 Arrow 表。

可以把问题先压缩成一张小表：

最容易踩坑的写法是只扩展了一个字段。例如文本被切成 12 个 chunk，却仍返回原来 4 条样本的 label。此时不是“某一列稍微慢了”，而是结果没有唯一的行数。

用最小诊断把不一致列找出来

把完整预处理函数临时缩成“观察输入、调用核心逻辑、检查输出”三段。下面的诊断器不执行模型推理，只负责在写入前报告哪一列改变了长度：

def inspect_batch(batch, transform):
    # 取第一列作为输入行数基准，空批次直接返回，避免 min/max 误判
    input_lengths = {name: len(values) for name, values in batch.items()}
    if not input_lengths:
        return {}
    # 所有输入列通常应先保持同一行数，先发现上游数据已坏的情况
    if len(set(input_lengths.values())) != 1:
        raise ValueError(f"输入列长度已不一致: {input_lengths}")
    output = transform(batch)
    output_lengths = {name: len(values) for name, values in output.items()}
    # batched 返回值必须能拼成一张表，所有输出列使用同一个目标长度
    if output_lengths and len(set(output_lengths.values())) != 1:
        raise ValueError(f"返回列长度不一致: {output_lengths}")
    return output

诊断结果重点看两组数字：输入是否都是 N，以及输出是否都是同一个 M。若输出列长度都等于 N，问题多半在某一列内部的嵌套结构或 Features 类型；若输出都等于 M，说明扩行逻辑是一致的；若只有一列不同，优先回到那一列的列表推导式。

不要把 writer_batch_size 当成长度修复开关。它只影响写入节奏，不能把一张行数不完整的表变成合法表。

remove_columns 应该怎样配

当 tokenizer 根据 text 生成 input_ids，同时不再保留原始文本时，可以这样写。函数执行时仍能读取 text，删除发生在映射结果合并时：

def tokenize_batch(batch, tokenizer):
    # tokenizer 对整批文本返回成组字段；这些字段共享批次数量
    encoded = tokenizer(batch["text"], truncation=True, padding=False)
    # 只返回要新增的字段，旧 text 由 map 的 remove_columns 处理
    return {
        "input_ids": encoded["input_ids"],
        "attention_mask": encoded["attention_mask"],
    }

tokenized = dataset.map(
    lambda batch: tokenize_batch(batch, tokenizer),
    batched=True,
    remove_columns=["text"],  # 映射函数仍可先读取 text
)

如果函数保留了 text，返回的新列就必须仍和这批文本拥有相同的外层行数。若是切块扩行，则应同时生成与 M 匹配的 source_id、label 等列，或者先删掉不再具有一一对应关系的旧列。

把修复固定成可复用模板

修复时按下面顺序排查，通常比反复增减并行参数更快：

1. 用一个很小的 batch_size 复现，记录每个输入和输出字段的外层长度。
2. 把长度断言放在映射函数返回前，先解决列之间的行数问题。
3. 确认输出列的嵌套层级和声明的 Features 一致，尤其是分页标签、二维 token 列和数组字段。
4. 决定是保留原列、替换原列，还是扩展为新行；不要让三种语义混在一次返回里。

切块扩行可以用同一个目标列表驱动多个字段，避免某个字段忘记同步：

def split_batch(batch, width=128):
    # 每个输入样本产生若干块，先统一收集所有目标行
    chunks, source_ids = [], []
    for source_id, text in enumerate(batch["text"]):
        parts = [text[i:i + width] for i in range(0, len(text), width)]
        chunks.extend(parts)
        source_ids.extend([source_id] * len(parts))
    # 两个返回列都使用同一个新长度，才能组成合法结果表
    return {"chunk": chunks, "source_id": source_ids}

最后专门覆盖四个边界：原样补列、单样本切成多行、删除原列、最后一个不足 batch_size 的批次。这样可以区分“长度契约错误”和“空批次/尾批次策略”两类问题。

常见问题

batched=True 返回值一定要和输入一样长吗？

不一定。可以扩行或减少行，但同一次返回中的各列必须共享同一个结果行数，且保留的原列也要和这个行数相容。

remove_columns 会让函数拿不到原始字段吗？

不会。它作用于映射结果合并阶段，函数仍可先读取被列出的输入字段。

为什么把 batch_size 调小后偶尔不报错？

这通常只是绕开了触发条件，例如某个尾批次没有走同一分支；长度契约问题仍在，应先用诊断器检查每个批次。