基于SBERT在百万级文本中高效匹配相似短语的技术咨询

问题背景 #

拥有含156个短语的列表，以及约1100万词的文本语料库，需为每个短语在语料库中找到最相似的片段，指定使用SBERT算法（基于sentence-transformers库）实现。当前遇到的问题：

• 现有文档仅覆盖两个列表间的相似度计算，未提及大文本中短语匹配的可行方案；
• 将语料拆分为句子列表后计算耗时极长，且仅能匹配完整句子，无法覆盖1-2词或跨多句的任意粒度相似片段。

当前实现代码：

phrase_list = ['Gregor Samse', 'in his bed into', 'horrible creature'...] # 共156个短语

very_long_string= 'One morning, when Gregor Samsa woke from troubled dreams, he found himself transformed in his bed into a horrible vermin. He lay on his armour-like back, and if he lifted his head a little he could see his brown belly, slightly domed and divided by arches into stiff sections. The bedding was hardly able to cover it and seemed ready to slide off any moment. His many legs, pitifully thin compared with the size of the rest of him, waved about helplessly as he looked...' 
# 语料约1100万词

# 将语料转换为句子列表
string_to_sent_list = very_long_string.split(".")

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer('all-mpnet-base-v2')

phrase_embeddings = model.encode(phrase_list, convert_to_tensor=True)
sent_embeddings = model.encode(string_to_sent_list, convert_to_tensor=True)

similarity_dict = {}
for i, phraselist_phrase in enumerate(phrase_list):
    similarities = util.cos_sim(phrase_embeddings[i], sent_embeddings)
    matches = [string_to_sent_list[j] for j, sim in enumerate(similarities[0]) if sim > 0.65] 
    similarity_dict[phraselist_phrase] = matches

print(similarity_dict)

# 注：原代码中dict没有to_csv方法，需用pandas转换后保存
# import pandas as pd
# pd.DataFrame.from_dict(similarity_dict, orient='index').to_csv('similarity dict.csv')

解答 #

一、替代实现方案 #

1. 优化语料片段生成逻辑

放弃仅按句子拆分的方式，改用滑动窗口生成多粒度片段，同时加入过滤减少计算量：

• 针对短短语（1-2词）：生成语料的1-3词滑动窗口片段（步长1）；针对长短语，调整窗口大小匹配短语长度；
• 针对跨句片段：将相邻2-3个句子拼接后再生成滑动窗口片段；
• 预过滤：先提取语料中的核心词（如名词、动词），仅保留包含phrase_list中至少一个核心词的片段，大幅减少候选片段数量。

2. 优化相似度计算效率

• 批量计算相似度矩阵：避免循环单个短语计算，一次性生成所有短语与所有片段的相似度矩阵，再批量筛选匹配项，减少重复计算开销：

  # 假设已生成候选片段列表corpus_fragments
  corpus_embeddings = model.encode(corpus_fragments, convert_to_tensor=True, device='cuda')  # 用GPU加速
  phrase_embeddings = model.encode(phrase_list, convert_to_tensor=True, device='cuda')
  
  # 一次性计算所有短语与片段的相似度
  sim_matrix = util.cos_sim(phrase_embeddings, corpus_embeddings)
  
  similarity_dict = {}
  threshold = 0.65
  top_k = 5  # 只保留top5最相似的片段，避免结果过多
  
  for idx, phrase in enumerate(phrase_list):
      # 获取top_k相似片段的索引
      top_indices = sim_matrix[idx].topk(k=top_k).indices
      # 筛选相似度超过阈值的片段
      matches = [corpus_fragments[i] for i in top_indices if sim_matrix[idx][i] > threshold]
      similarity_dict[phrase] = matches
  

• 使用轻量模型：替换all-mpnet-base-v2为all-MiniLM-L6-v2，编码速度提升3-4倍，语义相似度损失极小；
• GPU加速：确保模型运行在CUDA设备上，编码效率可提升10倍以上。

3. 换用向量检索库加速

对于百万级候选片段，用faiss做近似最近邻搜索，比全量计算相似度快几个数量级：

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer('all-MiniLM-L6-v2', device='cuda')
phrase_embeddings = model.encode(phrase_list, convert_to_tensor=True).cpu().numpy()
corpus_fragments = []  # 预先生成的候选片段列表
corpus_embeddings = model.encode(corpus_fragments, convert_to_tensor=True).cpu().numpy()

# 构建faiss索引（L2距离，归一化后与余弦相似度正相关）
index = faiss.IndexFlatL2(corpus_embeddings.shape[1])
index.add(corpus_embeddings)

# 转换余弦相似度阈值为L2距离阈值（归一化向量下：cos_sim = 1 - (L2²)/2）
threshold = 0.65
l2_threshold = np.sqrt(2 * (1 - threshold))

k = 5
distances, indices = index.search(phrase_embeddings, k)

similarity_dict = {}
for idx, phrase in enumerate(phrase_list):
    # 筛选符合阈值的片段
    valid_mask = distances[idx] < l2_threshold
    valid_indices = indices[idx][valid_mask]
    matches = [corpus_fragments[i] for i in valid_indices]
    similarity_dict[phrase] = matches

二、当前方法的匹配项问题 #

当前方法的匹配单位是完整句子，只会将相似度超过阈值的整个句子存入结果，不会单独保存同句子中的不同子短语匹配项。例如，若某句子同时包含"Gregor Samsa"和"horrible vermin"，当匹配"horrible creature"时，该句子会被作为一个整体加入匹配列表，但不会提取"horrible vermin"这个子片段。

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