Mapeamento em lote

A combinação da utilidade de Dataset.map() com o modo de lote é muito poderosa. Ela permite acelerar o processamento e controlar livremente o tamanho do dataset gerado.

Necessidade de velocidade

O objetivo principal do mapeamento em lote é acelerar o processamento. Muitas vezes, é mais rápido trabalhar com lotes de dados do que com exemplos individuais. Naturalmente, o mapeamento em lote se aplica bem à tokenização. Por exemplo, a biblioteca 🤗 Tokenizers funciona mais rápido com lotes porque paraleliza a tokenização de todos os exemplos em um lote.

Tamanho de entrada != tamanho de saída

A capacidade de controlar o tamanho do conjunto de dados gerado pode ser aproveitada em muitos casos de uso interessantes. Na seção map de como fazer, há exemplos de uso do mapeamento em lote para:

• Divida frases longas em trechos mais curtos.
• Aumente um conjunto de dados com tokens adicionais.

É útil entender como isso funciona, para que você possa criar suas próprias formas de usar o mapeamento em lote. Neste ponto, você pode estar se perguntando como controlar o tamanho do dataset gerado. A resposta é: a função mapeada não precisa retornar um lote de saída do mesmo tamanho.

Em outras palavras, a entrada da sua função mapeada pode ser um lote de tamanho N e retornar um lote de tamanho M. A saída M pode ser maior ou menor que N. Isso significa que você pode concatenar seus exemplos, dividi-los e até mesmo adicionar mais exemplos!

No entanto, lembre-se de que todos os valores no dicionário de saída devem conter o mesmo número de elementos que os outros campos no dicionário de saída. Caso contrário, não é possível definir o número de exemplos na saída retornada pela função mapeada. O número pode variar entre lotes sucessivos processados pela função mapeada. No entanto, para um único lote, todos os valores do dicionário de saída devem ter o mesmo comprimento (ou seja, o número de elementos).

Por exemplo, a partir de um conjunto de dados com 1 coluna e 3 linhas, se você usar map para retornar uma nova coluna com o dobro de linhas, você terá um erro. Nesse caso, você acaba com uma coluna com 3 linhas e uma coluna com 6 linhas. Como você pode ver, a tabela não será válida:

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

Para que seja válido, você deve remover uma das colunas:

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

Alternativamente, você pode sobrescrever a coluna existente para obter o mesmo resultado. Por exemplo, veja como duplicar cada linha no dataset sobrescrevendo a coluna "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]