Datasets documentation

批量映射

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

批量映射

将 Dataset.map() 的实用性与批处理模式相结合非常强大。它可以加快处理速度,并让你自由控制生成数据集的大小。

速度需求

批处理映射的主要目标是加速处理。通常,以数据批次为单位进行处理比逐条处理单个样本更快。自然,批处理映射非常适合用于分词。例如,🤗 Tokenizers 库在处理批次时速度更快,因为它会并行处理批次中所有示例的分词。

输入大小 != 输出大小

控制生成数据集大小的能力可用于许多有趣的用例。在如何映射部分,有使用批量映射的示例:

  • 将长句拆分为较短的片段。
  • 使用额外的标记来扩充数据集。

理解其工作原理会很有帮助,这样你就能想出使用批量映射的自定义方法。此时,你可能会想知道如何控制生成数据集的大小。答案是:映射函数返回的输出批次大小不必与输入相同。

换句话说,你的映射函数输入可以是一个大小为 N 的批次,并返回一个大小为 M 的批次。输出 M 可以大于或小于 N。这意味着你可以拼接你的示例,将其拆分,甚至添加更多示例!

但请记住,输出字典中的所有值必须包含与输出字典中其他字段相同数量的元素。否则,将无法确定映射函数返回的输出中的示例数量。该数量可以在映射函数处理的连续批次之间变化。但对于单个批次,输出字典中所有值的长度(即元素数量)应当相同。

例如,对于一个包含 1 列和 3 行的数据集,如果你使用 map 返回一个行数翻倍的新列,那么就会出错。 在这种情况下,你会得到一个包含 3 行的列和一个包含 6 行的列。如你所见,该表格将无效:

>>> from datasets import Dataset
>>> dataset = Dataset.from_dict({"a": [0, 1, 2]})
>>> dataset.map(lambda batch: {"b": batch["a"] * 2}, batched=True)  # new column with 6 elements: [0, 1, 2, 0, 1, 2]
'ArrowInvalid: Column 1 named b expected length 3 but got length 6'

要使其有效,您必须删除其中一列:

>>> from datasets import Dataset
>>> dataset = Dataset.from_dict({"a": [0, 1, 2]})
>>> dataset_with_duplicates = dataset.map(lambda batch: {"b": batch["a"] * 2}, remove_columns=["a"], batched=True)
>>> len(dataset_with_duplicates)
6

或者,您可以覆盖现有列以达到相同的效果。 例如,以下是通过覆盖列 "a" 来复制数据集中每一行的方法:

>>> from datasets import Dataset
>>> dataset = Dataset.from_dict({"a": [0, 1, 2]})
# overwrites the existing "a" column with duplicated values
>>> duplicated_dataset = dataset.map(
...     lambda batch: {"a": [x for x in batch["a"] for _ in range(2)]},
...     batched=True
... )
>>> duplicated_dataset
Dataset({
    features: ['a'],
    num_rows: 6
})
>>> duplicated_dataset["a"]
[0, 0, 1, 1, 2, 2]
Update on GitHub