使用HuggingFace Datasets库高效处理大规模数据集-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00388/article/details/148549181

使用HuggingFace Datasets库高效处理大规模数据集

notebooks Notebooks using the Hugging Face libraries 🤗 项目地址: https://gitcode.com/gh_mirrors/note/notebooks

引言：大数据处理的挑战

在自然语言处理(NLP)领域，研究人员和工程师经常需要处理海量文本数据。传统的数据加载方式往往需要将整个数据集一次性加载到内存中，这对于GB甚至TB级别的数据集来说几乎是不可能的任务。本文将介绍如何使用HuggingFace Datasets库来高效处理大规模数据集，而无需担心内存限制。

准备工作：安装必要库

在开始之前，我们需要安装一些必要的Python库：

pip install datasets evaluate transformers[sentencepiece]
pip install zstandard
pip install psutil

这些库包括：

datasets: HuggingFace提供的数据集处理库
transformers: HuggingFace的NLP模型库
zstandard: 用于处理压缩数据
psutil: 用于监控内存使用情况

加载大规模数据集

让我们从加载一个真实的医学文献数据集开始，这个数据集包含超过1500万篇PubMed文章的标题和摘要：

from datasets import load_dataset

data_files = "PUBMED_title_abstracts_2020_baseline.jsonl.zst"
pubmed_dataset = load_dataset("json", data_files=data_files, split="train")

查看数据集的基本信息：

pubmed_dataset

输出显示这个数据集包含15,518,009条记录，每条记录包含'meta'和'text'两个字段。我们可以查看第一条记录的内容：

pubmed_dataset[0]

内存使用分析

处理大数据时，内存管理至关重要。让我们看看加载这个数据集占用了多少内存：

import psutil

print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

令人惊讶的是，尽管数据集原始大小约为19.5GB，但内存占用却远小于这个数字：

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

这是因为Datasets库采用了智能的内存映射技术，不会一次性加载全部数据到内存中。

流式处理模式

对于真正的大规模数据集，我们可以使用流式处理模式(streaming mode)，这种方式几乎不占用内存：

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

流式模式下，我们可以逐个访问数据样本：

next(iter(pubmed_dataset_streamed))

数据处理管道

在流式模式下，我们可以构建高效的数据处理管道：

分词处理：使用预训练的分词器处理文本

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

数据洗牌：虽然数据是流式的，我们仍然可以进行洗牌操作

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

数据分割：获取数据集的前N个样本或跳过前N个样本

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

# 创建训练集和验证集
train_dataset = shuffled_dataset.skip(1000)
validation_dataset = shuffled_dataset.take(1000)

多数据集混合处理

Datasets库还支持混合多个流式数据集：

law_dataset_streamed = load_dataset(
    "json",
    data_files="FreeLaw_Opinions.jsonl.zst",
    split="train",
    streaming=True,
)

from itertools import islice
from datasets import interleave_datasets

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

处理更复杂的数据集结构

对于具有训练集、验证集和测试集分割的大型数据集，我们可以这样加载：

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

性能优化技巧

批量处理：通过批量处理数据可以提高处理速度

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