【RAG召回】bge实现向量相似度索引
sentence-transformers 是一个非常强大的 Python 框架，它可以将句子或段落转换成高质量、高信息密度的数字向量（称为“嵌入”或 Embeddings）。它厉害的地方在于，语义上相似的句子，其向量在空间中的距离也更近。
这使得我们能够超越简单的关键词匹配，去实现真正理解“意思”的搜索。例如，用户搜索“手提电脑”，我们能轻松地找出包含“笔记本电脑”的文档。
下面，我们同样通过一系列代码示例，快速上手 sentence-transformers。
sentence-transformers 全功能中文示例
本篇将展示如何使用 sentence-transformers 进行向量生成和相似度搜索。
准备工作
首先，你需要安装 sentence-transformers。它依赖于 PyTorch，通常会自动安装。同时，我们也会用到 scikit-learn 来辅助计算。
# 推荐同时安装 scikit-learn 以便使用其中的工具
pip install sentence-transformers scikit-learn
示例 1：将文本转换为向量
向量搜索的第一步，也是最关键的一步，就是将文本转换为数字向量。
from sentence_transformers import SentenceTransformer# 1. 加载一个预训练好的中文模型
# 'shibing624/text2vec-base-chinese' 是一个常用且效果不错的中文模型
# 第一次运行时会自动下载模型，请耐心等待
model = SentenceTransformer('shibing624/text2vec-base-chinese')# 2. 准备要编码的句子
sentences = ["今天天气真不错","今天天气很糟糕","我爱吃北京烤鸭"
]# 3. 使用 model.encode() 将句子转换为向量
embeddings = model.encode(sentences)print("向量的维度:", embeddings.shape)
print("\n第一个句子的向量 (部分展示):")
print(embeddings[0][:5]) # 向量很长，我们只看前5个维度
运行结果:
向量的维度: (3, 768)第一个句子的向量 (部分展示):
[-0.4398335 -0.43577313 0.0303108 -0.2334005 0.2187311 ]
可以看到，每个句子都被转换成了一个包含768个数字的向量。
示例 2：计算句子间的相似度
拿到向量后，我们就可以通过计算它们之间的“余弦相似度”（Cosine Similarity）来判断其语义的接近程度。分数范围在-1到1之间，越接近1代表越相似。
from sentence_transformers import SentenceTransformer, utilmodel = SentenceTransformer('shibing624/text2vec-base-chinese')# 我们有两个句子列表
sentences1 = ['一个男人在弹吉他', '一只猫在玩毛线球']
sentences2 = ['有个人正在演奏乐器', '那个动物在睡觉', '一个男人在弹奏吉他']# 将它们全部编码为向量
embeddings1 = model.encode(sentences1)
embeddings2 = model.encode(sentences2)# 2. 使用 util.cos_sim 计算相似度矩阵
# 这会计算 sentences1 中每个句子与 sentences2 中每个句子之间的相似度
cosine_scores = util.cos_sim(embeddings1, embeddings2)# 打印结果
for i in range(len(sentences1)):for j in range(len(sentences2)):print(f"'{sentences1[i]}' vs '{sentences2[j]}' 的相似度: {cosine_scores[i][j]:.4f}")print("-" * 30)运行结果:
'一个男人在弹吉他' vs '有个人正在演奏乐器' 的相似度: 0.7410
'一个男人在弹吉他' vs '那个动物在睡觉' 的相似度: 0.0526
'一个男人在弹吉他' vs '一个男人在弹奏吉他' 的相似度: 0.9472
------------------------------
'一只猫在玩毛线球' vs '有个人正在演奏乐器' 的相似度: -0.0163
'一只猫在玩毛线球' vs '那个动物在睡觉' 的相似度: 0.5212
'一只猫在玩毛线球' vs '一个男人在弹奏吉他' 的相似度: 0.0210
------------------------------
可以看到，语义相近的句子对（如“弹吉他”和“弹奏吉他”）获得了非常高的相似度分数。
示例 3：构建一个简单的语义搜索引擎
现在，我们可以整合以上步骤，构建一个和 rank-bm25 功能类似的搜索引擎。
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np# --- 准备阶段 ---
model = SentenceTransformer('shibing624/text2vec-base-chinese')corpus = ["北京是中国的首都，也是一座历史悠久的文化名城。","上海是中国的经济中心，拥有繁忙的港口和现代化的建筑。","深圳是中国科技创新的重要城市，被誉为“中国硅谷”。","广州的美食文化闻名全国，是粤菜的发源地。","学习人工智能技术需要扎实的数学基础和编程能力。","中国的历史源远流长，有许多著名的历史人物和事件。"
]# 1. 对语料库进行向量化 (这一步可以预先计算并存储，避免每次查询都重新计算)
print("正在对语料库进行向量化...")
corpus_embeddings = model.encode(corpus, convert_to_tensor=True)
print("向量化完成！")# --- 查询阶段 ---
def search(query, top_k=3):# 2. 将查询语句向量化query_embedding = model.encode(query, convert_to_tensor=True)# 3. 计算查询向量与所有语料库向量的余弦相似度# `util.cos_sim` 是 sentence-transformers 的内置函数，更高效cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]# 4. 寻找得分最高的 top_k 个结果# np.argpartition 比 np.argsort 更快，因为它只保证第k个位置的元素正确top_results_indices = np.argpartition(-cos_scores, range(top_k))[:top_k]print(f"\n查询: '{query}'")print(f"--- 最相关的 {top_k} 个结果 ---")# 按照分数从高到低排序并输出sorted_indices = top_results_indices[np.argsort(-cos_scores[top_results_indices])]for idx in sorted_indices:print(f"分数: {cos_scores[idx]:.4f} | 文档: {corpus[idx]}")# 执行搜索
search("中国的经济发展怎么样？")
search("历史古都")
运行结果:
正在对语料库进行向量化...
向量化完成！查询: '中国的经济发展怎么样？'
--- 最相关的 3 个结果 ---
分数: 0.6908 | 文档: 上海是中国的经济中心，拥有繁忙的港口和现代化的建筑。
分数: 0.6186 | 文档: 深圳是中国科技创新的重要城市，被誉为“中国硅谷”。
分数: 0.4485 | 文档: 学习人工智能技术需要扎实的数学基础和编程能力。查询: '历史古都'
--- 最相关的 3 个结果 ---
分数: 0.7711 | 文档: 北京是中国的首都，也是一座历史悠久的文化名城。
分数: 0.6695 | 文档: 中国的历史源远流长，有许多著名的历史人物和事件。
分数: 0.4418 | 文档: 广州的美食文化闻名全国，是粤菜的发源地。
即使查询中没有“北京”或“历史”这样的关键词，模型也能理解“历史古都”的含义并找到最相关的句子。
示例 4：在海量数据中挖掘相似句子对
如果你有一个庞大的文本列表，想从中找出所有语义相似的句子对（例如，寻找重复问题、相似评论），util.paraphrase_mining 函数会非常有用。
from sentence_transformers import SentenceTransformer, utilmodel = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2') # 使用一个多语言模型# 假设这是一个论坛的评论列表
comments = ['如何安装这个软件？','我的电脑蓝屏了，求助！','请问这个App怎么安装？','系统崩溃了，怎么办？','这个东西的价格是多少？','售价多少钱？','启动时遇到错误。'
]# 使用 paraphrase_mining 寻找相似度高于 0.8 的句子对
similar_pairs = util.paraphrase_mining(model, comments, show_progress_bar=True, score_threshold=0.8)print("\n--- 找到的相似句子对 (阈值 > 0.8) ---")
for score, i, j in similar_pairs:print(f"分数: {score:.4f}")print(f" - {comments[i]}")print(f" - {comments[j]}")print("-" * 20)运行结果:
Batches: 100%|██████████| 1/1 [00:00<00:00, 8.85it/s]--- 找到的相似句子对 (阈值 > 0.8) ---
分数: 0.9413- 这个东西的价格是多少？- 售价多少钱？
--------------------
分数: 0.9126- 如何安装这个软件？- 请问这个App怎么安装？
--------------------
分数: 0.8037- 我的电脑蓝屏了，求助！- 系统崩溃了，怎么办？
--------------------
这个功能对于数据清洗、文本聚类等任务非常实用。
总结与展望
sentence-transformers 将复杂的深度学习模型封装得极其易用，是实现现代语义搜索的首选工具。
-
优点:
- 理解语义: 能处理同义词、近义词，搜索结果更智能。
- 模型丰富: 拥有大量覆盖多语言、多领域的预训练模型。
- 功能全面: 不仅能搜索，还能用于聚类、信息挖掘等多种NLP任务。
-
下一步进阶:
当你的语料库达到数百万甚至上亿级别时，逐一计算余弦相似度会变得非常缓慢。这时，就需要专门的向量检索引擎了，例如 Faiss (来自Facebook) 或 Annoy (来自Spotify)。这些工具可以创建向量索引，实现**近似最近邻（ANN）**搜索，以微小的精度损失换来成千上万倍的搜索速度提升。
sentence-transformers生成的向量可以无缝对接到这些检索引擎中。