Hugging Face Datasets interleave_datasets 怎样按概率交错混合语料

多个语料源需要混成一条训练流时，datasets.interleave_datasets 可以按来源概率逐条选择样本。关键是把 probabilities、seed 和 stopping_strategy 分开理解：前者控制“下一条从谁取”，seed 控制随机序列是否可复现，停止策略则决定某个来源耗尽后是否结束或允许继续采样。

官方地址：https://huggingface.co/docs/datasets/en/package_reference/main_classes

想要可复现的按比例混合，先把概率列表与数据集列表一一对应，再固定 seed；默认策略遇到任一来源耗尽就停止，允许过采样时才改用 all_exhausted。

先确认概率列表代表什么

interleave_datasets 接收一个 Dataset 或 IterableDataset 列表。不给 probabilities 时，函数按来源轮换；传入概率后，每次从一个来源随机取一条。概率列表的顺序必须和数据集列表一致，且通常应表达相对权重，例如 [0.7, 0.2, 0.1] 表示 A 的选择机会高于 B 和 C，并不承诺最终输出中三类样本严格按 7:2:1 出现。

seed 只负责让随机选择序列可重复。它不会改变概率，也不会把不等长的来源自动补齐。检查时先问两个问题：概率是否与来源位置对应，下一次运行是否使用同一个 seed。

用最小示例观察交错结果

下面用三个很短的内存数据集演示按概率取样。示例故意固定 seed，并打印来源标签，方便把“选择来源”与“最终数量”区分开。

from datasets import Dataset, interleave_datasets

# 三个来源只放少量标签，便于观察交错顺序，而不是模拟真实语料规模。
news = Dataset.from_dict({"source": ["news-1", "news-2", "news-3"]})
code = Dataset.from_dict({"source": ["code-1", "code-2", "code-3"]})
qa = Dataset.from_dict({"source": ["qa-1", "qa-2", "qa-3"]})

# 概率列表按 [news, code, qa] 的位置对应，seed 让本次选择序列可复现。
mixed = interleave_datasets(
    [news, code, qa],
    probabilities=[0.7, 0.2, 0.1],
    seed=42,
)

# 默认 first_exhausted：任一来源耗尽时停止，输出长度不必等于三类的比例。
print(mixed["source"])

默认的 stopping_strategy 是 first_exhausted。因此，某个来源先被取完时，混合数据集就结束；概率较低的来源可能还没有取完，概率较高的来源也可能只消耗了一部分。不要把一次短样本输出当成长期分布统计。

三种停止策略怎么选

真正决定数据量和重复行为的是停止策略。官方文档列出三种值：first_exhausted、all_exhausted 和 all_exhausted_without_replacement。

如果选择 all_exhausted，低概率来源可能需要很长时间才被抽到，输出规模也可能明显膨胀。这个策略适合明确接受过采样的任务，不适合把输出长度当成固定预算的任务。

from datasets import Dataset, interleave_datasets

# 两个来源长度不同，用于突出“过采样”和“无放回”的差异。
short = Dataset.from_dict({"text": ["short-1", "short-2"]})
long = Dataset.from_dict({"text": ["long-1", "long-2", "long-3", "long-4"]})

# all_exhausted 允许短来源在需要时再次被抽取，直到两个来源都被覆盖。
oversampled = interleave_datasets(
    [short, long],
    probabilities=[0.5, 0.5],
    seed=7,
    stopping_strategy="all_exhausted",
)

# 无放回策略不重复消费样本，适合把每条记录视为一次性输入的场景。
without_replacement = interleave_datasets(
    [short, long],
    probabilities=[0.5, 0.5],
    seed=7,
    stopping_strategy="all_exhausted_without_replacement",
)

print(oversampled["text"])
print(without_replacement["text"])

streaming 数据要额外检查什么

当输入是 IterableDataset 时，返回结果仍然是可迭代数据集。你不应该用 map-style Dataset 的随机索引方式检查它，而要通过迭代器逐条读取。官方文档还提醒，流式数据的分片数会取各来源中最小的那个，以便保持并行读取的一致性。

from datasets import interleave_datasets, load_dataset

# streaming=True 只建立迭代入口，不把完整语料一次性加载到内存。
zh = load_dataset("allenai/c4", "zh", split="train", streaming=True)
en = load_dataset("allenai/c4", "en", split="train", streaming=True)

# 流式混合仍然可以按概率选择来源；固定 seed 便于重复同一选择逻辑。
stream = interleave_datasets(
    [zh, en],
    probabilities=[0.6, 0.4],
    seed=42,
)

# 只读取前几条做流程检查，避免把无限或超大数据流全部消费掉。
for index, row in enumerate(stream):
    print(row["text"][:80])
    if index == 4:
        break

流式场景的重点不是把所有样本数预先算出来，而是确认分片、停止条件和消费方式。尤其在分布式或 DataLoader worker 中，停止策略按进程应用，最终总样本数可能比单进程直觉更少。

排查概率混合的四个检查点

1. 先核对数据集列表与 probabilities 的位置，避免把标签和权重错配。
2. 固定 seed 复现同一条选择序列，再分别改变权重，观察问题来自概率还是数据耗尽。
3. 确认是否接受短来源重复：接受就考虑 all_exhausted，不接受就使用默认策略或无放回策略。
4. streaming 场景只迭代有限样本做检查，不用随机索引，也不把一次迭代长度当作完整数据规模。

一句话记忆：probabilities 决定“从谁取”，seed 决定“能否重现”，stopping_strategy 决定“何时停以及是否允许补样”。三者分别配置，排查时也分别确认。

相关问题

不传 probabilities 会怎样？

函数会按来源轮换取样，而不是按随机概率选择。适合需要固定轮转的场景；如果要表达不同来源权重，就显式传入与数据集顺序对应的概率列表。

概率总和必须等于 1 吗？

应把它写成一组合法的相对概率并保持与来源一一对应。实践中直接使用和为 1 的浮点列表最容易读和排查，也避免把权重语义和来源顺序混在一起。

为什么设置概率后输出仍然不均匀？

概率描述的是每次选择来源的机会，不是小样本上的硬配额；同时，数据集长度和停止策略会截断输出。固定 seed 后增加样本量，再结合停止策略判断分布是否符合预期。