Поток

Потоковая загрузка датасетов позволяет работать с датасетом без его скачивания. Данные передаются потоком по мере итерации по датасету. Это особенно полезно, когда:

• Не стоит ждать загрузки чрезвычайно большого набора данных.
• Размер набора данных превышает доступный объём дискового пространства на вашем компьютере.
• Вы хотите быстро просмотреть несколько образцов из набора данных.

Например, английская часть набора данных HuggingFaceFW/fineweb занимает 45 терабайт, но с помощью стриминга вы можете начать работу с ней мгновенно. Чтобы стримить набор данных, установите streaming=True в load_dataset(), как показано ниже:

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

Стриминг датасетов также позволяет работать с датасетом, состоящим из локальных файлов, без какого-либо преобразования. В этом случае данные стримятся из локальных файлов по мере итерации по датасету. Это особенно полезно, когда:

• Не стоит ждать, пока огромный локальный набор данных будет преобразован в формат Arrow.
• Размер конвертированных файлов превысит доступное место на диске вашего компьютера.
• Вы хотите быстро просмотреть несколько образцов из набора данных.
• Вы хотите загрузить только определённые столбцы или эффективно отфильтровать набор данных в формате Parquet.

Например, вы можете стримить локальный датасет из сотен сжатых файлов JSONL, таких как oscar-corpus/OSCAR-2201, чтобы использовать его сразу:

>>> 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 — это столбчатый формат, позволяющий потоково загружать только подмножество столбцов и игнорировать ненужные. Parquet также хранит метаданные, такие как статистика столбцов (на уровне файла и группы строк), что обеспечивает эффективную фильтрацию. Используйте аргументы columns и filters класса datasets.packaged_modules.parquet.ParquetConfig, чтобы потоково загружать наборы данных Parquet, выбирать столбцы и применять фильтры:

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, columns=["url", "date"])
>>> print(next(iter(dataset)))
{'url': 'http://%[email protected]/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}

Загрузка набора данных в режиме потоковой обработки создаёт новый экземпляр типа набора данных (вместо классического объекта Dataset), известного как IterableDataset. Этот специальный тип набора данных имеет собственный набор методов обработки, показанных ниже.

Объект IterableDataset полезен для итеративных задач, таких как обучение модели. Не следует использовать IterableDataset для задач, требующих случайного доступа к примерам, поскольку для их обработки необходимо пройтись по всем элементам с помощью цикла for. Чтобы получить последний пример в итерируемом наборе данных, нужно пройти по всем предыдущим примерам. Подробнее см. в руководстве по Dataset и IterableDataset.

Индексация столбцов

Иногда удобно итерировать по значениям конкретного столбца. К счастью, IterableDataset поддерживает индексацию по столбцам:

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

Преобразование из Dataset

Если у вас уже есть объект Dataset, вы можете преобразовать его в IterableDataset с помощью функции to_iterable_dataset(). Это даже быстрее, чем указание аргумента streaming=True в load_dataset(), поскольку данные стримятся из локальных файлов.

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

The to_iterable_dataset() function supports sharding when the IterableDataset is instantiated. This is useful when working with big datasets, and you’d like to shuffle the dataset or to enable fast parallel loading with a 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

Перемешивание

Как и обычный объект Dataset, объект IterableDataset можно перемешать с помощью IterableDataset.shuffle().

Аргумент buffer_size управляет размером буфера, из которого случайно выбираются примеры. Допустим, ваш набор данных содержит один миллион примеров, а вы установили buffer_size в десять тысяч. IterableDataset.shuffle() будет случайным образом выбирать примеры из первых десяти тысяч примеров в буфере. Выбранные примеры в буфере заменяются новыми. По умолчанию размер буфера составляет 1000.

>>> 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() заполняет буфер, используя до max_buffer_input_shards фрагментов за раз (по умолчанию 10), что значительно повышает качество перемешивания. Кроме того, если набор данных разбит на несколько файлов, порядок фрагментов также перемешивается.

Если вы используете потоковую загрузку набора данных Parquet, содержащего один или несколько файлов, вызовите reshard() перед перемешиванием. Это позволяет разбивать набор данных по группам строк, а не по файлам, что повышает качество перемешивания.

Перемешать

Иногда может потребоваться перемешивать датасет после каждой эпохи. Для этого нужно задавать разные значения seed для каждой эпохи. Используйте IterableDataset.set_epoch() между эпохами, чтобы сообщить датасету, какая эпоха выполняется.

Ваш seed фактически становится: initial seed + current epoch.

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

Разделение набора данных

Вы можете разделить набор данных одним из двух способов:

• IterableDataset.take() возвращает первые n примеров из набора данных:

>>> 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() пропускает первые n примеров в наборе данных и возвращает оставшиеся примеры:

>>> train_dataset = shuffled_dataset.skip(1000)

take и skip предотвращают будущие вызовы shuffle, поскольку они фиксируют порядок шардов. Вы должны shuffle набор данных перед его разделением.

Шард

🤗 Datasets поддерживает шардирование для разделения очень большого набора данных на заранее определённое количество фрагментов. Укажите параметр num_shards в shard(), чтобы определить, на сколько шардов разбить набор данных. Также необходимо указать с помощью параметра index, какой именно шард нужно вернуть.

Например, набор данных amazon_polarity содержит 4 фрагмента (в данном случае это 4 файла 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
})

После разбиения набора данных на два фрагмента в первом будет только 2 шарда:

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

Чтобы увеличить количество шардов набора данных, можно использовать IterableDataset.reshard():

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

Механизм ресхардирования зависит от формата файлов датасета. Например, для Parquet ресхардирование выполняется с использованием групп строк, а не одного файла на шард. См. описание работы для каждого формата в документации IterableDataset.reshard().

Если после ресхардирования в вашем наборе данных по-прежнему dataset.num_shards==1, используйте для его разбиения на части методы IterableDataset.skip() и IterableDataset.take().

Соединить

Раздельные наборы данных можно объединить, если у них совпадают типы столбцов. Объедините наборы данных с помощью 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
})

Шарды конкатенированного набора данных — это шарды входных наборов данных.

Также можно объединить два набора данных по горизонтали, установив axis=1, при условии, что у наборов одинаковое количество строк:

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

В этом случае объединённый набор данных содержит только один шард, чтобы избежать несовпадения границ шардов исходных наборов данных.

Чередование

interleave_datasets() может объединить IterableDataset с другими наборами данных, если у них совпадают типы столбцов. Сформированный набор данных возвращает примеры из исходных наборов по очереди.

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

Определите вероятности выборки для каждого из исходных наборов данных, чтобы точнее контролировать, как они выбираются и комбинируются. Укажите желаемые вероятности выборки в аргументе probabilities:

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

Около 80% итогового набора данных составляют es_dataset, а 20% — fr_dataset.

Вы также можете указать stopping_strategy. Стратегия по умолчанию, first_exhausted, является стратегией подвыборки, то есть построение набора данных останавливается, как только один из наборов исчерпывает свои примеры. Вы можете указать stopping_strategy=all_exhausted, чтобы применить стратегию оверсэмплинга. В этом случае построение набора данных останавливается, как только каждый пример из каждого набора был добавлен хотя бы один раз. На практике это означает, что если набор данных исчерпан, он возвращается к началу этого набора, пока не будет достигнут критерий остановки. Обратите внимание, что если вероятности выборки не указаны, новый набор данных будет содержать max_length_datasets*nb_dataset samples. Также существует stopping_strategy=all_exhausted_without_replacement, чтобы гарантировать, что каждый пример будет просмотрен ровно один раз.

Для обеспечения корректного параллелизма с использованием шардирования, шарды перемешанного набора данных содержат как минимум один шард каждого входного набора данных. Следовательно, уровень шардирования перемешанного набора данных равен минимальному уровню шардирования входных наборов данных.

Например, если входные наборы данных содержат соответственно 32, 48 и 128 фрагментов, то чередующийся набор данных будет содержать 32 = min(32, 48, 128) фрагментов, при этом каждый новый фрагмент будет включать 1 фрагмент из первого набора данных, 1–2 фрагмента из второго набора данных и 4 фрагмента из третьего набора данных.

Переименование, удаление и приведение типов

Следующие методы позволяют изменять столбцы набора данных. Эти методы полезны для переименования или удаления столбцов, а также для замены столбцов на новый набор признаков.

Переименовать

Используйте IterableDataset.rename_column(), когда нужно переименовать столбец в наборе данных. Функции, связанные с исходным столбцом, фактически переносятся под новым именем столбца, а не просто заменяют исходный столбец на месте.

Передайте в IterableDataset.rename_column() имя исходной колонки и новое имя колонки:

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

Удалить

Если нужно удалить одну или несколько колонок, передайте в IterableDataset.remove_columns() имя удаляемой колонки. Чтобы удалить несколько колонок, передайте список имён колонок:

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

Приведение

IterableDataset.cast() изменяет тип признаков одной или нескольких колонок. Этот метод принимает в качестве аргумента ваши новые Features. Приведённый ниже пример кода показывает, как изменить типы признаков ClassLabel и 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')}

Преобразование типов работает только в том случае, если исходный и новый типы признаков совместимы. Например, столбец с типом признака Value('int32') можно преобразовать в Value('bool'), если исходный столбец содержит только единицы и нули.

Используйте IterableDataset.cast_column(), чтобы изменить тип признака только одного столбца. Передайте имя столбца и его новый тип признака в качестве аргументов:

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

Карта

Аналогично функции Dataset.map() для обычного Dataset, библиотека 🤗 Datasets предоставляет IterableDataset.map() для обработки IterableDataset. IterableDataset.map() применяет обработку на лету по мере потоковой передачи примеров.

Это позволяет применять функцию обработки к каждому примеру в наборе данных, независимо или партиями. Эта функция даже может создавать новые строки и столбцы.

В следующем примере показано, как токенизировать IterableDataset. Функция должна принимать и возвращать dict:

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

Далее примените эту функцию к набору данных с помощью 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'}]

Рассмотрим ещё один пример, только на этот раз вы удалите столбцы с помощью IterableDataset.map(). Столбец удаляется только после того, как пример был передан в отображаемую функцию. Это позволяет отображаемой функции использовать содержимое столбцов до их удаления.

Укажите столбец для удаления с помощью аргумента remove_columns в 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...'}]

Пакетная обработка

IterableDataset.map() также поддерживает работу с батчами примеров. Чтобы оперировать батчами, установите batched=True. Размер батча по умолчанию составляет 1000, но его можно изменить с помощью аргумента batch_size. Это открывает возможности для многих интересных приложений, таких как токенизация, разбиение длинных предложений на более короткие фрагменты и аугментация данных.

Токенизация

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

См. другие примеры пакетной обработки в документации batched map processing. Они работают так же для итерируемых наборов данных.

Фильтр

Вы можете отфильтровать строки в наборе данных на основе функции-предиката с помощью Dataset.filter(). Она возвращает строки, соответствующие заданному условию:

>>> 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() также позволяет фильтровать по индексам, если установить 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...}]

Пакет

Метод batch преобразует ваш IterableDataset в итерируемый объект, состоящий из батчей. Это особенно полезно, когда вы хотите работать с батчами в цикле обучения или используете фреймворки, ожидающие батчевые входные данные.

Также существует вариант «Batch Processing» при использовании функции map для применения функции к батчам данных, о котором рассказано в разделе Map выше. Метод batch, описанный здесь, отличается и предоставляет более прямой способ создания батчей из вашего набора данных.

Метод batch можно использовать следующим образом:

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

В этом примере batched_dataset по-прежнему является IterableDataset, но теперь каждый возвращаемый элемент представляет собой батч из 32 образцов, а не один образец. Такая батчевка выполняется на лету при итерации по набору данных, что сохраняет свойство IterableDataset по экономному использованию памяти.

Метод batch также предоставляет параметр drop_last_batch. Если установить его в True, последний батч будет отброшен, если его размер меньше указанного batch_size. Это может быть полезно в сценариях, где ваша последующая обработка требует, чтобы все батчи были одинакового размера:

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

Потоковая передача в цикле обучения

IterableDataset можно интегрировать в цикл обучения. Сначала перемешайте набор данных:

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

Сохранение контрольной точки набора данных и возобновление итерации

Если цикл обучения останавливается, вы можете перезапустить обучение с того места, где оно было прервано. Для этого сохраните контрольную точку (checkpoint) вашей модели и оптимизаторов, а также загрузчика данных.

Наборы данных типа Iterable не предоставляют случайный доступ к конкретному индексу примера для возобновления работы, но вы можете использовать IterableDataset.state_dict() и IterableDataset.load_state_dict(), чтобы возобновить работу с контрольной точки, аналогично тому, как это делается для моделей и оптимизаторов:

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

Возвращает:

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

Под капотом итерируемый набор данных отслеживает текущий шард, из которого выполняется чтение, и индекс примера в текущем шарде, сохраняя эту информацию в state_dict.

Чтобы возобновить работу с контрольной точки, набор данных пропускает все сегменты, которые были прочитаны ранее, чтобы начать с текущего сегмента. Затем он читает сегмент и пропускает примеры, пока не достигнет точного примера из контрольной точки.

Поэтому перезапуск набора данных происходит достаточно быстро, так как он не перечитывает уже обработанные фрагменты. Тем не менее возобновление работы с набором данных обычно не происходит мгновенно, поскольку необходимо начать чтение с начала текущего фрагмента и пропустить примеры, пока не будет достигнута позиция контрольной точки.

Это можно использовать вместе с StatefulDataLoader из 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

При возобновлении выполнение продолжается ровно с того места, где была сохранена контрольная точка, за исключением случая использования .shuffle(): примеры из буферов перемешивания теряются при возобновлении, а буферы заполняются новыми данными.

Сохранить

Когда ваш итерируемый набор данных готов, вы можете сохранить его в формате Parquet как Hugging Face Dataset и использовать позже с помощью load_dataset().

Сохраните набор данных, передав в качестве аргумента push_to_hub() имя репозитория набора данных на Hugging Face, в который вы хотите его сохранить. Этот метод перебирает набор данных и последовательно загружает данные на Hugging Face:

dataset.push_to_hub("username/my_dataset")

По умолчанию IterableDataset.push_to_hub() записывает dataset.num_shards файлов. Установите num_shards, чтобы выбрать другое количество. Если количество выходных шардов больше, чем dataset.num_shards, каждый входной шард распределяется по своей группе выходных шардов путём записи примеров в них по кругу, за один проход по данным. Если примеров меньше, чем запрошено шардов, оставшиеся шарды будут пустыми файлами Parquet.

Если набор данных состоит из нескольких фрагментов (dataset.num_shards > 1), вы можете использовать несколько процессов для параллельной загрузки. Это особенно полезно, если вы применяли шаги map() или filter(), поскольку они выполняются быстрее при параллельной обработке:

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

Используйте функцию load_dataset() для повторной загрузки набора данных:

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

Экспорт

🤗 Datasets поддерживает экспорт, что позволяет работать с вашим набором данных в других приложениях. В таблице ниже перечислены поддерживаемые на данный момент форматы файлов для экспорта:

Например, экспортируйте ваш набор данных в файл CSV следующим образом:

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

Если у вас большой набор данных, вы можете сохранить по одному файлу на каждый шард, например

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