Stream

El streaming de datasets te permite trabajar con un dataset sin descargarlo. Los datos se transmiten mientras iteras sobre el dataset. Esto es especialmente útil cuando:

• No quieres esperar a que se descargue un conjunto de datos extremadamente grande.
• El tamaño del conjunto de datos supera la cantidad de espacio en disco disponible en tu computadora.
• Desea explorar rápidamente solo algunas muestras de un conjunto de datos.

Por ejemplo, la división en inglés del conjunto de datos HuggingFaceFW/fineweb tiene 45 terabytes, pero puedes usarlo al instante con streaming. Para transmitir un conjunto de datos, establece streaming=True en load_dataset() como se muestra a continuación:

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> print(next(iter(dataset)))
{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012\nWhen the first serious bout of tornadoes of 2012 blew through middle America in the middle of the night, they touched down in places hours from any AP bureau...', ...,
 'language_score': 0.9721424579620361, 'token_count': 717}

El streaming de conjuntos de datos también te permite trabajar con un conjunto de datos compuesto por archivos locales sin realizar ninguna conversión. En este caso, los datos se transmiten desde los archivos locales a medida que iteras sobre el conjunto de datos. Esto es especialmente útil cuando:

• No quieres esperar a que un conjunto de datos local extremadamente grande se convierta a Arrow.
• El tamaño de los archivos convertidos excedería la cantidad de espacio en disco disponible en su computadora.
• Desea explorar rápidamente solo algunas muestras de un conjunto de datos.
• Desea cargar solo ciertas columnas o filtrar de manera eficiente un conjunto de datos Parquet.

Por ejemplo, puede transmitir un conjunto de datos local de cientos de archivos JSONL comprimidos como oscar-corpus/OSCAR-2201 para usarlo de inmediato:

>>> from datasets import load_dataset
>>> data_files = {'train': 'path/to/OSCAR-2201/compressed/en_meta/*.jsonl.gz'}
>>> dataset = load_dataset('json', data_files=data_files, split='train', streaming=True)
>>> print(next(iter(dataset)))
{'id': 0, 'text': 'Founded in 2015, Golden Bees is a leading programmatic recruitment platform dedicated to employers, HR agencies and job boards. The company has developed unique HR-custom technologies and predictive algorithms to identify and attract the best candidates for a job opportunity.', ...

Parquet es un formato columnar que permite transmitir y cargar solo un subconjunto de columnas e ignorar las columnas no deseadas. Parquet también almacena metadatos, como estadísticas de columnas (a nivel de archivo y grupo de filas), lo que permite un filtrado eficiente. Utilice los argumentos columns y filters de datasets.packaged_modules.parquet.ParquetConfig para transmitir conjuntos de datos Parquet, seleccionar columnas y aplicar filtros:

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, columns=["url", "date"])
>>> print(next(iter(dataset)))
{'url': 'http://%20jwashington@ap.org/Content/Press-Release/2012/How-AP-reported-in-all-formats-from-tornado-stricken-regions', 'date': '2013-05-18T05:48:54Z'}
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, filters=[("language_score", ">=", 0.99)])
>>> print(next(iter(dataset)))
{'text': 'Everyone wishes for something. And lots of people believe they know how to make their wishes come true with magical thinking.\nWhat is it? "Magical thinking is a belief in forms of causation, with no known physical basis," said Professor Emily Pronin of Princeton...', ...,
 'language_score': 0.9900368452072144, 'token_count': 716}

Cargar un conjunto de datos en modo de transmisión crea una nueva instancia de tipo de conjunto de datos (en lugar del objeto clásico Dataset), conocido como IterableDataset. Este tipo especial de conjunto de datos tiene su propio conjunto de métodos de procesamiento que se muestran a continuación.

Un IterableDataset es útil para tareas iterativas, como el entrenamiento de un modelo. No se debe usar un IterableDataset para tareas que requieren acceso aleatorio a los ejemplos, ya que hay que iterar sobre todos ellos mediante un bucle for. Obtener el último ejemplo de un dataset iterable requeriría iterar sobre todos los ejemplos anteriores. Puede encontrar más detalles en la guía Dataset vs. IterableDataset.

Indexación de columnas

A veces es conveniente iterar sobre los valores de una columna específica. Afortunadamente, un IterableDataset admite el indexado por columnas:

>>> from datasets import load_dataset
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> print(next(iter(dataset["text"])))
Beginners BBQ Class Taking Place in Missoula!...

Convertir desde un Dataset

Si tienes un objeto Dataset existente, puedes convertirlo en un IterableDataset con la función to_iterable_dataset(). Esto es en realidad más rápido que establecer el argumento streaming=True en load_dataset() porque los datos se transmiten desde archivos locales.

>>> from datasets import load_dataset

# faster 🐇
>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset()

# slower 🐢
>>> iterable_dataset = load_dataset("ethz/food101", streaming=True)

La función to_iterable_dataset() admite el particionado (sharding) cuando se instancia el IterableDataset. Esto resulta útil al trabajar con conjuntos de datos grandes, cuando se desea barajar el conjunto de datos o habilitar una carga paralela rápida con un DataLoader de PyTorch.

>>> import torch
>>> from datasets import load_dataset

>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset(num_shards=64) # shard the dataset
>>> iterable_dataset = iterable_dataset.shuffle(buffer_size=10_000)  # shuffles the shards order and use a shuffle buffer when you start iterating
dataloader = torch.utils.data.DataLoader(iterable_dataset, num_workers=4)  # assigns 64 / 4 = 16 shards from the shuffled list of shards to each worker when you start iterating

Barajar

Al igual que con un objeto Dataset regular, también puedes barajar un IterableDataset con IterableDataset.shuffle().

El argumento buffer_size controla el tamaño del búfer desde el que se seleccionan ejemplos de forma aleatoria. Supongamos que su conjunto de datos tiene un millón de ejemplos y establece buffer_size en diez mil. IterableDataset.shuffle() seleccionará aleatoriamente ejemplos de los primeros diez mil ejemplos del búfer. Los ejemplos seleccionados en el búfer se reemplazan con nuevos ejemplos. De forma predeterminada, el tamaño del búfer es 1.000.

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> shuffled_dataset = dataset.shuffle(seed=42, buffer_size=10_000)

IterableDataset.shuffle() llena el búfer utilizando hasta max_buffer_input_shards fragmentos a la vez (el valor predeterminado es 10), lo que mejora considerablemente la calidad del mezclado. También mezcla el orden de los fragmentos si el conjunto de datos está dividido en varios archivos.

Si estás transmitiendo un conjunto de datos Parquet con solo uno o pocos archivos, usa reshard() antes de barajar. Esto es útil para dividir el conjunto de datos por grupo de filas en lugar de por archivo, mejorando la calidad del barajado.

Barajar

A veces puede que desees reordenar el conjunto de datos después de cada época. Esto requerirá que establezcas una semilla diferente para cada época. Usa IterableDataset.set_epoch() entre épocas para indicar al conjunto de datos en qué época te encuentras.

Tu semilla se convierte efectivamente en: initial seed + current epoch.

>>> for epoch in range(epochs):
...     shuffled_dataset.set_epoch(epoch)
...     for example in shuffled_dataset:
...         ...

Dividir el conjunto de datos

Puedes dividir tu conjunto de datos de una de dos formas:

• IterableDataset.take() devuelve los primeros n ejemplos de un conjunto de datos:

>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> dataset_head = dataset.take(2)
>>> list(dataset_head)
[{'text': "How AP reported in all formats from tor...},
 {'text': 'Did you know you have two little yellow...}]

• IterableDataset.skip() omite los primeros n ejemplos de un conjunto de datos y devuelve los ejemplos restantes:

>>> train_dataset = shuffled_dataset.skip(1000)

take y skip impiden futuras llamadas a shuffle porque fijan el orden de los fragmentos. Debes shuffle tu conjunto de datos antes de dividirlo.

Fragmento

🤗 Datasets admite el particionado (sharding) para dividir un conjunto de datos muy grande en un número predefinido de fragmentos. Especifique el parámetro num_shards en shard() para determinar el número de fragmentos en los que se dividirá el conjunto de datos. También deberá proporcionar el fragmento que desea devolver mediante el parámetro index.

Por ejemplo, el conjunto de datos amazon_polarity tiene 4 fragmentos (en este caso, son 4 archivos Parquet):

>>> from datasets import load_dataset
>>> dataset = load_dataset("fancyzhx/amazon_polarity", split="train", streaming=True)
>>> print(dataset)
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 4
})

Tras dividir el conjunto de datos en dos fragmentos, el primero solo tendrá 2 fragmentos:

>>> dataset.shard(num_shards=2, index=0)
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 2
})

Para aumentar el número de fragmentos de un conjunto de datos, puedes usar IterableDataset.reshard():

>>> dataset.reshard()
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 3600
})

El mecanismo de re-particionado depende del formato de archivo del conjunto de datos. Por ejemplo, en el caso de Parquet, se re-particiona utilizando grupos de filas en lugar de tener un archivo por partición. Consulte cómo funciona para cada formato en la documentación de IterableDataset.reshard().

Si tu conjunto de datos tiene dataset.num_shards==1 incluso después de redividirlo, debes dividirlo en fragmentos usando IterableDataset.skip() y IterableDataset.take() en su lugar.

Concatenar

Los conjuntos de datos separados pueden concatenarse si comparten los mismos tipos de columnas. Concatene los conjuntos de datos con concatenate_datasets():

>>> from datasets import concatenate_datasets, load_dataset

>>> stories = load_dataset("ajibawa-2023/General-Stories-Collection", split="train", streaming=True)
>>> stories = stories.select_columns(["text"])  # only keep the 'text' column
IterableDataset({
    features: Unknown,
    num_shards: 10
})

>>> wiki = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)
>>> wiki = wiki.select_columns(["text"])  # only keep the 'text' column
IterableDataset({
    features: ['text'],
    num_shards: 41
})

>>> bert_dataset = concatenate_datasets([stories, wiki])
IterableDataset({
    features: ['text'],
    num_shards: 51
})

Los fragmentos del dataset concatenado son los fragmentos de los datasets de entrada.

También puede concatenar dos conjuntos de datos horizontalmente estableciendo axis=1, siempre que los conjuntos de datos tengan el mismo número de filas:

>>> from datasets import IterableDataset
>>> stories_ids = IterableDataset.from_dict({"ids": list(range(num_stories))})
>>> stories_with_ids = concatenate_datasets([stories, stories_ids], axis=1)

En este caso, el conjunto de datos concatenado solo tiene 1 fragmento, para evitar terminar con fragmentos desalineados de los conjuntos de datos de entrada.

Interleaved

interleave_datasets() puede combinar un IterableDataset con otros conjuntos de datos si tienen los mismos tipos de columnas. El conjunto de datos combinado devuelve ejemplos alternos de cada uno de los conjuntos de datos originales.

>>> from datasets import interleave_datasets
>>> es_dataset = load_dataset('allenai/c4', 'es', split='train', streaming=True)
IterableDataset({
    features: Unknown,
    num_shards: 2048
})
>>> fr_dataset = load_dataset('allenai/c4', 'fr', split='train', streaming=True)
IterableDataset({
    features: Unknown,
    num_shards: 2048
})

>>> multilingual_dataset = interleave_datasets([es_dataset, fr_dataset])
IterableDataset({
    features: ['text', 'timestamp', 'url'],
    num_shards: 2048
})
>>> list(multilingual_dataset.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
 {'text': 'Le sacre de philippe ier, 23 mai 1059 - Compte Rendu...'}]

Defina las probabilidades de muestreo de cada uno de los conjuntos de datos originales para tener más control sobre cómo se muestrean y combinan. Establezca el argumento probabilities con las probabilidades de muestreo deseadas:

>>> multilingual_dataset_with_oversampling = interleave_datasets([es_dataset, fr_dataset], probabilities=[0.8, 0.2], seed=42)
>>> list(multilingual_dataset_with_oversampling.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
 {'text': 'Chevrolet Cavalier Usados en Bogota - Carros en Vent...'}]

Alrededor del 80% del conjunto de datos final está compuesto por el es_dataset, y el 20% por el fr_dataset.

También puede especificar stopping_strategy. La estrategia predeterminada, first_exhausted, es una estrategia de muestreo parcial, es decir, la construcción del conjunto de datos se detiene en cuanto uno de los conjuntos de datos se queda sin muestras. Puede especificar stopping_strategy=all_exhausted para ejecutar una estrategia de muestreo excesivo. En este caso, la construcción del conjunto de datos se detiene en cuanto todas las muestras de todos los conjuntos de datos se hayan añadido al menos una vez. En la práctica, esto significa que si un conjunto de datos se agota, volverá al principio de dicho conjunto de datos hasta que se alcance el criterio de parada. Tenga en cuenta que si no se especifican probabilidades de muestreo, el nuevo conjunto de datos tendrá max_length_datasets*nb_dataset samples. También existe stopping_strategy=all_exhausted_without_replacement para garantizar que cada muestra se vea exactamente una vez.

Para garantizar un paralelismo adecuado mediante el uso de sharding, los shards del dataset intercalado contienen al menos 1 shard de cada dataset de entrada. Por lo tanto, el nivel de sharding del dataset intercalado es el nivel de sharding mínimo de los datasets de entrada.

Por ejemplo, si los conjuntos de datos de entrada tienen respectivamente 32, 48 y 128 fragmentos, entonces el conjunto de datos intercalado tiene 32 = min(32, 48, 128) fragmentos, y cada nuevo fragmento tiene 1 fragmento del primer conjunto de datos, 1-2 fragmentos del segundo conjunto de datos y 4 fragmentos del tercer conjunto de datos.

Renombrar, eliminar y convertir

Los siguientes métodos permiten modificar las columnas de un conjunto de datos. Estos métodos son útiles para renombrar o eliminar columnas y para cambiar las columnas a un nuevo conjunto de características.

Renombrar

Use IterableDataset.rename_column() cuando necesite renombrar una columna en su conjunto de datos. Las características asociadas con la columna original se mueven de hecho bajo el nuevo nombre de columna, en lugar de simplemente reemplazar la columna original en su lugar.

Proporcione a IterableDataset.rename_column() el nombre de la columna original y el nuevo nombre de la columna:

>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.rename_column("text", "content")

Eliminar

Cuando necesite eliminar una o más columnas, proporcione a IterableDataset.remove_columns() el nombre de la columna a eliminar. Para eliminar más de una columna, proporcione una lista de nombres de columnas:

>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.remove_columns('timestamp')

Cast

IterableDataset.cast() cambia el tipo de característica de una o más columnas. Este método toma tus nuevas Features como argumento. El siguiente código de ejemplo muestra cómo cambiar los tipos de característica de ClassLabel y Value:

>>> from datasets import load_dataset
>>> dataset = load_dataset('nyu-mll/glue', 'mrpc', split='train', streaming=True)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['not_equivalent', 'equivalent']),
'idx': Value('int32')}

>>> from datasets import ClassLabel, Value
>>> new_features = dataset.features.copy()
>>> new_features["label"] = ClassLabel(names=['negative', 'positive'])
>>> new_features["idx"] = Value('int64')
>>> dataset = dataset.cast(new_features)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['negative', 'positive']),
'idx': Value('int64')}

El casting solo funciona si el tipo de característica original y el nuevo tipo de característica son compatibles. Por ejemplo, puedes hacer un casting de una columna con el tipo de característica Value('int32') a Value('bool') si la columna original solo contiene unos y ceros.

Use IterableDataset.cast_column() para cambiar el tipo de característica de una sola columna. Pasa el nombre de la columna y su nuevo tipo de característica como argumentos:

>>> dataset.features
{'audio': Audio(sampling_rate=44100, mono=True)}

>>> dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
>>> dataset.features
{'audio': Audio(sampling_rate=16000, mono=True)}

Map

De manera similar a la función Dataset.map() para un Dataset regular, 🤗 Datasets ofrece IterableDataset.map() para procesar un IterableDataset. IterableDataset.map() aplica el procesamiento al vuelo mientras se transmiten los ejemplos.

Permite aplicar una función de procesamiento a cada ejemplo de un conjunto de datos, de forma independiente o en lotes. Esta función incluso puede crear nuevas filas y columnas.

El siguiente ejemplo demuestra cómo tokenizar un IterableDataset. La función debe aceptar y devolver un dict:

>>> def add_prefix(example):
...     example['text'] = 'My text: ' + example['text']
...     return example

A continuación, aplique esta función al conjunto de datos con IterableDataset.map():

>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> updated_dataset = dataset.map(add_prefix)
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...',
  'timestamp': '2019-04-25 12:57:54',
  'url': 'https://klyq.com/beginners-bbq-class-taking-place-in-missoula/'},
 {'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...',
  'timestamp': '2019-04-21 10:07:13',
  'url': 'https://forums.macrumors.com/threads/restore-from-larger-disk-to-smaller-disk.1311329/'},
 {'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...',
  'timestamp': '2019-04-25 10:40:23',
  'url': 'https://awishcometrue.com/Catalogs/Clearance/Tweens/V1960-Find-A-Way'}]

Veamos otro ejemplo, pero esta vez eliminarás columnas con IterableDataset.map(). Al eliminar una columna, esta solo se elimina después de que el ejemplo se haya proporcionado a la función mapeada. Esto permite que la función mapeada utilice el contenido de las columnas antes de que se eliminen.

Especifique la columna a eliminar con el argumento remove_columns en IterableDataset.map():

>>> updated_dataset = dataset.map(add_prefix, remove_columns=["timestamp", "url"])
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...'},
 {'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...'},
 {'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...'}]

Procesamiento por lotes

IterableDataset.map() también permite trabajar con lotes de ejemplos. Para operar sobre lotes, establezca batched=True. El tamaño de lote predeterminado es 1000, pero puede ajustarlo con el argumento batch_size. Esto abre la puerta a muchas aplicaciones interesantes, como la tokenización, la división de oraciones largas en fragmentos más cortos y el aumento de datos.

Tokenización

>>> from datasets import load_dataset
>>> from transformers import AutoTokenizer
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
>>> def encode(examples):
...     return tokenizer(examples['text'], truncation=True, padding='max_length')
>>> dataset = dataset.map(encode, batched=True, remove_columns=["text", "timestamp", "url"])
>>> next(iter(dataset))
{'input_ids': [101, 4088, 16912, 22861, 4160, 2465, 2635, 2173, 1999, 3335, ..., 0, 0, 0],
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 0, 0]}

Vea otros ejemplos de procesamiento por lotes en la documentación de procesamiento por lotes con map. Funcionan de la misma manera para conjuntos de datos iterables.

Filtrar

Puede filtrar filas en el conjunto de datos en función de una función de predicado usando Dataset.filter(). Devuelve las filas que cumplen una condición especificada:

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', streaming=True, split='train')
>>> start_with_ar = dataset.filter(lambda example: example['text'].startswith('San Francisco'))
>>> next(iter(start_with_ar))
{'text': 'San Francisco 49ers cornerback Shawntae Spencer will miss the rest of the sea...}

Dataset.filter() también puede filtrar por índices si estableces with_indices=True:

>>> even_dataset = dataset.filter(lambda example, idx: idx % 2 == 0, with_indices=True)
>>> list(even_dataset.take(3))
[{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012 Whe...},
 {'text': 'Car Wash For Clara! Now is your chance to help! 2 year old Clara Woodward has...},
 {'text': 'Log In Please enter your ECode to log in. Forgotten your eCode? If you create...}]

Lote

El método batch transforma tu IterableDataset en un iterable de lotes. Esto resulta especialmente útil cuando deseas trabajar con lotes en tu bucle de entrenamiento o al utilizar frameworks que esperan entradas por lotes.

También existe la opción de “procesamiento por lotes” al usar la función map para aplicar una función a lotes de datos, lo cual se explica en la sección Map anterior. El método batch descrito aquí es diferente y ofrece una forma más directa de crear lotes a partir de su conjunto de datos.

Puedes usar el método batch de esta manera:

from datasets import load_dataset

# Load a dataset in streaming mode
dataset = load_dataset("some_dataset", split="train", streaming=True)

# Create batches of 32 samples
batched_dataset = dataset.batch(batch_size=32)

# Iterate over the batched dataset
for batch in batched_dataset:
    print(batch)
    break

En este ejemplo, batched_dataset sigue siendo un IterableDataset, pero cada elemento que se genera ahora es un lote de 32 muestras en lugar de una sola muestra. Este loteo se realiza sobre la marcha mientras se itera sobre el conjunto de datos, preservando la naturaleza eficiente en memoria de IterableDataset.

El método de lote también proporciona un parámetro drop_last_batch. Cuando se establece en True, descarta el último lote si es más pequeño que el batch_size especificado. Esto puede ser útil en escenarios donde el procesamiento posterior requiere que todos los lotes tengan el mismo tamaño:

batched_dataset = dataset.batch(batch_size=32, drop_last_batch=True)

Transmisión en un bucle de entrenamiento

IterableDataset se puede integrar en un bucle de entrenamiento. Primero, baraja el conjunto de datos:

>>> seed, buffer_size = 42, 10_000
>>> dataset = dataset.shuffle(seed, buffer_size=buffer_size)

>>> import torch
>>> from torch.utils.data import DataLoader
>>> from transformers import AutoModelForMaskedLM, DataCollatorForLanguageModeling
>>> from tqdm import tqdm
>>> dataset = dataset.with_format("torch")
>>> dataloader = DataLoader(dataset, collate_fn=DataCollatorForLanguageModeling(tokenizer))
>>> device = 'cuda' if torch.cuda.is_available() else 'cpu' 
>>> model = AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased")
>>> model.train().to(device)
>>> optimizer = torch.optim.AdamW(params=model.parameters(), lr=1e-5)
>>> for epoch in range(3):
...     dataset.set_epoch(epoch)
...     for i, batch in enumerate(tqdm(dataloader, total=5)):
...         if i == 5:
...             break
...         batch = {k: v.to(device) for k, v in batch.items()}
...         outputs = model(**batch)
...         loss = outputs[0]
...         loss.backward()
...         optimizer.step()
...         optimizer.zero_grad()
...         if i % 10 == 0:
...             print(f"loss: {loss}")

Guardar un punto de control del conjunto de datos y reanudar la iteración

Si tu bucle de entrenamiento se detiene, es posible que quieras reanudar el entrenamiento desde donde se interrumpió. Para ello, puedes guardar un punto de control (checkpoint) de tu modelo y optimizadores, así como de tu cargador de datos (data loader).

Los conjuntos de datos iterables no permiten el acceso aleatorio a un índice de ejemplo específico para reanudar desde ahí, pero puedes usar IterableDataset.state_dict() y IterableDataset.load_state_dict() para reanudar desde un punto de control, de manera similar a lo que puedes hacer con los modelos y los optimizadores:

>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> for idx, example in enumerate(iterable_dataset):
...     print(example)
...     if idx == 2:
...         state_dict = iterable_dataset.state_dict()
...         print("checkpoint")
...         break
>>> iterable_dataset.load_state_dict(state_dict)
>>> print(f"restart from checkpoint")
>>> for example in iterable_dataset:
...     print(example)

Devuelve:

{'a': 0}
{'a': 1}
{'a': 2}
checkpoint
restart from checkpoint
{'a': 3}
{'a': 4}
{'a': 5}

Por debajo, el dataset iterable mantiene un registro del shard actual que se está leyendo y del índice del ejemplo en el shard actual, y guarda esta información en el state_dict.

Para reanudar desde un punto de control, el conjunto de datos omite todos los fragmentos que se leyeron previamente para reiniciar desde el fragmento actual. Luego, lee el fragmento y omite ejemplos hasta llegar al ejemplo exacto del punto de control.

Por lo tanto, reiniciar un conjunto de datos es bastante rápido, ya que no volverá a leer los fragmentos que ya se han iterado. Sin embargo, reanudar un conjunto de datos generalmente no es instantáneo, ya que debe volver a leer desde el principio del fragmento actual y omitir ejemplos hasta llegar a la ubicación del punto de control.

Esto se puede usar con StatefulDataLoader de torchdata:

>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(iterable_dataset, batch_size=32, num_workers=4)
>>> # checkpoint
>>> state_dict = dataloader.state_dict()  # uses iterable_dataset.state_dict() under the hood
>>> # resume from checkpoint
>>> dataloader.load_state_dict(state_dict)  # uses iterable_dataset.load_state_dict() under the hood

Al reanudar, se continúa exactamente desde el punto en que se guardó el checkpoint, excepto si se usa .shuffle(): los ejemplos de los búferes de mezcla se pierden al reanudar y los búferes se rellenan con datos nuevos.

Guardar

Una vez que tu conjunto de datos iterable esté listo, puedes guardarlo como un Dataset de Hugging Face en formato Parquet y reutilizarlo más adelante con load_dataset().

Guarde su conjunto de datos proporcionando el nombre del repositorio de conjuntos de datos en Hugging Face al que desea guardarlo como argumento de push_to_hub(). Esto itera sobre el conjunto de datos y sube progresivamente los datos a Hugging Face:

dataset.push_to_hub("username/my_dataset")

Por defecto, IterableDataset.push_to_hub() escribe dataset.num_shards archivos. Establezca num_shards para elegir una cantidad diferente. Con más fragmentos de salida que dataset.num_shards, cada fragmento de entrada se distribuye en su propio grupo de fragmentos de salida escribiendo sus ejemplos en ellos de manera round-robin, en una sola pasada sobre los datos. Si hay menos ejemplos que los fragmentos solicitados, los fragmentos restantes son archivos Parquet vacíos.

Si el conjunto de datos consta de varios fragmentos (dataset.num_shards > 1), puede usar varios procesos para subirlo en paralelo. Esto es especialmente útil si aplicó pasos de map() o filter(), ya que se ejecutarán más rápido en paralelo:

dataset.push_to_hub("username/my_dataset", num_proc=8)

Use la función load_dataset() para recargar el conjunto de datos:

from datasets import load_dataset
reloaded_dataset = load_dataset("username/my_dataset")

Exportar

🤗 Datasets también permite la exportación, de modo que puedas trabajar con tu conjunto de datos en otras aplicaciones. La siguiente tabla muestra los formatos de archivo actualmente admitidos para la exportación:

Por ejemplo, exporte su conjunto de datos a un archivo CSV de la siguiente manera:

>>> dataset.to_csv("path/of/my/dataset.csv")

Si tienes un conjunto de datos grande, puedes guardar un archivo por cada fragmento, por ejemplo:

>>> num_shards = dataset.num_shards
>>> for index in range(num_shards):
...     shard = dataset.shard(index, num_shards)
...     shard.to_parquet(f"path/of/my/dataset/data-{index:05d}.parquet")