如何在流式模式下拆分Hugging Face数据集且不加载全量至内存？

当然有可行方案，流式模式下拆分的核心是惰性分流，无需将全量数据加载到内存，下面提供两种实用方法：

方法1：基于哈希值的确定性分流（适合无预设拆分的数据集） #

流式数据集是IterableDataset，无法直接使用train_test_split（该方法需要加载全量数据计算长度），但可以通过样本的唯一标识（如官方分配的id）生成哈希值，根据哈希结果将样本惰性分配到训练/验证集，整个过程仅在迭代时处理当前样本。

代码示例：

from datasets import load_dataset
import hashlib

# 加载流式数据集
dataset = load_dataset('squad', split='train', streaming=True)

def split_streaming_dataset(streaming_dataset, test_size=0.1):
    # 基于样本id生成哈希值，确保同一样本始终分到同一组
    def is_validation_sample(example):
        hash_obj = hashlib.md5(str(example['id']).encode('utf-8'))
        hash_int = int(hash_obj.hexdigest(), 16)
        # 通过取模控制验证集占比
        return hash_int % int(1 / test_size) == 0
    
    # 惰性过滤出训练/验证集
    val_dataset = streaming_dataset.filter(is_validation_sample)
    train_dataset = streaming_dataset.filter(lambda x: not is_validation_sample(x))
    
    return train_dataset, val_dataset

# 执行拆分（验证集占10%）
train_dataset, val_dataset = split_streaming_dataset(dataset, test_size=0.1)

# 验证效果（仅取少量样本查看）
print("训练集前5个样本id:")
for idx, example in enumerate(train_dataset.take(5)):
    print(f"{idx+1}: {example['id']}")

print("\n验证集前5个样本id:")
for idx, example in enumerate(val_dataset.take(5)):
    print(f"{idx+1}: {example['id']}")

方法2：直接加载官方预设拆分（最简单的场景） #

多数Hugging Face官方数据集已经预先划分好了训练/验证集，直接指定split参数加载流式版本即可：

代码示例：

from datasets import load_dataset

# 直接加载流式训练集和验证集
train_dataset = load_dataset('squad', split='train', streaming=True)
val_dataset = load_dataset('squad', split='validation', streaming=True)

注意事项 #

• 流式模式下无法使用len()获取数据集总大小（数据是按需加载的），若需统计数量，只能迭代计数，但这会遍历全量数据，失去流式的内存优势。
• 哈希分流时，务必使用样本的唯一且稳定的字段（如官方id），避免用随机数导致每次分配结果不一致。
• 调整验证集占比时，直接修改test_size参数即可，例如test_size=0.2表示20%数据进入验证集。

内容的提问来源于stack exchange，提问作者Charlie Parker