当前位置：首页 > news >正文

本地大模型编程实战(34)使用faiss实现语义检索

news 2025/9/18 14:24:09

FAISS（Facebook AI Similarity Search）是 Facebook AI 研究院开发的一款高效的向量相似性搜索库，专门用于优化大规模高维向量的近邻搜索任务，广泛应用于语义检索、图像检索、推荐系统等场景。
本文将详细介绍如何使用 FAISS 实现语义检索。具体内容包括：

将内容矢量化并保存为索引
在保存索引的同时，保存每条信息的相关信息metadata
通过矢量库进行语义检索

语义检索是一种基于内容含义（语义）而非表面关键词的信息检索技术，核心目标是让计算机 “理解” 用户查询的真实意图，从而返回含义相关的结果，而非仅匹配字面关键词的内容。
传统检索（如早期搜索引擎）依赖关键词匹配，例如用户搜索 “如何缓解头痛”，仅会返回包含 “缓解”“头痛” 等字面词汇的内容，若某篇文章写 “应对偏头痛的方法”（未出现 “缓解”“头痛”），则会被遗漏。
语义检索则通过语义嵌入（Embedding）技术解决这一问题：

将文本转化为 “语义向量”：用 AI 模型（如 BERT、BGE-M3 等）把用户的查询、待检索的文档（句子 / 段落）都转化为高维向量（类似 “数字指纹”），向量的相似度直接对应内容语义的相似度。
通过向量相似度匹配结果：计算 “查询向量” 与 “文档向量” 的相似度（如余弦相似度），排序后返回相似度最高的内容，实现 “含义相关即匹配”。

数据结构设计

数据项中包含metadata，可以存储任何重要内容。

documents = [{"text": "人工智能是计算机科学的一个分支，研究如何使机器具有智能","metadata": {"id": "doc_001","category": "人工智能基础","source": "教科书","publish_date": "2023-01-15"}},{"text": "机器学习是人工智能的核心技术，让计算机能从数据中学习","metadata": {"id": "doc_002","category": "机器学习","source": "技术博客","publish_date": "2023-03-20"}},{"text": "深度学习是机器学习的一个子领域，基于人工神经网络","metadata": {"id": "doc_003","category": "深度学习","source": "论文","publish_date": "2023-02-10"}},{"text": "自然语言处理专注于让计算机理解和生成人类语言","metadata": {"id": "doc_004","category": "自然语言处理","source": "技术文档","publish_date": "2023-04-05"}},{"text": "计算机视觉是人工智能的一个重要方向，处理图像和视频","metadata": {"id": "doc_005","category": "计算机视觉","source": "教程","publish_date": "2023-01-28"}}
]

初始化模型和存储路径

model = SentenceTransformer('BAAI/bge-m3',cache_folder=Path("./model")    # 模型会下载到这个目录，第一次执行会下载模型，比较慢。
)
prefix = "为这个句子生成表示以用于检索相关句子：" # 定义存储路径
data_folder_path = Path("./data")
if not data_folder_path.exists():data_folder_path.mkdir(parents=True, exist_ok=True)
index_path = data_folder_path / "faiss_index_with_metadata.index"
metadata_path = data_folder_path / "metadata.json"

上面代码中有两个要点：
使用bge-m3 ：因为在中文检索/相似度匹配场景中，bge-m3 远超 all-MiniLM-L6-v2
使用前缀 prefix：
在文本嵌入（Embedding）过程中添加特定前缀是指令微调（Instruction Tuning）思想在嵌入模型中的应用，主要作用是引导模型生成更符合特定任务需求的向量表示，从而提升下游任务（如语义检索）的效果。
当一个模型需要同时支持多种任务时，前缀可以作为 “任务标识符”：

检索任务：用"为检索生成表示：…"
聚类任务：用"为聚类生成表示：…"
问答匹配：用"为匹配问题和答案生成表示：…"

前缀的本质是通过自然语言指令引导模型聚焦任务需求，这是基于大语言模型时代的典型优化手段。对于 BGE-M3 这类支持指令微调的模型，使用官方推荐的前缀能显著提升检索效果；而对于早期不支持指令的模型（如原始 BERT），前缀可能无效甚至产生干扰。

生成数据库

下面使用 FAISS 生成嵌入数据库，并将其保存在本地文件夹中。

def process_documents(docs):# 提取文本并添加前缀texts = [prefix + doc["text"] for doc in docs]# 生成向量embeddings = model.encode(texts,normalize_embeddings=True).astype(np.float32)# 提取metadata列表（保持与向量顺序一致）metadatas = [doc["metadata"] for doc in docs]return embeddings, metadatasdef create_db():"""生成矢量库并保存"""if index_path.exists():print("数据库已创建")return# 处理文档生成向量和metadataembeddings, metadatas = process_documents(documents)print(f"向量形状: {embeddings.shape}")print(f"metadata数量: {len(metadatas)}")# 创建内积索引index = faiss.IndexFlatIP(embeddings.shape[1])index.add(embeddings)print(f"索引向量数量: {index.ntotal}")# 保存FAISS索引faiss.write_index(index, str(index_path))# 保存metadata（使用JSON格式）with open(metadata_path, "w", encoding="utf-8") as f:json.dump(metadatas, f, ensure_ascii=False, indent=2)print("\n数据保存完成：")print(f" - 索引: {index_path}")print(f" - Metadata: {metadata_path}")create_db()

保存json格式的metadata时，其顺序要与与向量顺序一致，这样在语义检索的时候，可以根据检索出来的矢量位置找到对应的 metadata 。

加载矢量库

与 metadata 一同加载。

def load_data():# 加载索引index = faiss.read_index(str(index_path))# 加载metadatawith open(metadata_path, "r", encoding="utf-8") as f:metadatas = json.load(f)print("\n数据加载完成：")print(f" - 索引向量数量: {index.ntotal}")print(f" - Metadata数量: {len(metadatas)}")return index, metadatasindex, metadatas = load_data()

语义检索

下面实现 FAISS语义检索 功能，在检索到最匹配文本的同时，也会找到关联的 metadata 。

def search_with_metadata(query, top_k=2):# 处理查询query_text = prefix + queryquery_embedding = model.encode([query_text],normalize_embeddings=True).astype(np.float32)# 搜索相似向量scores, indices = index.search(query_embedding, top_k)# 关联metadata并返回结果results = []for i in range(top_k):idx = indices[0][i]results.append({"score": float(scores[0][i]),"metadata": metadatas[idx],"text": documents[idx]["text"]  # 也可以从metadata中存储text})return results

我们用下面的代码做一下语义检索测试：

if __name__ == '__main__':    query = "什么是深度学习"results = search_with_metadata(query, top_k=2)print(f"\n查询: {query}")print("检索结果:")for i, result in enumerate(results, 1):print(f"\n结果 {i}:")print(f"相似度得分: {result['score']:.4f}")print(f"文本内容: {result['text']}")print(f"Metadata: {json.dumps(result['metadata'], ensure_ascii=False, indent=2)}")

执行上述代码后，将按照相似度高低返回两个结果：

结果 1:  
相似度得分: 0.8673
文本内容: 深度学习是机器学习的一个子领域，基于人工神经网络
Metadata: {"id": "doc_003","category": "深度学习","source": "论文","publish_date": "2023-02-10"
}结果 2:
相似度得分: 0.7643
文本内容: 机器学习是人工智能的核心技术，让计算机能从数据中学习
Metadata: {"id": "doc_002","category": "机器学习","source": "技术博客","publish_date": "2023-03-20"
}

靠谱，哈哈：）