バッチマッピング

Dataset.map()のユーティリティとバッチモードを組み合わせることは非常に強力です。これにより、処理を高速化し、生成されるデータセットのサイズを自由に制御できます。

スピードの必要性

バッチマッピングの主な目的は、処理を高速化することです。多くの場合、単一の例ではなく、データの一括（バッチ）で扱う方が高速です。当然ながら、バッチマッピングはトークナイゼーションに非常に適しています。例えば、🤗 Tokenizers ライブラリは、バッチ内のすべての例のトークナイゼーションを並列化するため、バッチで扱う方が高速に動作します。

入力サイズ != 出力サイズ

生成されるデータセットのサイズを制御できる機能は、多くの興味深いユースケースに活用できます。mapの使い方セクションには、バッチマッピングを使用して以下を行う例が示されています。

• 長い文を短いチャンクに分割します。
• 追加のトークンでデータセットを拡張します。

この仕組みを理解しておくことで、バッチマッピングを独自の手法で活用できるようになります。ここで、生成されるデータセットのサイズをどのように制御できるか疑問に思われるかもしれません。答えは、マッピング関数は、入力バッチと同じサイズの出力バッチを返す必要はないということです。

つまり、マッピング関数の入力はサイズNのバッチであり、出力はサイズMのバッチを返すことができます。出力のMは、Nより大きくも小さくもなります。これは、サンプルを連結したり、分割したり、さらにはサンプルを追加したりできることを意味します。

ただし、出力辞書のすべての値は、出力辞書の他のフィールドと同じ要素数を含む必要があります。そうしないと、マッピング関数が返す出力のサンプル数を定義できません。マッピング関数が処理する連続するバッチ間で、この数は変動する可能性があります。ただし、1つのバッチ内では、出力辞書のすべての値は同じ長さ（つまり、要素数）を持つ必要があります。

例えば、1列3行のデータセットから、mapを使用して行数が2倍になる新しい列を返すと、エラーが発生します。 この場合、3行の列が1つと、6行の列が1つになります。ご覧のとおり、テーブルは有効ではありません:

>>> 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'

有効にするには、列を1つ削除する必要があります:

>>> 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]