Hugging Face Datasets 流式读取大语料时如何切分样本

大语料不适合先完整下载再切片时，Hugging Face Datasets 的 streaming=True 可以边迭代边读取。但它返回的不是支持随机下标的普通 Dataset，而是 IterableDataset。所以“切分样本”不能简单理解成 dataset[1000:2000]：连续取一段用 take/skip，并行分给多个 worker 则优先用 shard。

先确认流式数据能做什么

下面的代码只展示读取边界，不把它当成本机执行截图。开启流式模式后，数据会随着 for 迭代逐步到达；不能用 len() 或随机下标来获得任意一条记录。这个特性适合训练、批处理和抽样预览，不适合频繁跳到某个绝对行号。

from datasets import load_dataset

# 只在迭代时读取数据，不预先把整个语料落到本地磁盘。
dataset = load_dataset(
    "allenai/c4",
    "en",
    split="train",
    streaming=True,
)

# 取一条记录查看字段；这里没有随机访问语义。
first = next(iter(dataset))
print(first.keys())

连续样本段用 shuffle、take 和 skip

如果需求是“先拿前 500 条做冒烟测试”或“跳过前 10 万条再处理”，可以把数据集看成一个有顺序的迭代器。需要随机顺序时，要在切分之前调用 shuffle：

# 先打乱 shard 顺序和缓冲区中的样本，再固定两个连续窗口。
shuffled = dataset.shuffle(seed=42, buffer_size=10_000)
preview = shuffled.take(500)
training_part = shuffled.skip(500).take(10_000)

for row in training_part:
    # 这里处理一个样本；生产代码可在此处做分词或写入批缓冲。
    text = row.get("text", "")
    if text:
        consume(text)

take 只保留开头的 N 条，skip 返回剩余部分。后者不会瞬移到第 N 条，开始迭代时仍要先走过被跳过的数据，因此跳过很大的偏移量会带来启动延迟。另一个常见坑是把 take/skip 放在 shuffle 前面：切分操作会锁定顺序，后续再调整 shard 顺序就不是你以为的随机切片。

并行任务用 shard 切互斥分片

训练或批处理需要多个 worker 同时消费时，按 worker 编号取分片更合适。num_shards 是总分片数，index 是当前 worker 的分片编号，编号从 0 开始：

def read_worker(dataset, worker_id, worker_count):
    # 固定总分片数和当前编号，让不同 worker 读取不同底层 shard。
    part = dataset.shard(num_shards=worker_count, index=worker_id)
    for row in part:
        # 每个 worker 只处理自己的分片，避免简单复制迭代器造成重复。
        yield row


for item in read_worker(dataset, worker_id=1, worker_count=4):
    # 这里接入当前 worker 的预处理逻辑。
    consume(item)

这种方式解决的是“谁负责哪些底层 shard”，不是精确的全局行号分页。若底层只有一个 shard，文档建议改用 skip/take 做窗口；如果是从已有 Dataset 转成流式数据，可用 to_iterable_dataset(num_shards=64) 预先增加可分配的 shard，再交给 PyTorch DataLoader 的 worker。

按目标选择切分方式

实际接入 DataLoader 时，还要确认 worker 数量、分片数和样本处理是否幂等。需要恢复中断位置时，不要保存一个“第几行”的假设；IterableDataset 提供 state_dict() 和 load_state_dict()，可以记录当前 shard 与 shard 内位置。

常见问题

take 和 skip 能否替代普通 Dataset 的切片？

只能替代顺序窗口，不能提供随机访问。尤其是 skip，大偏移量仍然需要逐条迭代过去。

先 shard 再 shuffle 可以吗？

如果目标是可预期的随机切分，先 shuffle 再 shard 或 take/skip；切分后再 shuffle 会受到顺序锁定语义影响。

shard 能保证每个 worker 获得完全相同数量的样本吗？

它按底层 shard 分配，不承诺按全局样本数精确均分。最后一个分片可能更短，业务侧应允许批次大小变化。

参考资料：https://huggingface.co/docs/datasets/en/stream；https://huggingface.co/docs/datasets/en/package_reference/loading_methods。