Hugging Face Datasets filter 按条件保留样本的操作方法

Hugging Face Datasets 的 Dataset.filter 适合在加载数据后按字段条件保留样本。最重要的边界是：它返回一个新的 Dataset，不会把原对象原地改小；因此应把结果接到新变量，再根据需要控制索引、缓存和并行参数。

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

把条件函数写成“当前样本是否保留”的判断，先保存返回对象，再用行数和字段检查结果。小数据集优先保持默认缓存和单进程，只有条件函数稳定、数据量足够大时再考虑 num_proc。

先分清原数据集和筛选结果

我在整理训练语料时，最容易犯的错是把 filter 当成列表的 remove：调用之后继续使用旧变量，却以为它已经只剩合格样本。官方处理指南明确说明处理方法返回新的 Dataset 对象，所以建议保留 raw_ds，把结果命名为 filtered_ds。

下面的示例只保留标签为 1 且文本长度至少为 20 的记录。条件函数返回 True 才会留下该行；代码中的注释解释了输入、返回值和缓存意图。

from datasets import load_dataset

# 读取一个公开数据集的训练切分，原对象后面仍可复用
raw_ds = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")

def keep_example(example):
    # 返回 True 表示保留当前样本，False 表示过滤掉
    text = example.get("text") or ""
    return example.get("label") == 1 and len(text.strip()) >= 20

# filter 返回新 Dataset，不要依赖原对象被原地修改
filtered_ds = raw_ds.filter(keep_example)

print(raw_ds.num_rows, filtered_ds.num_rows)

如果后续训练还需要原始数据，就不要覆盖 raw_ds。只有确定不再需要原对象时，才可以写成 raw_ds = raw_ds.filter(keep_example)；这只是变量重新绑定，不代表 API 做了原地更新。

需要行号时再打开 with_indices

默认回调只接收当前样本字典。如果条件和原始位置有关，例如只保留某个分片的行，使用 with_indices=True，让回调接收第二个参数。位置是数据进入当前 Dataset 后的索引，不能把它当成业务主键；如果数据经过排序、选择或拆分，业务 ID 仍应来自显式字段。

def keep_shard(example, index):
    # index 是当前 Dataset 的位置；业务主键仍应读取 example 的字段
    return index % 4 == 0 and bool(example.get("text"))

# 仅保留当前位置落在每四行一组中的样本
shard_preview = raw_ds.filter(keep_shard, with_indices=True)

如果只需要几个字段，可以用 input_columns 改变回调签名，减少条件函数对整行字典的依赖：

def keep_text_label(text, label):
    # input_columns 会按声明顺序传入字段值
    return label == 1 and len((text or "").strip()) >= 20

filtered_ds = raw_ds.filter(
    keep_text_label,
    input_columns=["text", "label"],
)

这两种写法不要混用错误的参数顺序：用了 input_columns 后，回调不再接收完整样本字典；还要确保列名确实存在，否则应先查看 raw_ds.column_names。

缓存和并行只在条件稳定后调整

filter 支持 load_from_cache_file、cache_file_name 和 num_proc。默认缓存可以避免完全相同的处理重复计算，但条件函数依赖外部变量、时间或随机数时，缓存命中可能让你误以为条件重新执行了。需要重新计算时，可以临时传入 load_from_cache_file=False，排查完成后再恢复默认策略。

filtered_ds = raw_ds.filter(
    keep_example,
    # 调试条件变化时暂时跳过旧缓存，正式流程可恢复为默认值
    load_from_cache_file=False,
    # 数据量和回调适合多进程时再打开，并确保回调可被序列化
    num_proc=2,
)

并行不是无条件加速：回调必须适合多进程，进程通信和数据写入也有成本。小数据集、复杂闭包或需要稳定调试时，先使用默认的单进程路径。若需要固定产物位置，可提供 cache_file_name；不要让多个不同条件的任务共用一个手工缓存文件名。

用三项检查确认筛选结果

筛选后先检查行数，再检查列名和一小段字段值。行数变少只说明有记录被排除，不代表条件逻辑一定正确；如果结果为零，优先检查字段类型、空字符串处理和标签编码。

最后记住边界：filter 负责筛选行，不负责把文本转成张量，也不负责修改字段格式。字段变换应交给 map，输出格式则另行处理。把职责拆开，后续训练代码会更容易复查。

相关问题

filter 和 map 有什么区别？ filter 决定保留哪些行，map 通常用于新增或变换字段；两者都返回新的数据集对象。

with_indices 是原始文件的行号吗？ 它表示当前 Dataset 视图中的位置，不保证等同于来源文件行号，也不能代替业务 ID。

什么时候不该开 num_proc？ 数据量较小、回调依赖不可序列化状态，或你正在排查条件逻辑时，先保持默认单进程。