Hugging Face的Datasets库：数据加载和处理

1. 总体介绍

在处理NLP模型或微调LLM时，你总是需要一个前置步骤：准备格式正确、足够干净的数据，并能够加载到训练管道中。Hugging Face的datasets库正好解决这个问题——它为你提供了一个统一的API来从Hub加载数据集，使用map/filter进行批量处理，并以优化的形式保存以供反复读取。

读完这篇文章后，你将能够从Hugging Face Hub加载数据集，理解Dataset/DatasetDict的结构，知道如何用tokenizer预处理文本，使用流处理处理大型数据集，以及编写测试来验证数据处理管道。

2. 安装

pip install datasets transformers

datasets库不严格依赖transformers，但第6部分的分词需要它。

import datasets
print(datasets.__version__)

3.0.1

3. 从Hugging Face Hub加载数据集

3.1. 基本的load_dataset

from datasets import load_dataset

reviews = load_dataset("imdb")
print(reviews)

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    unsupervised: Dataset({
        features: ['text', 'label'],
        num_rows: 50000
    })
})

load_dataset从Hub中加载数据集imdb的脚本/元数据，将压缩文件缓存到本地（~/.cache/huggingface/datasets），然后返回一个DatasetDict——即按split（train、test、unsupervised）分类的Dataset字典。

3.2. 仅获取特定的split

train_reviews = load_dataset("imdb", split="train")
print(type(train_reviews), len(train_reviews))

<class 'datasets.arrow_dataset.Dataset'> 25000

传递split="train"直接返回一个Dataset而不是DatasetDict，当你只需要一个集合时很方便。

注意： 如果Hub上的数据集需要config（例如glue），你必须传递config名称：load_dataset("glue", "mrpc")。缺少此参数将抛出错误要求选择config。

4. 探索数据集结构

sample = train_reviews[0]
print(sample.keys())
print(sample["label"], sample["text"][:60])

dict_keys(['text', 'label'])
0 I rented I AM CURIOUS-YELLOW from my video store because of...

通过索引访问返回普通Python dict。Dataset也支持切片，返回列表的dict：

batch = train_reviews[0:3]
print(batch["label"])

[0, 0, 0]

用features查看模式：

print(train_reviews.features)

{'text': Value(dtype='string', id=None),
 'label': ClassLabel(names=['neg', 'pos'], id=None)}

ClassLabel告诉你0对应标签neg，1是pos——当映射回标签名称时很有用。

5. 用map、filter、select处理数据

5.1. filter过滤行

positive_only = train_reviews.filter(lambda row: row["label"] == 1)
print(len(positive_only))

12500

filter接受一个为每一行返回True/False的函数，遍历整个数据集，只保留满足条件的行。

5.2. map转换数据

def add_length(row):
    row["text_length"] = len(row["text"])
    return row

train_with_length = train_reviews.map(add_length)
print(train_with_length[0]["text_length"])

1640

map将函数应用于每一行，返回包含新字段的dict（或更新旧字段）——结果合并回原始数据集。

5.3. map按批处理以加快速度

def add_length_batched(batch):
    batch["text_length"] = [len(t) for t in batch["text"]]
    return batch

train_fast = train_reviews.map(add_length_batched, batched=True, batch_size=1000)

使用batched=True，函数接收整个批量dict（每个字段是list）而不是单行——避免调用Python函数25,000次，对大型数据集的速度显著加快。

5.4. select获取子集

small_train = train_reviews.select(range(100))
print(len(small_train))

100

当你想在运行完整数据集之前在子集上快速测试管道时使用。

6. 用tokenizer预处理文本

最常见的现实场景：对文本进行分词以输入到模型。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_batch(batch):
    return tokenizer(batch["text"], truncation=True, max_length=256)

tokenized = train_reviews.map(tokenize_batch, batched=True)
print(tokenized.column_names)

['text', 'label', 'input_ids', 'attention_mask']

tokenize_batch接收一个批次并返回包含input_ids、attention_mask的dict——map自动将这些新字段合并到数据集中，保持旧字段不变。要删除不必要的字段（例如，原始text以节省RAM），使用remove_columns：

tokenized = train_reviews.map(
    tokenize_batch, batched=True, remove_columns=["text"]
)
print(tokenized.column_names)

['label', 'input_ids', 'attention_mask']

7. 分割split、shuffle、shard

7.1. 自己分割train/test

如果数据集只有一个split，使用train_test_split自己分割：

split_result = train_reviews.train_test_split(test_size=0.2, seed=42)
print(split_result)

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 20000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 5000
    })
})

seed=42确保分割结果在运行之间可重现。

7.2. Shuffle和shard

shuffled = train_reviews.shuffle(seed=42)
part = shuffled.shard(num_shards=4, index=0)
print(len(part))

6250

shard将数据集分成num_shards个相等的部分——当分配数据给分布式训练中的多个worker时很有用。

8. 流处理大型数据集

当数据集太大无法适应磁盘/RAM时，使用streaming=True：

streamed = load_dataset("imdb", split="train", streaming=True)
first_three = list(streamed.take(3))
print(len(first_three))

3

使用流处理，datasets不会将整个文件加载到磁盘，而是通过迭代器顺序读取每条记录。你不能使用len()或直接索引，但仍然可以以懒惰方式使用map/filter（只在迭代到时应用）。

注意： 流处理适合数百GB的数据集，如预训练集。对于小数据集（几百MB），通常正常加载（streaming=False）更快，因为它可以利用缓存和随机访问。

9. 保存和重新加载数据集

处理完成后，保存结果，这样就不必每次都再运行map：

tokenized.save_to_disk("data/imdb_tokenized")

from datasets import load_from_disk

reloaded = load_from_disk("data/imdb_tokenized")
print(len(reloaded))

25000

save_to_disk以Arrow格式写入数据集——重新加载非常快，因为不需要从原始源解析。如果你需要与pandas交互，直接转换：

df = tokenized.select(range(5)).to_pandas()
print(df.columns.tolist())

['label', 'input_ids', 'attention_mask']

10. 为数据处理管道编写测试

数据处理管道应该有测试来确保当你修改代码时，map/filter逻辑不会无声地中断。这是一个使用pytest的测试。

import pytest
from datasets import Dataset


def add_length(row):
    row["text_length"] = len(row["text"])
    return row


def test_add_length_computes_correct_value():
    data = Dataset.from_dict({"text": ["hello", "hi there"]})

    result = data.map(add_length)

    assert result[0]["text_length"] == 5
    assert result[1]["text_length"] == 8


def test_filter_keeps_only_positive_labels():
    data = Dataset.from_dict({"label": [0, 1, 1, 0]})

    positive = data.filter(lambda row: row["label"] == 1)

    assert len(positive) == 2
    assert all(row["label"] == 1 for row in positive)

Dataset.from_dict创建一个小的内存数据集，无需从Hub加载——非常适合单元测试。上面的两个测试检查map（创建新字段）和filter（按条件过滤）的正确行为，无需依赖互联网外的真实数据集。

pytest test_dataset_pipeline.py -v

test_dataset_pipeline.py::test_add_length_computes_correct_value PASSED
test_dataset_pipeline.py::test_filter_keeps_only_positive_labels PASSED

11. 常见错误和何时不应使用

• 忘记在处理大型数据集时使用batched=True：没有批处理的map运行缓慢，因为它为每一行单独调用Python函数。对于数十万行，如果逻辑允许向量化，总是优先使用batched=True。
• 在流处理数据集上使用随机索引：streamed[5]会抛出错误，因为流处理数据集是迭代器，不支持随机访问。仅使用.take()、.skip()或顺序迭代。
• 在shuffle/train_test_split时不设置seed：结果在运行之间不可重现，当比较实验之间的指标时难以调试。
• 缓存失控膨胀：每次用不同的函数调用map时，datasets都会写一个新的缓存文件。如果你的机器磁盘较小，使用load_dataset(..., keep_in_memory=True)或定期删除~/.cache/huggingface/datasets。
• 不应该为复杂的关系数据使用datasets（多个表相互连接，需要事务）。这个库针对ML中使用的平面表/文本数据进行了优化，不是数据库的替代品。

12. 结论

datasets库为你提供了一个干净的API，通过几个load_dataset、map、filter命令就能从“从Hub加载原始数据集”到“数据集已分词，准备好训练”。对于小数据集，正常加载并利用Arrow缓存；对于大型数据集，切换到streaming=True以避免将所有内容加载到磁盘。当为真实项目构建数据处理管道时，别忘了使用Dataset.from_dict为你的map/filter函数编写测试——这有助于你及早发现预处理逻辑何时被无意中更改。