如何对句子嵌入向量余弦相似度结果排序并返回对应输入句子

核心逻辑是保留两个语料库的原始句子索引，将相似度值与对应的两个原句绑定后再排序，完整可运行代码如下：

import tensorflow_hub as hub
from sklearn.metrics.pairwise import cosine_similarity

# 加载预训练模型
embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

# 分开保存原始句子和对应向量，不要覆盖原句变量
sentences1 = [
    "I'd like an apple juice",
    "An apple a day keeps the doctor away",
    "Eat apple every day",
    "We buy apples every week",
    "We use machine learning for text classification",
    "Text classification is subfield of machine learning"
]
embeddings1 = embed(sentences1)

sentences2 = [
    "I'd like an orange juice",
    "An orange a day keeps the doctor away",
    "Eat orange every day",
    "We buy orange every week",
    "We use machine learning for document classification",
    "Text classification is some subfield of machine learning"
]
embeddings2 = embed(sentences2)

# 计算相似度矩阵
similarity_matrix = cosine_similarity(embeddings1, embeddings2)

# 组装所有句对和对应相似度
all_pairs = []
for idx1, s1 in enumerate(sentences1):
    for idx2, s2 in enumerate(sentences2):
        all_pairs.append({
            "cosine_similarity": round(similarity_matrix[idx1][idx2], 4),
            "corpus1_sentence": s1,
            "corpus2_sentence": s2
        })

# 按相似度降序排序，reverse改为False就是升序
sorted_pairs = sorted(all_pairs, key=lambda x: x["cosine_similarity"], reverse=True)

# 示例输出：打印相似度前5的句对
for pair in sorted_pairs[:5]:
    print(f"相似度：{pair['cosine_similarity']}")
    print(f"语料库1句子：{pair['corpus1_sentence']}")
    print(f"语料库2句子：{pair['corpus2_sentence']}")
    print("-"*50)

扩展用法 #

• 如果只需要保留高相似度句对，可以在组装时加过滤条件，比如只保留相似度大于0.7的结果：

  if similarity_matrix[idx1][idx2] > 0.7:
      all_pairs.append(...)
  

• 如果需要对语料库1的每个句子单独匹配语料库2中最相似的句子，不需要全局排序，可以在内层循环结束后单独排序当前idx1对应的所有句对即可。

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