00240 NLP Course - Big data? 🤗 Datasets to the rescue!

前言

如今，发现自己正在处理多吉字节数据集的情况并不少见，特别是如果你计划从头开始预训练像 BERT 或 GPT-2 这样的变压器模型。在这些情况下，即使是加载数据也可能是一个挑战。例如，用于预训练 GPT-2 的 WebText 语料库包含超过800万个文档和40GB的文本——将这些数据加载到你的笔记本电脑的内存中，很可能会让它崩溃！

幸运的是，🤗 Datasets 被设计用来克服这些限制。它通过将数据集视为内存映射文件，让你摆脱内存管理问题，并通过流式传输语料库中的条目，让你摆脱硬盘限制。

在这一节中，我们将通过一个名为 Pile 的巨大 825 GB 语料库来探索 🤗 Datasets 的这些特性。让我们开始吧！

src link: https://huggingface.co/learn/nlp-course/chapter5/4

Operating System: Ubuntu 22.04.4 LTS

参考文档

1. NLP Course - Big data? 🤗 Datasets to the rescue!

什么是 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 ...'}

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

内存映射的魔力

在 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"Number of files in dataset : {pubmed_dataset.dataset_size}")
size_gb = pubmed_dataset.dataset_size / (1024**3)
print(f"Dataset size (cache file) : {size_gb:.2f} GB")

Number of files in dataset : 20979437051
Dataset size (cache file) : 19.54 GB

太好了——尽管数据集几乎有 20 GB 大，我们还是能够使用更少的内存来加载和访问数据集！

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

如果你熟悉 Pandas，这个结果可能会让你感到惊讶，因为 Wes Kinney 有一个著名的经验法则，即你通常需要比数据集大小多 5 到 10 倍的内存。那么 🤗 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 的执行时间。通常，你能够以每秒几百兆字节到几吉字节的速度遍历一个数据集。这对于绝大多数应用程序来说都很好，但有时你需要处理一个太大以至于无法存储在你笔记本电脑硬盘上的数据集。例如，如果我们尝试下载整个 Pile，我们将需要 825 GB 的可用磁盘空间！为了处理这些情况，🤗 Datasets 提供了一个流式处理功能，允许我们即时下载和访问元素，而无需下载整个数据集。让我们来看看这是如何工作的。

💡 在 Jupyter 笔记本中，你还可以使用 %%timeit 魔法函数来计算单元格的执行时间。

流式数据集

要启用数据集流式处理，你只需要向 load_dataset() 函数传递 streaming=True 参数。例如，让我们再次加载 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...'}

✏️ 尝试一下！使用像 mc4 或 oscar 这样的大型 Common Crawl 语料库来创建一个流式多语言数据集，代表你选择的国家中语言的使用比例。例如，瑞士的四种国家语言是德语、法语、意大利语和罗曼什语，所以你可以尝试根据它们的使用比例对 Oscar 子集进行采样，以创建一个瑞士语料库。

现在，你已经拥有了加载和处理各种形状和大小的数据集所需的所有工具——但除非你非常幸运，否则在你的 NLP 旅程中总会有一个时刻，你需要实际创建一个数据集来解决手头的问题。这就是下一节的主题！

结语

第二百四十篇博文写完，开心！！！！

今天，也是充满希望的一天。