优化嵌套循环计算耗时：KeyBERT关键词提取性能提升

问题背景 #

我有一个包含数千行文本的DataFrame，需要用KeyBERT结合4种不同词性（POS）模式提取符合规则的关键词。原嵌套循环和itertools.product优化后的版本，在大数据集下耗时极久，急需更高效的优化方案。

原始基础代码 #

所需依赖与初始化 #

from keybert import KeyBERT
from sentence_transformers import SentenceTransformer
from keyphrase_vectorizers import KeyphraseCountVectorizer
import pandas as pd

sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
st_model = KeyBERT(model=sentence_model)

测试数据 #

data = pd.DataFrame({'text':[
    'Machine learning (ML) is a type of artificial intelligence (AI) that allows software applications to become more accurate at predicting outcomes without being explicitly programmed to do so. Machine learning algorithms use historical data as input to predict new output values.',
    'Physics is the natural science that studies matter, its fundamental constituents, its motion and behavior through space and time, and the related entities of energy and force. Physics is one of the most fundamental scientific disciplines, with its main goal being to understand how the universe behaves.',
    'Chemistry is the branch of science that deals with the properties, composition, and structure of elements and compounds, how they can change, and the energy that is released or absorbed when they change.'
]})

高效优化方案 #

方案1：批量文本处理 + Vectorizer复用 #

KeyBERT的extract_keywords支持直接传入文本列表，无需逐行循环调用模型。一次性处理所有文本，配合预初始化的Vectorizer，大幅减少模型调用的开销。

patterns = ['<J.*>*<N.*>+', '<V.*>+', '<N.*>*<V.*>+', '<J.*>*<N.*>*<V.*>+']
# 预初始化所有POS模式对应的Vectorizer
vecz = [KeyphraseCountVectorizer(pos_pattern=p) for p in patterns]

pt = []
# 每个Vectorizer对应一次批量文本处理
for vec in vecz:
    # 传入整个文本列的列表，一次性处理所有行
    batch_results = st_model.extract_keywords(
        data.text.tolist(),
        stop_words="english",
        vectorizer=vec,
        use_mmr=True,
        diversity=0.4
    )
    # 将批量结果合并到总列表中
    pt.extend(batch_results)

方案2：多进程并行处理 #

关键词提取属于CPU密集型任务，利用多进程并行处理所有文本-POS模式组合，充分利用多核CPU资源，缩短总耗时。

from multiprocessing import Pool

# 定义子进程执行的任务函数（子进程内初始化模型，避免序列化问题）
def process_single_task(args):
    text, pattern = args
    sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
    st_model = KeyBERT(model=sentence_model)
    vec = KeyphraseCountVectorizer(pos_pattern=pattern)
    return st_model.extract_keywords(
        text,
        stop_words="english",
        vectorizer=vec,
        use_mmr=True,
        diversity=0.4
    )

patterns = ['<J.*>*<N.*>+', '<V.*>+', '<N.*>*<V.*>+', '<J.*>*<N.*>*<V.*>+']
# 生成所有待处理的(text, pattern)组合
tasks = [(text, pattern) for text in data.text.tolist() for pattern in patterns]

# 启动多进程池（进程数可根据CPU核心数调整）
with Pool(processes=4) as pool:
    pt = pool.map(process_single_task, tasks)

方案3：精简模型与计算参数 #

通过调整模型和提取参数，减少不必要的计算量：

• 改用轻量模型：当前使用的all-MiniLM-L6-v2已是轻量模型，若仍需提速，可尝试更小的distiluse-base-multilingual-cased-v2
• 减少候选词数量：通过top_n参数限制MMR计算的候选词数量（默认10，可根据需求调整为15以内）
• 关闭MMR（若允许）：如果不需要关键词多样性，可移除use_mmr=True和diversity参数，改用默认的余弦相似度排序，速度会显著提升

# 调整参数后的示例
batch_results = st_model.extract_keywords(
    data.text.tolist(),
    stop_words="english",
    vectorizer=vec,
    top_n=12,  # 减少候选词数量
    # use_mmr=True, diversity=0.4  # 不需要多样性时注释此行
)

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