Пакетное отображение

Сочетание полезности Dataset.map() с режимом пакетной обработки очень мощно. Это позволяет ускорить обработку и свободно контролировать размер создаваемого набора данных.

Нужна скорость

Основная цель batch mapping — ускорить обработку. Часто работа с батчами данных оказывается быстрее, чем с отдельными примерами. Естественно, batch mapping хорошо подходит для токенизации. Например, библиотека 🤗 Tokenizers работает быстрее с батчами, поскольку параллелизует токенизацию всех примеров в батче.

Размер входа != размер выхода

Возможность контролировать размер генерируемого набора данных можно использовать во многих интересных сценариях. В разделе «Как map» приведены примеры использования batch mapping для:

• Разбивайте длинные предложения на более короткие фрагменты.
• Добавьте в набор данных дополнительные токены.

Понимание принципа работы этого механизма поможет вам придумать собственные способы использования batch mapping. На этом этапе вы, возможно, задаетесь вопросом, как можно контролировать размер генерируемого датасета. Ответ: функция, применяемая при маппинге, не обязана возвращать выходную batch того же размера.

Другими словами, вход вашей отображаемой функции может представлять собой батч размера 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]