HuggingfaceNLP笔记5.3Big data? Datasets to the rescue!

如今，处理数GB的数据集已司空见惯，特别是当你计划从头开始预训练BERT或GPT-2这样的变压器时。在这种情况下，甚至加载数据都可能是一个挑战。例如，用于预训练GPT-2的WebText语料库包含超过800万个文档和40GB的文本——将这些内容加载到笔记本电脑的RAM中可能会让其崩溃！

幸运的是，🤗 数据集（🤗 Datasets）设计时就考虑到了这些限制。它通过将数据集视为内存映射文件，解决了内存管理问题，通过流式处理语料库中的条目，解决了硬盘限制。在本节中，我们将使用一个巨大的825GB语料库——Pile，来探索这些功能。

What is the Pile?

Pile是由EleutherAI创建的一个英语文本语料库，用于训练大规模语言模型。它包含了广泛的数据集，涵盖了科学论文、GitHub代码库和过滤后的网络文本。训练语料库分为14GB的块，你也可以下载其中的部分组件。我们先来看看PubMed摘要数据集，它是一个包含1500万份生物医学出版物摘要的语料库，这些出版物来自PubMed。数据集采用JSON Lines格式，并使用zstandard库进行了压缩，首先我们需要安装它：

!pip install zstandard

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

from datasets import load_dataset

# 这可能需要几分钟时间，请在等待时去泡杯茶或咖啡吧 :)
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会自动解压加载数据集所需的文件。如果你想节省硬盘空间，可以在load_dataset()的download_config参数中传入DownloadConfig(delete_extracted=True)。有关更多详细信息，请参阅文档。

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

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

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

The magic of memory mapping

在Python中，一个简单测量内存使用的方法是使用psutil库，它可以通过pip安装：

!pip install psutil

它提供了一个Process类，可以让我们检查当前进程的内存使用情况：

import psutil

# 将Process.memory_info()的值从字节转换为兆字节
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

很好——尽管数据集将近20GB大，但我们仍然能够使用较少的RAM来加载和访问它！

📝 动手试试！ 从《Pile》预处理组件中选择一个大于你笔记本或台式机RAM的子集，使用🤗 Datasets加载它，并测量使用的内存量。为了得到准确的测量结果，你可能需要在新进程中进行。有关每个子集的解压缩大小，可以在《Pile》论文的表1中找到。

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

内存映射文件还可以在多个进程之间共享，这使得Dataset.map()等方法可以并行化，而无需移动或复制数据集。在这些功能的背后，是Apache Arrow内存格式和pyarrow库的支持，它们使得数据加载和处理速度极快。（有关Apache Arrow和与Pandas的比较，可以参考Dejan Simic的博客文章。）现在，让我们运行一个小速度测试，通过遍历PubMed摘要数据集的所有元素来查看效果：

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"遍历了{len(pubmed_dataset)}个示例（大约{size_gb:.1f}GB），耗时{time:.1f}s，即{size_gb/time:.3f}GB/s"
)

'遍历了15518009个示例（大约19.5GB）耗时64.2s，即0.304GB/s'

这里我们使用了Python的timeit模块来测量code_snippet的执行时间。通常情况下，你可以在几十分之一到几GB/s的速度下遍历数据集。这对大多数应用来说已经足够快，但有时你可能需要处理的数据太大，甚至无法存放在笔记本的硬盘上。例如，如果我们尝试下载整个Pile，我们需要825GB的免费磁盘空间！为了解决这种情况，🤗 Datasets提供了流式处理功能，允许我们在不下载整个数据集的情况下下载和访问元素。让我们看看它是如何工作的。

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

Streaming datasets

要启用数据集流式处理，只需在 load_dataset() 函数中传入 streaming=True 参数。例如，我们再次加载 PubMed 摘要数据集，但以流式模式：

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 摘要数据集中选择前 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()函数划分训练集和验证集

# 跳过前1000个例子，其余用于训练集
train_dataset = shuffled_dataset.skip(1000)
# 前1000个例子用于验证集
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 摘要数据集中的示例结合起来。

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()函数，从联合数据集中选择前两个例子，可以看到它们与两个源数据集的前两个例子匹配。

如果你想流式处理整个Pile（825 GB）数据集，可以这样获取准备好的文件：

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)
pile_dataset["train"]中的第一个样本：`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子集中采样，来创建一个瑞士语料库。

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