Stream

O streaming de datasets permite trabalhar com um dataset sem baixá-lo. Os dados são transmitidos conforme você itera sobre o dataset. Isso é especialmente útil quando:

• Você não quer esperar o download de um conjunto de dados extremamente grande.
• O tamanho do conjunto de dados excede a quantidade de espaço em disco disponível no seu computador.
• Você deseja explorar rapidamente apenas alguns exemplos de um conjunto de dados.

Por exemplo, a divisão em inglês do conjunto de dados HuggingFaceFW/fineweb tem 45 terabytes, mas você pode usá-lo instantaneamente com streaming. Para transmitir um conjunto de dados, defina streaming=True em load_dataset(), conforme mostrado abaixo:

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

O streaming de datasets também permite trabalhar com um dataset feito de arquivos locais sem realizar nenhuma conversão. Nesse caso, os dados são transmitidos dos arquivos locais à medida que você itera sobre o dataset. Isso é especialmente útil quando:

• Você não quer esperar que um conjunto de dados local extremamente grande seja convertido para Arrow.
• O tamanho dos arquivos convertidos excederia a quantidade de espaço em disco disponível no seu computador.
• Você deseja explorar rapidamente apenas alguns exemplos de um conjunto de dados.
• Você deseja carregar apenas certas colunas ou filtrar um dataset Parquet de forma eficiente.

Por exemplo, você pode transmitir um conjunto de dados local com centenas de arquivos JSONL compactados, como oscar-corpus/OSCAR-2201, para usá-lo imediatamente:

>>> 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 é um formato colunar que permite transmitir e carregar apenas um subconjunto de colunas e ignorar colunas indesejadas. O Parquet também armazena metadados, como estatísticas de colunas (no nível do arquivo e do grupo de linhas), permitindo filtragem eficiente. Use os argumentos columns e filters de datasets.packaged_modules.parquet.ParquetConfig para transmitir conjuntos de dados Parquet, selecionar colunas e 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}

Carregar um conjunto de dados em modo de streaming cria uma nova instância de tipo de conjunto de dados (em vez do objeto clássico Dataset), conhecido como IterableDataset. Esse tipo especial de conjunto de dados possui seu próprio conjunto de métodos de processamento, mostrados abaixo.

Um IterableDataset é útil para tarefas iterativas, como o treinamento de um modelo. Você não deve usar um IterableDataset para tarefas que exigem acesso aleatório a exemplos, pois é necessário iterar por todos eles usando um laço for. Obter o último exemplo em um dataset iterável exigiria iterar por todos os exemplos anteriores. Você pode encontrar mais detalhes no guia Dataset vs. IterableDataset.

Indexação de colunas

Às vezes, é conveniente iterar sobre os valores de uma coluna específica. Felizmente, um IterableDataset suporta indexação por coluna:

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

Converter de um Dataset

Se você já tem um objeto Dataset, pode convertê-lo em um IterableDataset com a função to_iterable_dataset(). Isso é, na verdade, mais rápido do que definir o argumento streaming=True em load_dataset(), pois os dados são transmitidos a partir de arquivos locais.

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

A função to_iterable_dataset() suporta sharding quando o IterableDataset é instanciado. Isso é útil ao trabalhar com grandes conjuntos de dados, quando você deseja embaralhar o conjunto de dados ou habilitar o carregamento paralelo rápido com um PyTorch DataLoader.

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

Embaralhar

Assim como um objeto Dataset comum, você também pode embaralhar um IterableDataset com IterableDataset.shuffle().

O argumento buffer_size controla o tamanho do buffer a partir do qual os exemplos são amostrados aleatoriamente. Suponha que seu conjunto de dados tenha um milhão de exemplos e você defina o buffer_size como dez mil. O IterableDataset.shuffle() selecionará aleatoriamente exemplos dos primeiros dez mil exemplos no buffer. Os exemplos selecionados no buffer são substituídos por novos exemplos. Por padrão, o tamanho do buffer é 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() preenche o buffer usando até max_buffer_input_shards shards por vez (o padrão é 10), o que melhora significativamente a qualidade da embaralhamento. Ele também embaralha a ordem dos shards se o conjunto de dados estiver dividido em vários arquivos.

Se você estiver transmitindo um dataset Parquet com apenas um ou poucos arquivos, use reshard() antes de embaralhar. Isso é útil para dividir o dataset por grupo de linhas em vez de por arquivo, melhorando a qualidade do embaralhamento.

Embaralhar

Às vezes, você pode querer reembaralhar o dataset após cada epoch. Isso exigirá que você defina uma seed diferente para cada epoch. Use IterableDataset.set_epoch() entre os epochs para informar ao dataset em qual epoch você está.

Sua semente efetivamente se torna: initial seed + current epoch.

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

Dividir o conjunto de dados

Você pode dividir seu conjunto de dados de uma das duas formas:

• IterableDataset.take() retorna os primeiros n exemplos de um conjunto de dados:

>>> 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 os primeiros n exemplos de um conjunto de dados e retorna os exemplos restantes:

>>> train_dataset = shuffled_dataset.skip(1000)

take e skip impedem futuras chamadas para shuffle porque fixam a ordem dos shards. Você deve shuffle seu dataset antes de dividi-lo.

Shard

🤗 O Datasets oferece suporte a sharding para dividir um conjunto de dados muito grande em um número predefinido de blocos. Especifique o parâmetro num_shards em shard() para determinar o número de shards em que o conjunto de dados será dividido. Você também precisará fornecer o shard que deseja retornar com o parâmetro index.

Por exemplo, o conjunto de dados amazon_polarity tem 4 shards (neste caso, são 4 arquivos 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
})

Após dividir o conjunto de dados em dois blocos, o primeiro terá apenas 2 shards:

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

Para aumentar o número de shards de um dataset, você pode usar IterableDataset.reshard():

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

O mecanismo de resharding depende do formato do arquivo do conjunto de dados. Por exemplo, para Parquet, ele faz o resharding usando grupos de linhas em vez de ter um arquivo por shard. Veja como funciona para cada formato na documentação de IterableDataset.reshard().

Se o seu conjunto de dados tiver dataset.num_shards==1 mesmo após o resharding, você deve dividi-lo em blocos usando IterableDataset.skip() e IterableDataset.take() em vez disso.

Concatenar

Conjuntos de dados separados podem ser concatenados se compartilharem os mesmos tipos de colunas. Concatene conjuntos de dados com 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
})

Os shards do dataset concatenado são os shards dos datasets de entrada.

Você também pode concatenar dois conjuntos de dados horizontalmente definindo axis=1, desde que os conjuntos de dados tenham o mesmo número de linhas:

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

Nesse caso, o conjunto de dados concatenado possui apenas 1 shard, para evitar terminar com shards desalinhados dos conjuntos de dados de entrada.

Interlevar

interleave_datasets() pode combinar um IterableDataset com outros conjuntos de dados, desde que tenham os mesmos tipos de colunas. O conjunto de dados combinado retorna exemplos alternados de cada um dos conjuntos de dados originais.

>>> 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 as probabilidades de amostragem de cada um dos conjuntos de dados originais para ter mais controle sobre como cada um deles é amostrado e combinado. Defina o argumento probabilities com as probabilidades de amostragem desejadas:

>>> 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...'}]

Cerca de 80% do conjunto de dados final é composto pelo es_dataset, e 20% pelo fr_dataset.

You can also specify the stopping_strategy. The default strategy, first_exhausted, is a subsampling strategy, i.e the dataset construction is stopped as soon one of the dataset runs out of samples. You can specify stopping_strategy=all_exhausted to execute an oversampling strategy. In this case, the dataset construction is stopped as soon as every samples in every dataset has been added at least once. In practice, it means that if a dataset is exhausted, it will return to the beginning of this dataset until the stop criterion has been reached. Note that if no sampling probabilities are specified, the new dataset will have max_length_datasets*nb_dataset samples. There is also stopping_strategy=all_exhausted_without_replacement to ensure that every sample is seen exactly once.

Para garantir o paralelismo adequado usando sharding, os shards do dataset intercalado contêm pelo menos 1 shard de cada dataset de entrada. Portanto, o nível de sharding do dataset intercalado é o nível de sharding mínimo dos datasets de entrada.

Por exemplo, se os conjuntos de dados de entrada tiverem, respectivamente, 32, 48 e 128 shards, então o conjunto de dados intercalado terá 32 = min(32, 48, 128) shards, e cada novo shard terá 1 shard do primeiro conjunto de dados, 1-2 shards do segundo conjunto de dados e 4 shards do terceiro conjunto de dados.

Renomear, remover e converter

Os métodos a seguir permitem modificar as colunas de um conjunto de dados. Esses métodos são úteis para renomear ou remover colunas e para alterar colunas para um novo conjunto de recursos.

Renomear

Use IterableDataset.rename_column() quando precisar renomear uma coluna no seu dataset. Os recursos associados à coluna original são, na verdade, movidos para o novo nome da coluna, em vez de apenas substituir a coluna original no local.

Forneça para IterableDataset.rename_column() o nome da coluna original e o novo nome da coluna:

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

Remover

Quando você precisa remover uma ou mais colunas, forneça a IterableDataset.remove_columns() o nome da coluna a ser removida. Remova mais de uma coluna fornecendo uma lista de nomes de colunas:

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

Cast

IterableDataset.cast() altera o tipo de recurso de uma ou mais colunas. Este método recebe seus novos Features como argumento. O código de exemplo a seguir mostra como alterar os tipos de recurso de ClassLabel e 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')}

O casting só funciona se o tipo de feature original e o novo tipo de feature forem compatíveis. Por exemplo, você pode fazer o casting de uma coluna com o tipo de feature Value('int32') para Value('bool') se a coluna original contiver apenas uns e zeros.

Use IterableDataset.cast_column() para alterar o tipo de recurso de apenas uma coluna. Passe o nome da coluna e seu novo tipo de recurso 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

Assim como a função Dataset.map() para um Dataset regular, o 🤗 Datasets oferece IterableDataset.map() para processar um IterableDataset. IterableDataset.map() aplica o processamento em tempo real à medida que os exemplos são transmitidos.

Permite aplicar uma função de processamento a cada exemplo em um conjunto de dados, de forma independente ou em lotes. Essa função pode até criar novas linhas e colunas.

O exemplo a seguir demonstra como tokenizar um IterableDataset. A função precisa aceitar e retornar um dict:

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

A seguir, aplique esta função ao conjunto de dados com 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'}]

Vamos analisar outro exemplo, mas desta vez você removerá colunas com IterableDataset.map(). Ao remover uma coluna, ela só é removida depois que o exemplo tiver sido fornecido à função mapeada. Isso permite que a função mapeada utilize o conteúdo das colunas antes que sejam removidas.

Especifique a coluna a ser removida com o argumento remove_columns em 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...'}]

Processamento em lote

IterableDataset.map() também suporta o processamento de lotes de exemplos. Para operar em lotes, defina batched=True. O tamanho padrão do lote é 1000, mas você pode ajustá-lo com o argumento batch_size. Isso abre caminho para muitas aplicações interessantes, como tokenização, divisão de frases longas em trechos menores e aumento de dados.

Tokenização

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

Veja outros exemplos de processamento em lote na documentação de processamento em lote com map. Eles funcionam da mesma forma para conjuntos de dados iteráveis.

Filtrar

Você pode filtrar linhas no dataset com base em uma função predica usando Dataset.filter(). Ele retorna as linhas que correspondem a uma condição 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() também pode filtrar por índices se você definir 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

O método batch transforma seu IterableDataset em um iterável de lotes. Isso é particularmente útil quando você deseja trabalhar com lotes no seu loop de treinamento ou ao usar frameworks que esperam entradas em lote.

Há também a opção “Batch Processing” ao usar a função map para aplicar uma função a lotes de dados, que é discutida na seção Map acima. O método batch descrito aqui é diferente e oferece uma maneira mais direta de criar lotes a partir do seu conjunto de dados.

Você pode usar o método batch da seguinte forma:

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

Neste exemplo, batched_dataset ainda é um IterableDataset, mas cada item gerado agora é um lote de 32 amostras em vez de uma única amostra. Esse loteamento é feito em tempo real à medida que você itera sobre o conjunto de dados, preservando a natureza de eficiência de memória do IterableDataset.

O método de lote também fornece um parâmetro drop_last_batch. Quando definido como True, ele descartará o último lote se for menor que o batch_size especificado. Isso pode ser útil em cenários em que o processamento subsequente exige que todos os lotes tenham o mesmo tamanho:

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

Stream em um loop de treinamento

IterableDataset pode ser integrado a um loop de treinamento. Primeiro, embaralhe o conjunto de dados:

>>> 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}")

Salvar um checkpoint de dataset e retomar a iteração

Se o seu loop de treinamento parar, você pode querer reiniciar o treinamento de onde ele estava. Para isso, você pode salvar um checkpoint do seu modelo e dos otimizadores, bem como do seu carregador de dados.

Os conjuntos de dados iteráveis não oferecem acesso aleatório a um índice específico de exemplo para retomar a partir dele, mas você pode usar IterableDataset.state_dict() e IterableDataset.load_state_dict() para retomar a partir de um checkpoint, da mesma forma que você pode fazer para modelos e otimizadores:

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

Retorna:

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

Por baixo dos panos, o dataset iterável mantém o controle do shard atual sendo lido e do índice do exemplo no shard atual e armazena essas informações no state_dict.

Para retomar de um checkpoint, o dataset ignora todos os shards que foram lidos anteriormente para reiniciar a partir do shard atual. Em seguida, ele lê o shard e ignora exemplos até alcançar o exemplo exato do checkpoint.

Portanto, reiniciar um dataset é bastante rápido, pois ele não relerá os shards que já foram iterados. Ainda assim, retomar um dataset geralmente não é instantâneo, pois ele precisa reiniciar a leitura a partir do início do shard atual e pular exemplos até alcançar a localização do checkpoint.

Isso pode ser usado com o StatefulDataLoader do 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

Ao retomar, o processo continua exatamente de onde o checkpoint foi salvo, exceto quando .shuffle() é usado: os exemplos dos buffers de embaralhamento são perdidos ao retomar e os buffers são reabastecidos com novos dados.

Salvar

Quando seu dataset iterável estiver pronto, você pode salvá-lo como um Dataset do Hugging Face no formato Parquet e reutilizá-lo mais tarde com load_dataset().

Salve seu dataset fornecendo o nome do repositório do dataset no Hugging Face para o qual deseja salvá-lo como argumento para push_to_hub(). Isso itera sobre o dataset e faz o upload dos dados progressivamente para o Hugging Face:

dataset.push_to_hub("username/my_dataset")

Por padrão, IterableDataset.push_to_hub() grava dataset.num_shards arquivos. Defina num_shards para escolher uma contagem diferente. Com mais shards de saída do que dataset.num_shards, cada shard de entrada é distribuído sobre seu próprio grupo de shards de saída, gravando seus exemplos neles de forma round-robin, em uma única passagem sobre os dados. Se houver menos exemplos do que os shards solicitados, os shards restantes serão arquivos Parquet vazios.

Se o conjunto de dados consiste em vários shards (dataset.num_shards > 1), você pode usar vários processos para enviá-lo em paralelo. Isso é especialmente útil se você aplicou etapas de map() ou filter(), pois elas serão executadas mais rapidamente em paralelo:

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

Use a função load_dataset() para recarregar o conjunto de dados:

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

Exportar

🤗 O Datasets também suporta exportação, para que você possa trabalhar com seu conjunto de dados em outros aplicativos. A tabela a seguir mostra os formatos de arquivo atualmente suportados para exportação:

Por exemplo, exporte seu conjunto de dados para um arquivo CSV da seguinte forma:

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

Se você tiver um conjunto de dados grande, pode salvar um arquivo por shard, por exemplo:

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