大数据？🤗 Datasets 来拯救！

如今，您经常会遇到处理多千兆字节数据集的情况，特别是如果您计划从头开始预训练像 BERT 或 GPT-2 这样的 Transformer。在这些情况下，即使是加载数据也可能是一个挑战。例如，用于预训练 GPT-2 的 WebText 语料库包含超过 800 万个文档和 40 GB 的文本——将其加载到笔记本电脑的 RAM 中很可能会让它“心脏病发作”！

幸运的是，🤗 Datasets 旨在克服这些限制。它通过将数据集视为内存映射文件，将您从内存管理问题中解放出来；通过流式传输语料库中的条目，将您从硬盘限制中解放出来。

在本节中，我们将使用一个庞大的 825 GB 语料库，称为 the Pile，来探索 🤗 Datasets 的这些功能。让我们开始吧！

什么是 Pile？

Pile 是一个由 EleutherAI 创建的英文文本语料库，用于训练大型语言模型。它包含各种各样的数据集，涵盖科学文章、GitHub 代码库和过滤后的网页文本。训练语料库以 14 GB 分块提供，您还可以下载其中一些单独组件。让我们首先看看 PubMed Abstracts 数据集，这是一个包含来自 PubMed 的 1500 万篇生物医学出版物摘要的语料库。该数据集采用 JSON Lines 格式，并使用 zstandard 库进行压缩，因此我们首先需要安装它

!pip install zstandard

接下来，我们可以使用我们在第 2 节中学到的远程文件加载方法加载数据集

from datasets import load_dataset

# This takes a few minutes to run, so go grab a tea or coffee while you wait :)
data_files = "https://the-eye.eu/public/AI/pile_preliminary_components/PUBMED_title_abstracts_2019_baseline.jsonl.zst"
pubmed_dataset = load_dataset("json", data_files=data_files, split="train")
pubmed_dataset

Dataset({
    features: ['meta', 'text'],
    num_rows: 15518009
})

我们可以看到我们的数据集中有 15,518,009 行和 2 列——这真是太多了！

✎ 默认情况下，🤗 Datasets 会解压加载数据集所需的文件。如果您想节省硬盘空间，可以将 DownloadConfig(delete_extracted=True) 传递给 load_dataset() 的 download_config 参数。有关更多详细信息，请参阅文档。

让我们检查第一个示例的内容

pubmed_dataset[0]

{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

好的，这看起来像一篇医学文章的摘要。现在让我们看看我们已经使用了多少 RAM 来加载数据集！

内存映射的魔力

一种衡量 Python 中内存使用情况的简单方法是使用 psutil 库，可以通过 pip 安装如下：

!pip install psutil

它提供了一个 Process 类，允许我们检查当前进程的内存使用情况，如下所示：

import psutil

# Process.memory_info is expressed in bytes, so convert to megabytes
print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

RAM used: 5678.33 MB

这里的 rss 属性指的是驻留集大小，它是进程在 RAM 中占用的内存比例。这个测量也包括 Python 解释器和我们加载的库所使用的内存，所以加载数据集实际使用的内存会小一些。为了进行比较，让我们使用 dataset_size 属性看看数据集在磁盘上的大小。由于结果与之前一样以字节表示，我们需要手动将其转换为千兆字节：

print(f"Dataset size in bytes: {pubmed_dataset.dataset_size}")
size_gb = pubmed_dataset.dataset_size / (1024**3)
print(f"Dataset size (cache file) : {size_gb:.2f} GB")

Dataset size in bytes : 20979437051
Dataset size (cache file) : 19.54 GB

太棒了——尽管它接近 20 GB，我们仍能以少得多的 RAM 加载和访问数据集！

✏️ 试一试！从 Pile 中选择一个大于您的笔记本电脑或台式机 RAM 的子集，使用 🤗 Datasets 加载它，并测量使用的 RAM 量。请注意，为了获得准确的测量结果，您需要在新进程中执行此操作。您可以在 Pile 论文的表 1 中找到每个子集的解压缩大小。

如果你熟悉 Pandas，这个结果可能会让你感到惊讶，因为 Wes McKinney 有一句著名的经验法则，即你通常需要的数据集大小的 5 到 10 倍的 RAM。那么 🤗 Datasets 是如何解决这个内存管理问题的呢？🤗 Datasets 将每个数据集视为一个内存映射文件，它提供了 RAM 和文件系统存储之间的映射，允许库访问和操作数据集的元素，而无需将其完全加载到内存中。

内存映射文件还可以跨多个进程共享，这使得像 Dataset.map() 这样的方法可以在不需要移动或复制数据集的情况下进行并行化。在底层，所有这些功能都由 Apache Arrow 内存格式和 pyarrow 库实现，这使得数据加载和处理速度极快。（有关 Apache Arrow 和与 Pandas 的比较的更多详细信息，请查看 Dejan Simic 的博客文章。）为了实际了解这一点，让我们通过迭代 PubMed Abstracts 数据集中的所有元素来运行一个小速度测试

import timeit

code_snippet = """batch_size = 1000

for idx in range(0, len(pubmed_dataset), batch_size):
    _ = pubmed_dataset[idx:idx + batch_size]
"""

time = timeit.timeit(stmt=code_snippet, number=1, globals=globals())
print(
    f"Iterated over {len(pubmed_dataset)} examples (about {size_gb:.1f} GB) in "
    f"{time:.1f}s, i.e. {size_gb/time:.3f} GB/s"
)

'Iterated over 15518009 examples (about 19.5 GB) in 64.2s, i.e. 0.304 GB/s'

这里我们使用了 Python 的 timeit 模块来测量 code_snippet 的执行时间。通常，您能够以每秒零点几 GB 到几 GB 的速度迭代数据集。这对于绝大多数应用程序来说都非常有效，但有时您必须处理一个甚至无法存储在笔记本电脑硬盘上的数据集。例如，如果我们尝试下载整个 Pile，我们将需要 825 GB 的可用磁盘空间！为了处理这些情况，🤗 Datasets 提供了一个流式传输功能，允许我们即时下载和访问元素，而无需下载整个数据集。让我们看看它是如何工作的。

💡 在 Jupyter 笔记本中，您还可以使用%%timeit 魔术函数来计时单元格。

流式数据集

要启用数据集流式传输，您只需将 streaming=True 参数传递给 load_dataset() 函数。例如，让我们再次加载 PubMed Abstracts 数据集，但以流式传输模式：

pubmed_dataset_streamed = load_dataset(
    "json", data_files=data_files, split="train", streaming=True
)

与本章其他地方我们熟悉的 Dataset 不同，使用 streaming=True 返回的对象是 IterableDataset。顾名思义，要访问 IterableDataset 的元素，我们需要对其进行迭代。我们可以按如下方式访问流式数据集的第一个元素：

next(iter(pubmed_dataset_streamed))

{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

流式数据集的元素可以使用 IterableDataset.map() 即时处理，这在训练期间（如果您需要对输入进行标记化）很有用。这个过程与我们在第 3 章中标记化数据集的过程完全相同，唯一的区别是输出是一个接一个地返回的

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
tokenized_dataset = pubmed_dataset_streamed.map(lambda x: tokenizer(x["text"]))
next(iter(tokenized_dataset))

{'input_ids': [101, 4958, 5178, 4328, 6779, ...], 'attention_mask': [1, 1, 1, 1, 1, ...]}

💡 为了加快流式分词的速度，您可以传递 batched=True，正如我们在上一节中看到的那样。它将批量处理示例；默认批量大小为 1,000，并且可以使用 batch_size 参数指定。

您还可以使用 IterableDataset.shuffle() 打乱流式数据集，但与 Dataset.shuffle() 不同，这只会在预定义的 buffer_size 中打乱元素

shuffled_dataset = pubmed_dataset_streamed.shuffle(buffer_size=10_000, seed=42)
next(iter(shuffled_dataset))

{'meta': {'pmid': 11410799, 'language': 'eng'},
 'text': 'Randomized study of dose or schedule modification of granulocyte colony-stimulating factor in platinum-based chemotherapy for elderly patients with lung cancer ...'}

在这个例子中，我们从缓冲区中的前 10,000 个例子中随机选择了一个例子。一旦一个例子被访问，它在缓冲区中的位置就会被语料库中的下一个例子填充（例如，在上述情况下是第 10,001 个例子）。你也可以使用 IterableDataset.take() 和 IterableDataset.skip() 函数从流式数据集中选择元素，它们的作用类似于 Dataset.select()。例如，要选择 PubMed Abstracts 数据集中的前 5 个例子，我们可以这样做：

dataset_head = pubmed_dataset_streamed.take(5)
list(dataset_head)

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'pmid': 11409575, 'language': 'eng'},
  'text': 'Clinical signs of hypoxaemia in children with acute lower respiratory infection: indicators of oxygen therapy ...'},
 {'meta': {'pmid': 11409576, 'language': 'eng'},
  'text': "Hypoxaemia in children with severe pneumonia in Papua New Guinea ..."},
 {'meta': {'pmid': 11409577, 'language': 'eng'},
  'text': 'Oxygen concentrators and cylinders ...'},
 {'meta': {'pmid': 11409578, 'language': 'eng'},
  'text': 'Oxygen supply in rural africa: a personal experience ...'}]

类似地，您可以使用 IterableDataset.skip() 函数从打乱的数据集中创建训练和验证拆分，如下所示：

# Skip the first 1,000 examples and include the rest in the training set
train_dataset = shuffled_dataset.skip(1000)
# Take the first 1,000 examples for the validation set
validation_dataset = shuffled_dataset.take(1000)

让我们通过一个常见的应用来结束我们对数据集流式处理的探索：将多个数据集组合在一起创建一个单一的语料库。🤗 Datasets 提供了一个 interleave_datasets() 函数，它将 IterableDataset 对象的列表转换为单个 IterableDataset，新数据集的元素通过在源示例之间交替获得。当您尝试组合大型数据集时，此函数特别有用，因此作为一个示例，让我们流式传输 Pile 的 FreeLaw 子集，这是一个 51 GB 的美国法院法律意见数据集

law_dataset_streamed = load_dataset(
    "json",
    data_files="https://the-eye.eu/public/AI/pile_preliminary_components/FreeLaw_Opinions.jsonl.zst",
    split="train",
    streaming=True,
)
next(iter(law_dataset_streamed))

{'meta': {'case_ID': '110921.json',
  'case_jurisdiction': 'scotus.tar.gz',
  'date_created': '2010-04-28T17:12:49Z'},
 'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}

这个数据集足够大，足以让大多数笔记本电脑的内存吃紧，但我们却能够轻松加载和访问它！现在让我们使用 interleave_datasets() 函数将 FreeLaw 和 PubMed Abstracts 数据集中的示例组合起来

from itertools import islice
from datasets import interleave_datasets

combined_dataset = interleave_datasets([pubmed_dataset_streamed, law_dataset_streamed])
list(islice(combined_dataset, 2))

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'case_ID': '110921.json',
   'case_jurisdiction': 'scotus.tar.gz',
   'date_created': '2010-04-28T17:12:49Z'},
  'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}]

这里我们使用了 Python 的 itertools 模块中的 islice() 函数来选择组合数据集中的前两个示例，我们可以看到它们与两个源数据集中的第一个示例相匹配。

最后，如果您想完整地流式传输 825 GB 的 Pile，您可以按如下方式获取所有准备好的文件：

base_url = "https://the-eye.eu/public/AI/pile/"
data_files = {
    "train": [base_url + "train/" + f"{idx:02d}.jsonl.zst" for idx in range(30)],
    "validation": base_url + "val.jsonl.zst",
    "test": base_url + "test.jsonl.zst",
}
pile_dataset = load_dataset("json", data_files=data_files, streaming=True)
next(iter(pile_dataset["train"]))

{'meta': {'pile_set_name': 'Pile-CC'},
 'text': 'It is done, and submitted. You can play “Survival of the Tastiest” on Android, and on the web...'}

✏️ 试一试！使用大型 Common Crawl 语料库，如 mc4 或 oscar，创建代表您选择的国家/地区语言口语比例的流式多语言数据集。例如，瑞士的四种国语是德语、法语、意大利语和罗曼什语，因此您可以尝试通过根据其口语比例对 Oscar 子集进行抽样来创建瑞士语料库。

您现在拥有加载和处理各种形状和大小数据集所需的所有工具——但除非您特别幸运，否则您的自然语言处理之旅总会遇到需要实际创建数据集来解决手头问题的时候。这就是下一节的主题！

< > 在 GitHub 上更新

LLM 课程

大数据？🤗 Datasets 来拯救！

什么是 Pile？

内存映射的魔力

流式数据集