如何高效统计化学品在海量文献摘要中的出现频率？

问题描述 #

现有两个大型数据集：

• 包含约10万个化学品名称的列表：

chemicals_list = ["chemical1", "chemical2", ..., "chemical100000"]

• 包含约5000万篇文献摘要的列表：

abstracts_list = ["abstract1 is very very very long", "abstract2 is very very very VERY long", ..., "abstract50000000 is pretty long as well"]

需要构建一个frequency_dict，将每个化学品映射到包含该化学品的摘要数量（文档频率）。

现有实现的问题 #

当前使用两层嵌套循环的实现时间复杂度为O(M*N)（M为化学品数量，N为摘要数量），计算量达到5e12次，运行速度极慢，且存在逻辑错误：同一篇摘要中同一化学品出现多次时会重复计数，不符合统计文档频率的需求。

原代码：

frequency_dict = {}
for c in chemicals_list:
    exact_entity = f' {c} ' # 试图匹配完整实体，但会漏掉摘要开头/结尾的情况
    for abstract_text in abstracts_list: 
        if exact_entity in abstract_text:
            if c in frequency_dict.keys():
                frequency_dict[c] += 1
            else: 
                frequency_dict[c] = 1

优化方案 #

一、CPU层面优化：Aho-Corasick多模式匹配 #

核心思路是遍历一次所有摘要，在单篇摘要中一次性匹配所有化学品，避免重复遍历摘要。使用Aho-Corasick自动机可以在O(NL + MK)的时间复杂度内完成匹配（L为摘要平均长度，K为化学品平均长度），效率远超原方法。

实现步骤：

1. 预处理化学品，生成三种匹配模式（覆盖摘要开头、中间、结尾的情况，确保匹配完整实体）；
2. 构建Aho-Corasick自动机，将所有匹配模式与对应化学品关联；
3. 遍历每篇摘要，用自动机找到所有匹配的化学品，去重后更新计数。

代码示例：

import ahocorasick

# 1. 预处理匹配模式：映射模式到对应的化学品
pattern_to_chemical = {}
for c in chemicals_list:
    # 三种模式覆盖不同位置的匹配需求
    pattern_to_chemical[f' {c} '] = c   # 中间位置
    pattern_to_chemical[f'{c} '] = c    # 摘要开头
    pattern_to_chemical[f' {c}'] = c    # 摘要结尾

# 2. 构建Aho-Corasick自动机
automaton = ahocorasick.Automaton()
for pattern, chemical in pattern_to_chemical.items():
    automaton.add_word(pattern, chemical)
automaton.make_automaton()

# 3. 初始化频率字典，避免重复判断键是否存在
frequency_dict = {c: 0 for c in chemicals_list}

# 4. 遍历所有摘要统计频率
for abstract in abstracts_list:
    matched_chemicals = set()
    # 自动机一次性找出所有匹配的化学品
    for end_idx, chemical in automaton.iter(abstract):
        matched_chemicals.add(chemical)
    # 每篇摘要中同一化学品只计数一次
    for c in matched_chemicals:
        frequency_dict[c] += 1

二、GPU加速方案：利用RAPIDS库并行处理 #

如果GPU资源可用，可使用RAPIDS（NVIDIA的GPU加速数据科学库）进一步提升速度，利用GPU的并行计算能力处理大规模摘要数据。

实现步骤：

1. 将数据集转换为cuDF格式（RAPIDS的DataFrame）；
2. 使用cuml库的Aho-Corasick实现进行多模式匹配；
3. 分块处理摘要（避免GPU内存溢出），统计每个化学品的文档频率。

代码示例：

import cudf
from cuml import AhoCorasick

# 1. 预处理化学品与匹配模式
patterns = []
chemicals = []
for c in chemicals_list:
    patterns.extend([f' {c} ', f'{c} ', f' {c}'])
    chemicals.extend([c, c, c])

# 2. 构建GPU版Aho-Corasick自动机
ac = AhoCorasick(cudf.Series(patterns))

# 3. 初始化频率统计数组
chemicals_unique = cudf.Series(chemicals_list).unique()
frequency_df = cudf.Series(0, index=chemicals_unique)

# 4. 分块处理摘要（根据GPU内存调整chunk_size）
chunk_size = 100000
abstracts_df = cudf.Series(abstracts_list)

for i in range(0, len(abstracts_df), chunk_size):
    chunk = abstracts_df[i:i+chunk_size]
    # 匹配当前块中的所有化学品
    matches = ac.match(chunk)
    if len(matches) == 0:
        continue
    # 提取匹配对应的化学品并去重
    matched_chemicals = cudf.Series(chemicals)[matches[:, 1]].unique()
    # 更新频率计数
    frequency_df.loc[matched_chemicals] += 1

# 转换为Python字典
frequency_dict = frequency_df.to_dict()

额外优化点 #

• 预处理摘要：统一转换为小写（如果大小写不敏感），减少匹配的模式数量；
• 内存优化：如果摘要数据过大，可采用逐行读取文件的方式，避免一次性加载所有摘要到内存；
• 去重处理：确保同一篇摘要中同一化学品只计数一次，这是统计文档频率的核心要求。

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