流

数据集流式处理（Dataset streaming）允许你在不下载数据集的情况下对其进行操作。当你遍历数据集时，数据会以流式方式传输。这在以下情况下特别有用：

您不希望等待一个超大数据集的下载。
数据集大小超出了您计算机上可用的磁盘空间。
您想快速浏览数据集中的几个样本。

例如，HuggingFaceFW/fineweb 数据集的英文部分大小为 45 TB，但你可以使用流式处理立即使用它。如下所示，在 load_dataset() 中设置 streaming=True 即可流式加载数据集：

>>> 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 还存储元数据，例如列统计信息（在文件和行组级别），从而实现高效的过滤。使用 datasets.packaged_modules.parquet.ParquetConfig 的 columns 和 filters 参数来流式传输 Parquet 数据集、选择列并应用过滤器：

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

以流式模式加载数据集会创建一个新类型的数据集实例（而非经典的 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 对象，可以使用 to_iterable_dataset() 函数将其转换为 IterableDataset。这实际上比在 load_dataset() 中设置 streaming=True 参数更快，因为数据是从本地文件流式读取的。

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

Shuffle

与普通的 Dataset 对象一样，你也可以使用 IterableDataset.shuffle() 来打乱 IterableDataset。

buffer_size 参数用于控制从中随机采样示例的缓冲区大小。假设你的数据集包含一百万个示例，并将 buffer_size 设置为 10,000。IterableDataset.shuffle() 将从缓冲区中的前 10,000 个示例中随机选取示例。缓冲区中被选中的示例会被新示例替换。默认情况下，缓冲区大小为 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() 每次最多使用 max_buffer_input_shards 个分片（默认为 10）来填充缓冲区，从而显著提升洗牌质量。如果数据集被分片到多个文件中，它还会打乱这些分片的顺序。

如果您正在流式处理仅包含一个或少数几个文件的 Parquet 数据集，请在执行 shuffle 之前使用 reshard()。这有助于按行组（row group）而非按文件对数据集进行分片，从而提升 shuffle 的质量。

重新洗牌

有时您可能希望在每个 epoch 之后重新打乱数据集。这需要您为每个 epoch 设置不同的随机种子。在 epoch 之间使用 IterableDataset.set_epoch() 来告知数据集当前所处的 epoch。

您的种子实际上变为：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 支持分片（sharding），可将非常大的数据集划分为预定义数量的块。在 shard() 中指定 num_shards 参数，以确定将数据集拆分为多少个分片。此外，还需要通过 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)

在这种情况下，拼接后的数据集只有 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，以确保每个样本恰好被查看一次。

为确保使用分片实现正确的并行性，交错数据集的分片至少包含每个输入数据集的 1 个分片。因此，交错数据集的分片级别是输入数据集分片级别的最小值。

例如，如果输入数据集分别有 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')}

类型转换仅在原始特征类型与新特征类型兼容时才能生效。例如，如果原始列仅包含 1 和 0，则可以将特征类型为 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 的 Dataset.map() 函数类似，🤗 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() 来移除列。当你移除一列时，只有在示例被提供给映射函数之后，该列才会被移除。这使得映射函数可以在列被移除之前使用这些列的内容。

在 IterableDataset.map() 中，使用 remove_columns 参数指定要删除的列：

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

请参阅批量映射处理文档中的其他批处理示例。它们对可迭代数据集的用法相同。

筛选

您可以使用 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

batch 方法会将你的 IterableDataset 转换为一个批次迭代器。当你希望在训练循环中处理批次，或使用期望接收批次输入的训练框架时，这一点尤为有用。

使用 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 小于指定的 batch_size，则会将其丢弃。在下游处理要求所有 batch 大小一致的场景中，这非常有用：

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

在训练循环中流式处理

IterableDataset 可以集成到训练循环中。首先，对数据集进行洗牌：

Pytorch

Hide Pytorch content

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

保存数据集检查点并恢复迭代

如果训练循环中断，您可能希望从中断处恢复训练。为此，您可以保存模型、优化器以及数据加载器的检查点。

可迭代数据集无法通过特定示例索引进行随机访问以从中恢复，但你可以使用 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 中。

要从检查点恢复，数据集会跳过之前已读取的所有分片，从当前分片重新开始。然后它会读取该分片，并跳过示例，直到到达检查点中对应的确切示例。

因此，重启数据集的速度相当快，因为它不会重新读取已经迭代过的分片。不过，恢复数据集通常并非瞬间完成，因为它需要从当前分片的开头重新开始读取，并跳过示例，直到到达检查点位置。

这可以与 torchdata 中的 StatefulDataLoader 一起使用：

>>> 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()：恢复时，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	IterableDataset.to_csv()
JSON	IterableDataset.to_json()
Parquet	IterableDataset.to_parquet()
SQL	IterableDataset.to_sql()
内存中的 Python 对象	IterableDataset.to_pandas(), `IterableDataset.to_polars()` or IterableDataset.to_dict()

例如，您可以将数据集导出为 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")

Update on GitHub