跳至主要内容

如何在 MongoDB 中实现语义搜索

学习如何使用 Python 在 MongoDB 中实现语义搜索。生成向量嵌入、创建 Vector Search 索引,并运行 $vectorSearch 查询。
· 9分钟 阅读

使用 AI 探索

这正是常规文本搜索力不从心的时刻。以下是实际中的表现:

  • 客服门户在用户输入“打不开”而非帮助文章中使用的准确表述时,返回无结果。
  • 电商平台在购物者搜索“冬天保暖的东西”时找不到结果,因为没有商品描述包含这些精确词语。
  • 知识库错过了正确的文档,因为用户提问的表述与文档写法不同。

语义搜索能解决这个问题。它理解查询背后的含义与意图,而不是只匹配精确词语。

本教程将演示如何在 MongoDB 中使用 Python 和免费的nomic-embed-text-v1嵌入模型来实现语义搜索。

注意:nomic-embed-text-v1 加载进内存约 0.27 GB,因此在运行本教程中的任意脚本前,请确保您的机器至少有 1 GB 的可用内存(RAM)。首次运行会从 Hugging Face 下载模型,所需时间取决于您的网络连接。后续运行将从本地缓存加载模型。CPU 推理显著慢于 GPU 推理。在没有独立 GPU 的机器上,模型将在 CPU 上运行并使用 RAM 而非显存(VRAM),因此生成嵌入可能需要更久。

本教程您将学到:

  • 什么是向量嵌入,以及它们如何表示含义
  • 如何在 MongoDB 中生成并存储嵌入
  • 如何创建 MongoDB Vector Search 索引
  • 如何将用户查询转换为向量
  • 如何运行 $vectorSearch 聚合管道并解读结果

您可以在GitHub 仓库中找到本教程的所有代码示例。

文本搜索 vs 语义搜索

文本搜索会匹配包含查询中精确词语的文档。语义搜索则匹配与查询具有相同含义的文档,即使它们使用完全不同的词。

下表展示了针对查询“heart problems(心脏问题)”两种方式各自的返回结果:

文档文本文本搜索语义搜索文本搜索返回的原因
"...heart problems in adults..."是是包含精确词语“heart problems”
"...cardiac conditions and symptoms..."否是不包含精确词语“heart problems”
"...risk factors for heart disease..."否是不包含精确词语“heart problems”
"...chest pain and shortness of breath..."否是不包含精确词语“heart problems”

语义搜索的核心概念

在 MongoDB 中,语义搜索的三个核心概念有助于理解本教程各步骤的原理。

向量嵌入

嵌入模型将文本转换为固定长度的数字列表,即向量。向量中的每个数字代表一个意义维度。

含义相近的两段文本会生成数值上彼此接近的向量。“cardiac conditions(心脏状况)”与“heart problems(心脏问题)”在向量空间中彼此接近,尽管它们并不共享词语。这种接近性使语义搜索成为可能。

本教程使用 nomic-embed-text-v1 嵌入模型,该模型免费、开源,并完全在本地运行。首次运行脚本会自动从 Hugging Face 下载模型,后续运行从本地缓存加载。该模型每个输入会生成 768 个数值。

向量搜索索引

向量搜索索引会告知 MongoDB 哪个字段保存嵌入、应期望多少维度,以及用于比较的相似度函数。在运行任何 $vectorSearch 查询之前,您必须先创建此索引。

$vectorSearch 聚合阶段

$vectorSearch 是执行搜索的聚合阶段。它接收查询向量,搜索已索引的字段,并按语义相似度对文档进行排序返回。您可以像其他聚合管道一样将其与 $project 等阶段串联。

现在开始动手吧。

前置条件

开始之前,请确保已准备好以下内容:

  • 已安装 Python 3.8 或更高版本
  • 已设置 M0(免费层)集群的 MongoDB Atlas 账户
  • pymongo 4.7 或更高版本、sentence-transformers 和 einops Python 包已安装
  • 具备 Python 和 MongoDB 集合的基础知识
  • 您的 Atlas 连接字符串,可在 Atlas UI 的 **Database > Connect > Drivers** 下获取
  • 在运行脚本前,已在 Atlas 中放行您的 IP 地址。请在 Atlas UI 的 **Security > Network Access** 中添加您当前的 IP 地址。

设置您的项目

为您的项目创建一个文件夹,并在终端中进入该目录:

mkdir mongodb-semantic-search
cd mongodb-semantic-search

创建并激活虚拟环境,以隔离项目依赖:

python -m venv venv
source venv/bin/activate

在 Windows 上,使用以下命令激活虚拟环境:

venv\Scripts\activate

现在安装所需的包:

pip install pymongo sentence-transformers einops

您将为本教程的每个步骤创建一个 Python 文件。所有文件都放在 mongodb-semantic-search 文件夹中。

创建嵌入工具文件

创建名为 embedding_utils.py 的文件。该文件包含本教程中使用的两个嵌入函数:

from sentence_transformers import SentenceTransformer

# Load the free, open-source embedding model.
# The model downloads from Hugging Face on first run and saves locally.
# trust_remote_code=True is required by this model:
model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)


def get_embedding(text, precision="float32"):
    # Use this function when embedding text you plan to store in MongoDB:
    return model.encode(text, precision=precision).tolist()


def get_query_embedding(text, precision="float32"):
    # Use this function when embedding a user's search query:
    return model.encode(text, precision=precision).tolist()

警告:trust_remote_code=True 允许从 Hugging Face 下载的模型特定 Python 代码在您的机器上执行。若 Hugging Face 上的模型源代码被入侵或被更新为恶意代码,这些代码将自动在您的环境中运行。在生产部署中使用该参数前请务必谨慎。

生成并存储向量嵌入

创建名为 generate_embeddings.py 的文件。该文件从 embedding_utils.py 导入嵌入函数,为每个文档的 text 字段生成向量,并将包含嵌入的完整文档存入 MongoDB:

from embedding_utils import get_embedding
from pymongo import MongoClient

# Replace the placeholder with your Atlas connection string:
mongodb_client = MongoClient(
    "mongodb+srv://<USERNAME>:<PASSWORD>@<HOST>/",
    appname="devrel-tutorial-python-semantic-search"
)

collection = mongodb_client["sample_db"]["documents"]

# Sample data:
sample_documents = [
    {
        "title": "MongoDB Atlas",
        "text": "MongoDB Atlas is a fully managed cloud database."
    },
    {
        "title": "Vector Search",
        "text": "Vector search finds results based on semantic meaning."
    },
    {
        "title": "Nomic AI",
        "text": "nomic-embed-text-v1 is a free, open-source embedding model."
    },
]

docs_to_insert = []

for doc in sample_documents:
    embedding = get_embedding(doc["text"])
    docs_to_insert.append({
        "title": doc["title"],
        "text": doc["text"],
        # The vector lives alongside your original data in the same document:
        "embedding": embedding
    })

# Drop the collection before each run to avoid inserting duplicate documents:
collection.drop()

result = collection.insert_many(docs_to_insert)
print(f"Inserted {len(result.inserted_ids)} documents with embeddings.")

mongodb_client.close()

在上述代码中,collection.drop() 会在每次运行前删除集合中的所有文档和索引。若不加此行重复运行 generate_embeddings.py,会插入重复文档,导致 $vectorSearch 多次返回同一文档。丢弃集合也会删除向量搜索索引,因此每次重新运行 generate_embeddings.py 后都必须重新运行 create_vector_index.py。

运行脚本:

python generate_embeddings.py

您将在终端看到如下输出:

<All keys matched successfully>
Inserted 3 documents with embeddings.

现在 MongoDB 中的每个文档都在 embedding 字段中同时保存了原始文本及其 768 维向量。该向量与数据存放在同一文档中,因此查询时无需单独存储或查找。

创建 MongoDB 向量搜索索引

创建名为 create_vector_index.py 的文件。该文件在 embedding 字段上定义并创建向量搜索索引,使 MongoDB 能对您的集合运行 $vectorSearch 查询。

索引定义需要三个字段:

  • path:保存嵌入的字段(本教程为 embedding)
  • numDimensions:必须与模型的输出大小一致(nomic-embed-text-v1 为 768)
  • similarity:比较函数(对于该模型,cosine 正确)

numDimensions 的取值必须与您的嵌入模型完全一致。若不匹配,索引创建将失败:

from pymongo.mongo_client import MongoClient
from pymongo.operations import SearchIndexModel
import time

# Replace the placeholder with your Atlas connection string:
mongodb_client = MongoClient(
    "mongodb+srv://<USERNAME>:<PASSWORD>@<HOST>/",
    appname="devrel-tutorial-python-semantic-search"
)

# Point to the same database and collection you used in the previous step:
database = mongodb_client["sample_db"]
collection = database["documents"]

# Define the vector search index.
# The three required fields tell MongoDB what to index and how to compare vectors:
search_index_model = SearchIndexModel(
    definition={
        "fields": [
            {
                "type": "vector",
                "path": "embedding",      # The field name from generate_embeddings.py
                "numDimensions": 768,     # nomic-embed-text-v1 always outputs 768 dimensions
                "similarity": "cosine"    # Recommended similarity function for this model
            }
        ]
    },
    name="vector_index",
    type="vectorSearch"
)

result = collection.create_search_index(model=search_index_model)
print("New search index named " + result + " is building.")

# Poll every five seconds until the index is ready to accept queries:
print("Polling to check if the index is ready. This may take up to a minute.")

predicate = lambda index: index.get("queryable") is True

while True:
    indices = list(collection.list_search_indexes(result))
    if len(indices) and predicate(indices[0]):
        break
    time.sleep(5)

print(result + " is ready for querying.")

mongodb_client.close()

运行脚本:

python create_vector_index.py

您将在终端看到如下输出:

New search index named vector_index is building.
Polling to check if the index is ready. This may take up to a minute.
vector_index is ready for querying.

索引的构建最多可能需要一分钟。轮询循环每 5 秒检查一次,仅当 MongoDB 确认索引可查询时才会退出。在看到 “vector_index is ready for querying.” 之前请勿进行下一步。

将搜索查询转换为向量

创建名为 generate_query_vector.py 的文件。该文件从 embedding_utils.py 导入 get_query_embedding() 函数。运行 generate_query_vector.py 以确认嵌入模型能正确加载并生成 768 维向量:

  • 它定义了下一步将导入的 get_query_embedding() 函数。
  • 您也可单独运行它,验证模型工作是否正常。

此处必须与第一步使用相同的模型。不同模型会生成处在不同数值空间的向量,导致比较毫无意义:

from embedding_utils import get_query_embedding

user_query = "I need an automated, scalable system for serious information storage"
query_vector = get_query_embedding(user_query)
print(f"Query: '{user_query}'")
print(f"Vector dimensions: {len(query_vector)}")
print(f"First 5 values: {query_vector[:5]}")

运行脚本:

python generate_query_vector.py

您将看到类似如下的输出:

<All keys matched successfully>
Query: 'I need an automated, scalable system for serious information storage'
Vector dimensions: 768
First 5 values: [0.0025914530269801617, 0.09862980246543884, -0.023092379793524742, -0.0171672236174345, -0.05548065900802612]

在上述输出中,Vector dimensions: 768 说明模型输出与已存储的嵌入维度一致。

运行 $vectorSearch 查询

创建名为 run_vector_search.py 的文件。该文件从 embedding_utils.py 导入 get_query_embedding() 函数并将用户的搜索词转换为向量。run_vector_search.py 运行一个 $vectorSearch 聚合管道,并按语义相似度输出结果:

from embedding_utils import get_query_embedding
from pymongo import MongoClient

# Replace the placeholder with your Atlas connection string:
mongodb_client = MongoClient(
    "mongodb+srv://<USERNAME>:<PASSWORD>@<HOST>/",
    appname="devrel-tutorial-python-semantic-search"
)

collection = mongodb_client["sample_db"]["documents"]

# Generate a query vector from the user's search input:
user_query = "I need an automated, scalable system for serious information storage"
query_vector = get_query_embedding(user_query)

# Define the $vectorSearch aggregation pipeline:
pipeline = [
    {
        "$vectorSearch": {
            "index": "vector_index",      # The index created in create_vector_index.py
            "path": "embedding",          # The field that holds your stored vectors
            "queryVector": query_vector,  # The vector generated from the user's query
            "numCandidates": 150,         # How many neighbors MongoDB considers
            "limit": 3                    # How many results to return
        }
    },
    {
        "$project": {
            "_id": 0,
            "title": 1,
            "text": 1,
            "score": {
                "$meta": "vectorSearchScore"  # Relevance score for each result
            }
        }
    }
]

results = collection.aggregate(pipeline)

print(f"\nTop results for query: '{user_query}'\n")

for doc in results:
    print(f"Title: {doc['title']}")
    print(f"Text:  {doc['text']}")
    print(f"Score: {doc['score']:.4f}")
    print()

mongodb_client.close()

在上述代码中,有两个参数控制搜索行为:

  • numCandidates 设置 MongoDB 在收敛到最终结果前所考察的初始候选池大小。取值越大,召回率越高,但耗时略增。
  • limit 设置返回结果的数量。本教程中,limit 设为 3,numCandidates 设为 150。常见的起点是将 numCandidates 设为 limit 的 10–15 倍。

现在运行脚本:

python run_vector_search.py

您将看到类似如下的输出:

<All keys matched successfully>

Top results for query: 'I need an automated, scalable system for serious information storage.'

Title: MongoDB Atlas
Text:  MongoDB Atlas is a fully managed cloud database.
Score: 0.7211

Title: Nomic AI
Text:  nomic-embed-text-v1 is a free, open-source embedding model.
Score: 0.6818

Title: Vector Search
Text:  Vector search finds results based on semantic meaning.
Score: 0.6642

即便查询 “I need an automated, scalable system for serious information storage” 与 “MongoDB Atlas is a fully managed cloud database.” 并无词语重合,“MongoDB Atlas” 仍获得最高分。之所以返回该结果,是因为它们的向量在含义上很接近:“……自动化、可扩展的信息存储系统” 在语义上映射到 “……全托管的云数据库”。对同一查询进行文本搜索将返回零结果,因为没有任何文档包含查询中的精确词语。这正是语义搜索应有的效果。

要点回顾

  • MongoDB Vector Search 可在 M0 免费层集群上运行,无需付费计划。
  • 为存储嵌入与为查询生成向量必须使用同一嵌入模型。混用模型会导致结果毫无意义。
  • numDimensions 在索引定义中的取值必须与嵌入模型的输出尺寸完全一致。nomic-embed-text-v1 始终生成 768 维。
  • input_type="document" 可为存储优化嵌入,input_type="query" 则为检索优化。请在各阶段使用正确类型。
  • numCandidates 控制 MongoDB 在收敛到最终 limit 结果前所覆盖的搜索范围。更大的取值可提升召回,但会增加查询时间。
  • vectorSearchScore 按语义相似度对结果排序。结果无需包含查询中的任何精确词语。分数范围因模型与数据集而异,但以下边界可作为 nomic-embed-text-v1 配合余弦相似度的参考起点:
    • 0.9 及以上:几乎完全相同的含义。文档与查询在语义上几乎一致。
    • 0.7 到 0.9:高度相关。文档与查询意图明显相关。
    • 0.5 到 0.7:中等相关。文档与主题相关,但表述或上下文不同。
    • 低于 0.5:相关性较弱。关联度低,结果可能不够有用。

您可以在GitHub 仓库中找到本教程的所有代码示例。

延伸阅读

常见问题

实现语义搜索需要付费的 Atlas 计划吗?

不需要。本教程的四个步骤均可在免费的 M0 层集群上运行。

如果查询使用的嵌入模型与存储文档使用的不同,会发生什么?

结果将毫无意义。不同模型会将文本映射到不同的数值空间,向量不可比较。请始终在索引与查询时使用同一模型。

`numCandidates` 与 `limit` 有何区别?

numCandidates 表示 MongoDB 在搜索过程中考察的向量数量;limit 表示返回的前若干条结果。更高的 numCandidates 值可提升结果质量,但查询会稍慢。常见的起点是将 numCandidates 设为 limit 的 10 到 15 倍。

我可以将语义搜索与文本搜索一起使用吗?

可以。MongoDB 支持混合搜索,可在同一管道中组合 $vectorSearch 与 $search。

语义搜索能用于英语之外的语言吗?

取决于您的嵌入模型。nomic-embed-text-v1 主要基于英文文本训练。若涉及多语言用例,请选择在您的数据所含语言上训练的多语言嵌入模型。

主题