如何从datasets.arrow_dataset.Dataset中筛选hypothesis列含指定关键词的行子集

• 误用了pandas的filter方法：该方法仅用于筛选列/行标签，不支持按行内容条件过滤
• 逻辑判断写反：原代码写的searchfor in example['hypothesis']是判断整个关键词列表是否为字符串的子串，该条件永远为False，因此返回空结果

方案1：直接使用Dataset原生filter方法（推荐，无需转DataFrame，性能更高） #

该方案不需要转换数据集格式，适合处理大规模数据集：

import re

searchfor = ['she', 'he','his','her','him','boys','woman','Woman','girl','men','man','female','girls']
# 构造带单词边界的正则规则，避免匹配到类似the中的he这类误判结果
# 加re.IGNORECASE可实现大小写不敏感匹配，此时可以删除searchfor里重复的大小写变体如Woman
pattern = re.compile(r'\b(' + '|'.join(searchfor) + r')\b')

# 直接筛选符合条件的行
filtered_dataset = crop_dataset.filter(lambda example: pattern.search(example['hypothesis']) is not None)

如果不需要严格整词匹配，仅需要判断字符串包含任意关键词子串，可以简化为：

searchfor = ['she', 'he','his','her','him','boys','woman','Woman','girl','men','man','female','girls']
filtered_dataset = crop_dataset.filter(lambda example: any(word in example['hypothesis'] for word in searchfor))

方案2：转Pandas DataFrame的正确写法 #

如果确实需要使用Pandas处理，可以用apply构造布尔索引筛选行：

import re
import datasets

df_temp = crop_dataset.to_pandas()
searchfor = ['she', 'he','his','her','him','boys','woman','Woman','girl','men','man','female','girls']
pattern = re.compile(r'\b(' + '|'.join(searchfor) + r')\b')

# 按条件筛选行
rows_where_found = df_temp[df_temp['hypothesis'].apply(lambda x: pattern.search(x) is not None)]
# 如需转回Dataset格式可执行以下代码
# filtered_dataset = datasets.Dataset.from_pandas(rows_where_found)

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