跳到主要内容

数据截至 (上游 commit cd6a3406572e)

HF Datasets — 架构与原理

30 秒导读: 🤗 Datasets 是 Hugging Face 的训练数据加载与预处理库。它做两件事:一是 load_dataset("org/name") 一行把 Hub 上(或本地)的数据变成可直接喂给训练循环的对象;二是用 map/filter/shuffle 做大规模预处理,结果按"函数指纹"缓存在磁盘上,第二次运行直接复用。底层靠 Apache Arrow 的内存映射,所以语料比 RAM 大也能工作;另有一套 IterableDataset 支持不下载、边下边读的流式训练。

1. 这是什么(零基础也能懂)​

一句话定义​

HF Datasets 是一个把"数据集"变成带类型的、可随机访问、可流式迭代的表的 Python 库:底层是一张 Apache Arrow 列式表,上面包了一层缓存、指纹、格式转换的 API。

它要解决谁的什么问题​

假设你要训练一个模型,数据是 500 GB 的 JSONL 语料。你会立刻撞上三堵墙:

  • 装不下:全读进内存直接爆掉。
  • 预处理贵:分词、清洗一遍要几小时;改一行代码重跑脚本,不想再等几小时。
  • 下载慢:有时候你只想边下边训,不想先把 500 GB 落盘。

HF Datasets 的三件套正好对应这三堵墙:Arrow 内存映射(数据留在磁盘,按需换页)、指纹缓存(预处理结果按内容寻址,重复运行直接命中)、流式 IterableDataset(远程文件逐块读)。

它能做什么​

能力入口
一行加载 Hub / 本地数据集(CSV、JSON、Parquet、Arrow、文本、imagefolder 等)load_dataset()(src/datasets/load.py:1480)
惰性流式读取,不下载全量load_dataset(..., streaming=True)
大规模预处理:map / filter / shuffle / select / sort / train_test_splitDataset 的方法(src/datasets/arrow_dataset.py:718)
多进程预处理map(num_proc=N)
预处理结果磁盘缓存,重跑零开销自动(指纹机制,src/datasets/fingerprint.py)
多数据源按比例混合interleave_datasets()(src/datasets/combine.py:18)
转 NumPy / Pandas / PyTorch / TensorFlow / JAX / Polars 格式set_format()
存取自有格式save_to_disk() / load_from_disk() / push_to_hub()
检索索引add_faiss_index() / add_elasticsearch_index()(src/datasets/search.py)

用起来什么样​

# 典型用法(摘自 README 与 docstring 示例)
from datasets import load_dataset

ds = load_dataset("rajpurkar/squad", split="train")   # 下载+转成 Arrow+内存映射打开
ds = ds.map(lambda ex: {"len": len(ex["question"])})  # 预处理,结果写进磁盘缓存
ds = ds.shuffle(seed=42).select(range(1000))          # 打乱+选子集(只写索引映射)
ds.set_format("torch")                                 # 取出时自动转 torch.Tensor

for batch in ds:                                       # 喂给训练循环
    ...

# 流式:不下载,边下边读
ds = load_dataset("allenai/c4", "en", split="train", streaming=True)
for example in ds.shuffle(seed=42, buffer_size=10_000):
    ...

一句话直觉​

把 HF Datasets 想成"数据集世界的 Git + 虚拟内存"。 磁盘上的 Arrow 文件是对象库,内存映射让它像内存一样被随机访问;每次 map/filter 相当于一次 commit——结果文件以"指纹"为名,指纹由输入数据 + 变换函数的内容决定,同样的输入同样的函数永远命中同一个缓存文件。

2. 顶层全景(它大概怎么转)​

2.1 一张图看全局​

怎么读这张图: 自上而下是 load_dataset() 一次调用的四个阶段;①② 只有第一次跑才做重活,③④ 是之后每次运行的常态路径。

load_dataset("rajpurkar/squad", split="train")
        │
        ▼
┌──────────────────────────────┐
│ ① 路由:工厂链                 │  load.py · dataset_module_factory
│ 本地目录? Hub 仓库? 什么格式?    │  → 选中一个打包好的格式构建器
└──────────────┬───────────────┘     (parquet / json / csv / ...)
               ▼
┌──────────────────────────────┐
│ ② 构建:下载 + 写成 Arrow      │  builder.py · download_and_prepare
│ 缓存已存在则整段跳过           │  → <cache>/名字/配置/版本/哈希/*.arrow
└──────────────┬───────────────┘
               ▼
┌──────────────────────────────┐
│ ③ 打开:内存映射,不进 RAM     │  table.py · MemoryMappedTable
│ Dataset._data 指向磁盘文件     │  → 随机访问由 OS 换页支撑
└──────────────┬───────────────┘
               ▼
┌──────────────────────────────┐
│ ④ 变换:map / filter / shuffle │  arrow_dataset.py + fingerprint.py
│ 结果以"指纹"命名缓存           │  → cache-<新指纹>.arrow,下次直接命中
└──────────────────────────────┘

旁路:streaming=True → IterableDataset(迭代器树),跳过 ②③,不落地

2.2 部件一句话职责​

部件干什么在哪个文件
load_dataset总入口:路由 → 构建 → 打开src/datasets/load.py:1480
dataset_module_factory判断"数据在哪、什么格式",返回构建器模块src/datasets/load.py:966
DatasetBuilder下载原始文件、生成 Arrow 缓存目录src/datasets/builder.py:214
ArrowWriter把例子成批写成 Arrow IPC 流文件src/datasets/arrow_writer.py:487
ArrowReader / ReadInstruction按 split 指令(如 train[:10%])读缓存文件src/datasets/arrow_reader.py:285、:456
Datasetmap 式数据集;_data(主表)+ _indices(可选索引映射)src/datasets/arrow_dataset.py:718
MemoryMappedTable磁盘 Arrow 文件的内存映射视图,pickle 时只存路径+变换历史src/datasets/table.py:1055
ConcatenationTable多张表的逻辑拼接(行/列两个方向),不复制数据src/datasets/table.py:1339
Hasher / update_fingerprintxxhash + dill 给"函数+参数"算指纹src/datasets/fingerprint.py:196、:253
IterableDataset流式数据集,包着一棵惰性迭代器树src/datasets/iterable_dataset.py:2503
_BaseExamplesIterable迭代器树节点基类(map/filter/shuffle 各是一种节点)src/datasets/iterable_dataset.py:202
Features列类型 schema(文本、类别、图像、音频……)src/datasets/features/features.py:1863
DataFilesDict把 data/train-*.parquet 这类模式解析成"哪个 split 用哪些文件"src/datasets/data_files.py:648

2.3 主线走一遍​

走 load_dataset("rajpurkar/squad", split="train") 这条路:

  1. 路由(src/datasets/load.py:966 dataset_module_factory):path 不是本地目录、不是打包格式名,形如 org/name → 走 Hub 路径,先下载 README(数据集卡)并把 revision 钉死成 commit hash(src/datasets/load.py:1121-1130),按仓库里的数据文件后缀选中 Parquet 构建器。
  2. 构建(src/datasets/builder.py:702 download_and_prepare):先查 <cache>/squad/配置/版本/哈希/ 下有没有 dataset_info.json——有就直接返回(第二次运行的全部秘密在 :842-848 这六行);没有就下载 Parquet、逐 split 写成 Arrow,整个过程在 <目录>.incomplete 临时目录里做,成功后一次性 rename(:859-876)。
  3. 打开(src/datasets/builder.py:1076 _as_dataset → src/datasets/arrow_reader.py:317 read_table):默认 MemoryMappedTable.from_file——pa.memory_map 打开文件,数据一页都不进 RAM,直到你真的去读某一行。
  4. 变换:之后每次 ds.map(fn) 都会算出新指纹,写成 cache-<指纹>.arrow;再跑一遍同样的 fn 直接 Dataset.from_file 打开旧文件(src/datasets/arrow_dataset.py:3461-3470)。

这条线最值得记住的一点: ② 的产物(版本化缓存目录)和 ④ 的产物(指纹命名缓存文件)是两套独立的缓存——前者按"数据集身份"寻址,后者按"计算内容"寻址。第 2 章和第 4 章分别讲它们。

3. 阅读地图(建议顺序)​

四章由浅入深。时间有限就读 01 → 02——Arrow 内存映射和指纹缓存是这个库区别于"普通 pandas 封装"的全部要害。

顺序章节讲什么适合谁
101-arrow-storage.md为什么数据比内存大也能随机访问:memory-map、replay 机制、ConcatenationTable所有人必读,这是地基
202-map-cache-fingerprint.md预处理缓存怎么做到"重跑零开销":指纹、缓存命中、多进程、索引映射天天用 map/filter 的人
303-iterable-streaming.md流式训练:迭代器树、两级 shuffle、断点恢复、分布式切分语料太大不想下载的人
404-hub-loading.md一行 load_dataset 背后的工厂链、版本化缓存目录、数据文件模式匹配想搞清"数据集怎么被找到"的人

4. 巧妙之处(可借鉴的技术)​

先给速览,每条在对应章节有详细展开:

  1. pickle 只存"路径 + 操作历史"。 MemoryMappedTable.__getstate__ 返回的是 {"path": ..., "replays": [...]},反序列化时重新 mmap 文件并把 slice/filter/cast 这些操作"重放"一遍(src/datasets/table.py:1087-1107)。数据永远不被复制进 pickle 流——这就是多进程 map 传 Dataset 不爆内存的原因。→ 第 1 章
  2. 把用户函数本身哈希进缓存键。 Hasher 用 dill 把函数(含字节码)序列化后喂给 xxhash(src/datasets/fingerprint.py:196-223),于是"改了函数"自动等于"换了缓存键",旧缓存自然失效,不需要任何手动版本号。→ 第 2 章
  3. shuffle 不搬数据,只写一列 uint64 索引。 shuffle/select/filter 的产物是一张单列索引表盖在原表上(src/datasets/arrow_dataset.py:4586-4591),打乱 1 亿行只写 800 MB 索引而不碰原始数据;代价是随机读变慢,官方文档自己标注"up to 10x slower"(:4851-4855 的 docstring),并给出 flatten_indices 这条"落盘重写"的回头路。→ 第 2 章
  4. shuffle 拆成两级:shard 顺序 + 固定大小 buffer。 IterableDataset.shuffle 先打乱数据源顺序、再交错最多 10 个 shard 填充、最后用 buffer_size 的内存 buffer 做近似打乱(src/datasets/iterable_dataset.py:3738-3826),内存占用与语料大小无关,这是 TB 级语料也能"洗开"的关键。→ 第 3 章
  5. 缓存写入全部走"临时文件 + 原子改名"。 map 的每个分片先写 NamedTemporaryFile,成功后 shutil.move 到缓存路径(src/datasets/arrow_dataset.py:4044-4049);builder 整目录同理(.incomplete 目录)。中途 Ctrl-C 永远不会留下半个缓存文件被下次误命中。→ 第 2、4 章
  6. 数据集卡 README 兼任配置文件。 Hub 仓库的 README YAML(configs: 字段)直接参与决定"哪些文件进哪个 split、用什么构建器"(src/datasets/load.py:585-620),元数据和数据在同一个 git 仓库里同版本演化。→ 第 4 章

5. 边界与局限​

诚实清单,细节散在各章:

  • 它不是数据清洗框架。 没有内置去重、PII、质量打分这些语料级流水线;预训练规模的清洗去看 datatrove 和 dolma。
  • 指纹机制有失效面。 函数必须能被 dill 序列化,否则指纹退化为随机值、缓存永不命中,只警告一次(src/datasets/fingerprint.py:259-276);函数引用的全局变量不进指纹,改了全局变量缓存不会失效(inferred:指纹只 hash 函数自身与显式参数,见第 2 章)。
  • indices mapping 是甜点也是陷阱。 shuffle/select 之后随机访问慢一个量级,训练前要么 flatten_indices 重写,要么干脆换 IterableDataset。
  • 流式模式没有随机访问、没有精确长度。 IterableDataset 不支持 ds[i],shuffle 是近似的;skip/take 之后 shard 顺序被冻结,shuffle 质量下降(src/datasets/iterable_dataset.py:2008-2009)。
  • 多进程 map 有进程开销与序列化成本。 num_proc 对秒级函数不值;且分片数会跟着缓存时的 num_proc 走,换并行度可能重算(src/datasets/arrow_dataset.py:3488-3502)。
  • 数据集脚本时代已结束。 2024 年起 Hub 上的 Python 加载脚本被移除,load_dataset 遇到 .py 直接报错(src/datasets/load.py:1066-1069)——只认数据文件 + README YAML。

6. 横向对比​

同书架(training-data 一带)的兄弟项目,取舍各不相同:

库定位与 HF Datasets 的关键差异
HF Datasets(本库)通用加载 + 预处理 + 分发单机/单机多进程为主;缓存按"函数指纹";Hub 生态是护城河
datatrove预训练语料清洗流水线面向"几 TB 网页文本过滤去重",流水线算子 + 分布式执行,不做随机访问
dolmaOLMo 的预训练数据工具集围绕"产出 Dolma 语料"的固定工具链(tag/mix/decontaminate),不是通用加载器
olmo完整预训练代码数据侧是训练 loop 的 DataLoader,用 HF Datasets 之外的自有格式(MMAP/JSONL.gz)

一句话:HF Datasets 是"进训练循环之前"的最后一公里;语料规模大到要先做工业级清洗时,上游接 datatrove/dolma,产出落回 Parquet/Hub,再回到本库做加载与 shuffle。

7. 代码地图(入口级)​

每章末尾有自己的细粒度地图,这里只列从零开始读源码的五个入口:

主题文件路径符号名
总入口src/datasets/load.pyload_dataset、load_dataset_builder、dataset_module_factory
map 式数据集主体src/datasets/arrow_dataset.pyDataset、Dataset.map、Dataset._map_single
Arrow 存储层src/datasets/table.pyMemoryMappedTable、InMemoryTable、ConcatenationTable、IndexedTableMixin
指纹与缓存src/datasets/fingerprint.pyHasher、update_fingerprint、fingerprint_transform
流式数据集src/datasets/iterable_dataset.pyIterableDataset、_BaseExamplesIterable、BufferShuffledExamplesIterable