从 0 到 1 构建语义搜索引擎：用 FAISS 实现智能文本检索

佑瞻

于 2025-05-29 10:52:10 发布

阅读量1.1k

点赞数 25

CC 4.0 BY-SA版权

分类专栏： Transformers 文章标签：搜索引擎 faiss Transformers

本文链接：https://blog.csdn.net/The_Thieves/article/details/148305498

Transformers 专栏收录该内容

26 篇文章

订阅专栏

在日常的文本处理与信息检索场景中，无论是企业内部的海量文档、论坛社区的问答数据，还是学术领域的论文语料，我们常常面临一个核心痛点：传统关键词搜索只能匹配字面内容，却无法理解文本背后的语义关联。比如搜索 “如何离线下载数据”，可能会漏掉 “本地数据集加载方法”“无网络环境数据获取” 等表述相似但关键词不同的内容。这种 “语义鸿沟” 让高效信息检索成为技术落地的一大挑战。

今天，我们就来聊聊如何利用 FAISS（Facebook AI 相似性搜索库）构建一个语义搜索引擎，通过文本嵌入技术将自然语言转化为可计算的向量表示，让机器 “理解” 语义而非简单匹配关键词。我们将以 GitHub issues 和评论数据集为例，手把手演示从数据处理到搜索实战的全流程，无论你是处理技术文档、客服对话还是其他文本数据，都能从中获取可复用的方法论。

一、明确目标：让搜索从 “关键词匹配” 升级为 “语义理解”

我们的核心任务是：以 GitHub 开源社区的 issues 和评论为样本，构建一个能根据用户查询语义，快速返回最相关技术讨论的搜索引擎。具体来说，需要实现三个关键步骤：

数据预处理：清洗无效信息，将非结构化文本转化为适合模型处理的格式；
文本嵌入生成：用预训练模型将文本转化为高维向量，作为语义表示；
FAISS 高效搜索：利用向量索引技术实现毫秒级相似性检索。

接下来，我们就从数据处理开始，逐步拆解实现过程。

二、数据预处理：清洗数据，为搜索筑牢地基

1. 加载原始数据，过滤无效信息

首先，使用 Hugging Face 的datasets库加载 GitHub issues 数据集。初始数据包含 2855 条记录，但其中包含大量拉取请求（与问题解答无关）和无评论的 issue（无法提供有效答案），需要先过滤：

python

from datasets import load_dataset

# 加载原始数据集（包含标题、正文、评论等字段）
issues_dataset = load_dataset("lewtun/github-issues", split="train")
# 过滤条件：非拉取请求且至少有1条评论
issues_dataset = issues_dataset.filter(
    lambda x: (x["is_pull_request"] == False and len(x["comments"]) > 0)
)
print(f"过滤后有效issue数量：{len(issues_dataset)}")  # 输出：771

2. 精简列，聚焦核心信息

数据集中有 20 + 冗余列（如仓库 URL、标签链接等），仅保留 4 个核心字段：标题（title）、正文（body）、评论（comments）、原始链接（html_url）：

python

columns_to_keep = ["title", "body", "comments", "html_url"]
issues_dataset = issues_dataset.remove_columns(
    [col for col in issues_dataset.column_names if col not in columns_to_keep]
)

3. 展开评论：让每条评论成为独立搜索单元

原始数据中comments是列表格式（每个 issue 包含多条评论），需将每条评论拆分为单独行，便于后续以 “评论 + 上下文” 为单位进行检索：

python

# 转换为Pandas DataFrame，利用explode展开评论列
issues_dataset.set_format("pandas")
df = issues_dataset[:]
comments_df = df.explode("comments", ignore_index=True)  # 每条评论生成新行
# 转回Dataset格式，此时数据量从771扩展到2842条
from datasets import Dataset
comments_dataset = Dataset.from_pandas(comments_df)

4. 过滤噪声：移除无意义短评论

像 “cc @user”“+1” 这类短评论（字数≤15）对搜索无意义，添加评论字数列并过滤：

python

# 添加评论字数列（按空格拆分后的单词数）
comments_dataset = comments_dataset.map(
    lambda x: {"comment_length": len(x["comments"].split())}
)
# 保留字数>15的评论，最终得到2098条有效数据
comments_dataset = comments_dataset.filter(lambda x: x["comment_length"] > 15)

5. 拼接文本：构建完整语义上下文

为了让模型捕捉评论的完整语境，将标题、正文、评论拼接为text字段，作为嵌入生成的输入：

python

def concatenate_text(examples):
    return {
        "text": f"{examples['title']}\n{examples['body']}\n{examples['comments']}"
    }
comments_dataset = comments_dataset.map(concatenate_text)

三、创建文本嵌入：给文本装上 “数字指纹”

1. 选择适合语义搜索的预训练模型

我们选用sentence-transformers/multi-qa-mpnet-base-dot-v1模型，它专为 “短查询→长文档” 的非对称搜索场景设计，能高效捕捉跨长度文本的语义关联：

python

from transformers import AutoTokenizer, AutoModel
import torch

model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt).to(torch.device("cuda"))  # GPU加速

2. 定义嵌入生成函数（CLS 池化技术）

利用 BERT 类模型的[CLS] token 作为文本整体表示，该 token 能融合全句语义信息：

python

def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]  # 取第一个token（CLS）的向量

def get_embeddings(text_list):
    # 分词并填充到固定长度
    encoded_input = tokenizer(
        text_list, padding=True, truncation=True, max_length=512, return_tensors="pt"
    ).to(torch.device("cuda"))
    # 模型推理与池化
    model_output = model(**encoded_input)
    return cls_pooling(model_output).cpu().detach().numpy()  # 转numpy格式

3. 批量生成文本嵌入向量

通过Dataset.map()对 2098 条数据并行处理，生成 768 维的嵌入向量，存储为embeddings列：

python

comments_dataset = comments_dataset.map(
    lambda x: {"embeddings": get_embeddings([x["text"]])[0]}
)

四、FAISS 实战：从向量到高效检索的关键一跃

1. 构建 FAISS 索引：让千万向量秒级可查

FAISS 通过高效数据结构（如倒排索引、PQ 量化）加速高维向量搜索，直接调用datasets库接口即可创建索引：

python

comments_dataset.add_faiss_index(column="embeddings")  # 对embeddings列建立索引

datasets 库底层封装了 FAISS 库的核心功能，并为 Dataset 类内置了与 FAISS 集成的接口。
datasets 库的 Dataset 类提供了一系列与向量检索相关的高层接口，其中 add_faiss_index 专门用于在数据集的某一列（通常是存储嵌入向量的列，如 embeddings）上构建 FAISS 索引。
核心逻辑：
该方法会自动识别目标列（如 embeddings）中的向量数据（要求是 NumPy 数组格式，维度一致）。底层调用 FAISS 库创建对应的索引对象（如 faiss.IndexFlatL2、faiss.IndexHNSWFlat 等），并将索引与数据集绑定。

2. 示例查询：“如何离线加载数据集”

以技术场景中常见的 “离线数据加载” 问题为例，演示搜索流程：

python

def semantic_search(question, top_k=5):
    # 生成查询嵌入
    query_emb = get_embeddings([question])
    # 执行最近邻搜索，返回相似度评分和匹配样本
    scores, samples = comments_dataset.get_nearest_examples(
        "embeddings", query_emb, k=top_k
    )
    # 整理结果为DataFrame并按评分排序
    import pandas as pd
    results_df = pd.DataFrame(samples)
    results_df["similarity_score"] = scores
    return results_df.sort_values("similarity_score", ascending=False)

# 执行搜索
query = "How can I load a dataset without internet?"
results = semantic_search(query, top_k=3)

关键参数：
- column="embeddings"：指定使用哪一列的向量进行搜索（即之前创建索引的列）。
- k=top_k：返回前 k 个最相似的结果，通过调整 k 平衡精度和召回率。
返回值：
- scores：查询向量与匹配文档向量的相似度评分（值越大越相似，本文使用点积作为度量）。
- samples：匹配到的文档样本（包含标题、评论、链接等原始信息）。

3. 查看搜索结果：语义匹配 vs 关键词匹配

传统搜索若关键词不匹配会直接漏过相关内容，而语义搜索能捕捉 “offline”“without internet”“local” 等同义表达。以下是部分输出：

python

for idx, row in results.iterrows():
    print(f"【相似度】{row['similarity_score']:.2f}")
    print(f"【标题】{row['title']}")
    print(f"【核心评论】{row['comments'][:80]}...（完整内容见链接）")
    print(f"【链接】{row['html_url']}\n")

plaintext

【相似度】25.51
【标题】Discussion using datasets in offline mode
【核心评论】Requiring online connection is a deal breaker...Could you please elaborate on how that should look like?
【链接】https://github.com/huggingface/datasets/issues/824

五、避坑指南：让搜索效果提升 30% 的实用技巧

模型选择策略：
- 短查询→长文档：选multi-qa-mpnet-base-dot-v1（本文所用）；
- 文档→文档相似性：选all-MiniLM-L6-v2（轻量且对称匹配）；
- 多语言场景：用paraphrase-multilingual-MiniLM-L12-v2支持 50 + 语言。
数据清洗细节：
- 评论过滤阈值建议设为 15 词以上，太短的文本语义信息不足；
- 拼接文本时保留标题和正文作为上下文，能显著提升长文档匹配精度。
FAISS 参数调优：
- k值控制返回结果数量，默认 5 适合精准搜索，业务场景可设为 10-20 提升召回率；
- 大规模数据（>10 万条）建议使用IVFFlat等量化索引，平衡速度与精度。