AI Engineering Academy · 课时

使用 NumPy 进行语义搜索

使用 NumPy 构建纯 Python 语义搜索系统,计算查询嵌入与文档嵌入集合之间的余弦相似度。

第 3 / 4 课13 个步骤

使用 NumPy 进行语义搜索 是 CoddyKit 上的免费 AI Engineering Academy 课时。 这是第 3 节课,共 4 节。 你可以在下方免费阅读本课时的完整内容 — 然后在浏览器中使用内置代码编辑器和全天候 AI 导师进行实践。 这是 AI Engineering Academy 学习路径的一部分,你的进度在网页和 CoddyKit 应用中同步。 AI Engineering Academy 课程共包含 4 节课。

无需数据库的语义搜索

语义搜索根据含义而不是关键词重合程度,查找与查询最相关的文档。最简单的实现方式是使用 NumPy,在内存中计算查询嵌入与所有文档嵌入之间的余弦相似度——无需外部数据库。

这种方法适用于最多数万份文档,也非常适合在投入建设向量数据库之前进行原型开发。

构建文档语料库

首先收集文档,并使用 OpenAI API 为每份文档生成一个嵌入。将嵌入存储为二维 NumPy 数组,其中每一行代表一个文档向量。同时保留一个与之对应的文档文本列表,这样找到最佳匹配项后,就能取回原始内容。

import numpy as np
from openai import OpenAI

client = OpenAI()

documents = [
    'Python is a high-level programming language.',
    'NumPy provides fast numerical computing for Python.',
    'Embeddings represent text as dense vectors.',
    'Cosine similarity measures angle between vectors.',
    'RAG combines retrieval with language generation.'
]

response = client.embeddings.create(
    model='text-embedding-3-small',
    input=documents
)

corpus_embeddings = np.array([item.embedding for item in response.data])
print(f'Corpus shape: {corpus_embeddings.shape}')  # (5, 1536)

为用户查询生成嵌入

用户提交搜索查询时,请使用与文档相同的模型为查询生成嵌入。混用模型——例如使用 text-embedding-3-small 处理文档、使用 text-embedding-3-large 处理查询——会使向量处于不同的空间中,从而产生毫无意义的相似度分数。

from openai import OpenAI
import numpy as np

client = OpenAI()

query = 'How do I compute similarity between text?'

response = client.embeddings.create(
    model='text-embedding-3-small',   # must match corpus model
    input=query
)

query_embedding = np.array(response.data[0].embedding)
print(f'Query vector shape: {query_embedding.shape}')  # (1536,)

使用 NumPy 计算余弦相似度

要在一次操作中计算查询与每份文档之间的相似度,请计算查询向量与文档向量矩阵的点积。由于 OpenAI 嵌入都经过 L2 归一化,这对所有文档同时计算时就等于余弦相似度——复杂度为 O(n * d),其中 n 是文档数量,d 是维度。

import numpy as np

# corpus_embeddings: (n_docs, 1536)
# query_embedding: (1536,)

def semantic_search(query_vec, corpus_vecs):
    # Matrix-vector dot product: shape (n_docs,)
    similarities = corpus_vecs @ query_vec
    return similarities

# Example call (assuming pre-computed embeddings)
# sims = semantic_search(query_embedding, corpus_embeddings)
# print(sims)  # array of similarity scores, one per document

对结果排序并检索前 K 项

使用 np.argsort 按相似度分数降序排列文档,然后切片取出前 k 个索引。这样即可获得最相关文档的索引,再使用这些索引从对应的文档列表中查找原始文本。

import numpy as np

def get_top_k(query_vec, corpus_vecs, documents, k=3):
    similarities = corpus_vecs @ query_vec
    # argsort gives ascending order; [::-1] reverses to descending
    ranked_indices = np.argsort(similarities)[::-1]
    top_k_indices = ranked_indices[:k]
    return [
        {'text': documents[i], 'score': float(similarities[i])}
        for i in top_k_indices
    ]

# results = get_top_k(query_embedding, corpus_embeddings, documents, k=3)
# for r in results:
#     print(f'{r["score"]:.4f}: {r["text"]}')

完整的语义搜索示例

将所有步骤组合起来:为语料库生成嵌入、为查询生成嵌入、计算相似度,并返回排序后的结果。这一完整模式是每个 RAG 检索步骤的核心,即使底层使用向量数据库替代 NumPy,也是如此。

import numpy as np
from openai import OpenAI

client = OpenAI()

docs = [
    'Embeddings map text to numerical vectors.',
    'Python lists store ordered collections.',
    'Cosine similarity compares vector directions.',
    'RAG retrieves documents to ground LLM answers.',
    'Dictionaries store key-value pairs in Python.'
]

corpus_resp = client.embeddings.create(model='text-embedding-3-small', input=docs)
corpus = np.array([d.embedding for d in corpus_resp.data])

query = 'finding similar text using angles'
q_resp = client.embeddings.create(model='text-embedding-3-small', input=query)
q_vec = np.array(q_resp.data[0].embedding)

scores = corpus @ q_vec
for i in np.argsort(scores)[::-1][:3]:
    print(f'{scores[i]:.3f}: {docs[i]}')

设置分数阈值

前 k 个结果并不一定都相关——有时最佳匹配在语义上仍然不够合适。请添加分数阈值,过滤掉相似度较低的结果。余弦相似度常用的阈值是 0.70–0.80,但您应使用真实查询,结合具体领域进行校准。

import numpy as np

def search_with_threshold(query_vec, corpus_vecs, documents, k=5, threshold=0.75):
    similarities = corpus_vecs @ query_vec
    ranked = np.argsort(similarities)[::-1][:k]
    results = []
    for i in ranked:
        if similarities[i] >= threshold:
            results.append({'text': documents[i], 'score': float(similarities[i])})
    return results

# Only returns documents above the minimum similarity threshold

NumPy 搜索的性能特征

NumPy 相似度搜索每次查询的时间复杂度为 O(n * d),其中 n 是文档数量,d 是嵌入维度。对于 1536 维嵌入:

  • 10,000 份文档:在现代 CPU 上每次查询约 5 毫秒
  • 100,000 份文档:每次查询约 50 毫秒
  • 1,000,000 份文档:约 500 毫秒,速度过慢,应切换到向量数据库

NumPy 非常适合原型开发,但内置功能不支持近似搜索、筛选或持久化。

将嵌入持久化到磁盘

每次运行都重新计算嵌入会浪费 API 调用次数和资金。请将语料库嵌入和文档文本保存到磁盘,这样只有在语料库发生变化时才需要重新生成嵌入。

np.save 可以高效地存储嵌入矩阵,您还可以将文档列表保存为 JSON。启动时加载这两个文件,而不是调用 API。

import numpy as np
import json

# Save
np.save('/tmp/corpus_embeddings.npy', corpus_embeddings)
with open('/tmp/corpus_docs.json', 'w') as f:
    json.dump(documents, f)

# Load
corpus_embeddings = np.load('/tmp/corpus_embeddings.npy')
with open('/tmp/corpus_docs.json') as f:
    documents = json.load(f)

print(f'Loaded {len(documents)} docs, shape {corpus_embeddings.shape}')

增量处理新文档

有新文档到达时,无需为整个语料库重新生成嵌入。只需为新文档生成嵌入,并使用 np.vstack 将其向量追加到现有矩阵中。请记得按相同顺序将新文本追加到文档列表中。

import numpy as np
from openai import OpenAI

client = OpenAI()

# Assume these exist from a previous session:
# corpus_embeddings: (n, 1536)
# documents: list of strings

new_docs = ['New document about vector search.']
resp = client.embeddings.create(model='text-embedding-3-small', input=new_docs)
new_vecs = np.array([item.embedding for item in resp.data])

corpus_embeddings = np.vstack([corpus_embeddings, new_vecs])
documents.extend(new_docs)
print(f'Corpus now has {len(documents)} documents')

内存搜索的局限性

与专用向量数据库相比,NumPy 语义搜索存在明显局限:

  • 不支持持久化——所有内容都存放在 RAM 中,重启后会丢失
  • 不支持元数据筛选——无法按日期、类别或作者筛选结果
  • 仅支持线性扫描——没有近似最近邻索引
  • 不支持并发访问——不适合多用户生产部署

这些局限说明,在生产环境的 RAG 系统中应使用专用向量数据库。

快速检查

测试您对本课中 AI 工程概念的理解。

课程回顾

在本课中,您学到了:针对 L2 归一化嵌入计算矩阵点积,可以通过一次操作为整个语料库计算余弦相似度;反转 np.argsort 的结果即可检索最相似的前 k 个文档;以及NumPy 搜索非常适合原型开发,但不支持持久化和元数据筛选。接下来,我们将使用聚类和 UMAP,发现嵌入集合中的主题结构。

免费开始

用 AI 导师学习 Python — 免费

在浏览器中编写并运行真实代码,获得全天候 AI 导师的即时帮助,并在网页或应用中继续学习。

课程
30
课程
120

常见问题解答

「使用 NumPy 进行语义搜索」课时是免费的吗?

是的 — 「使用 NumPy 进行语义搜索」的完整文本可在网页上免费阅读。要进行交互式练习(内置代码编辑器和全天候 AI 导师)并解锁 AI Engineering Academy 课程的其余内容,请升级到 CoddyKit PRO。 AI Engineering Academy 课程共包含 4 节课。

「使用 NumPy 进行语义搜索」这节课中我会学到什么?

使用 NumPy 构建纯 Python 语义搜索系统,计算查询嵌入与文档嵌入集合之间的余弦相似度。 你通过在浏览器中直接运行的动手代码来练习 AI Engineering Academy,全天候 AI 导师会在你学习这节课的过程中回答你的问题。

学习 AI Engineering Academy 需要有经验吗?

无需任何先前经验。CoddyKit 上的 AI Engineering Academy 课程适合初学者到高级学习者,你可以从这里开始或从头开始,按照自己的节奏学习。 这是第 3 节课,共 4 节。

「使用 NumPy 进行语义搜索」课时需要多长时间?

大多数 CoddyKit 课程大约需要 5–10 分钟。每节课都很精短且互动,所以你能稳步进步,并在网页和应用中从离开的地方继续。

我能在这节 AI Engineering Academy 课中编写并运行代码吗?

能。每节 AI Engineering Academy 课都包含内置代码编辑器,你可以在浏览器中直接编写并运行真实代码,并获得即时 AI 反馈 — 无需本地设置。

此课程中的所有课时

  1. 什么是向量嵌入?
  2. 使用 OpenAI 生成嵌入
  3. 使用 NumPy 进行语义搜索
  4. 聚类并可视化嵌入
← 返回 AI Engineering Academy