趋近智
让我们把之前学习的知识点结合起来,构建一个小型、实用的语义搜索应用。本练习整合一个嵌入 (embedding)模型、一个向量 (vector)数据库客户端和一个简单的网页框架,体现一个完整但简易的搜索流程,从用户查询到相关结果。我们将使用在本地容易配置的组件,让您可以专注于它们之间的配合。
在本次实践中,我们将使用:
sentence-transformers 库,搭配 all-MiniLM-L6-v2 等预训练 (pre-training)模型。这个模型高效,能为句子和短段落提供高质量的嵌入。ChromaDB。我们将使用其 Python 客户端进行本地持久化存储,这简化了本示例的配置。FastAPI。一个现代的 Python 网页框架,易于使用并能自动生成交互式 API 文档。首先,请确保您已安装所需的库。您可以使用 pip 进行安装:
pip install sentence-transformers chromadb fastapi uvicorn[standard] python-multipart Jinja2
sentence-transformers:用于加载嵌入 (embedding)模型和生成向量 (vector)。chromadb:用于与 Chroma 向量数据库交互的客户端库。fastapi:用于创建我们 API 端点的网页框架。uvicorn:一个用于运行我们的 FastAPI 应用的 ASGI 服务器。python-multipart:FastAPI 处理表单数据所需(尽管我们可能使用 JSON)。Jinja2:FastAPI 用于可选 HTML 模板,如果需要,通常包含在 FastAPI 的依赖项中。让我们创建一个脚本(index_data.py)来准备一些示例数据,生成嵌入 (embedding),并将它们索引到 ChromaDB 中。
# index_data.py
import chromadb
from sentence_transformers import SentenceTransformer
# --- Configuration ---
MODEL_NAME = 'all-MiniLM-L6-v2'
COLLECTION_NAME = "docs_collection"
PERSIST_DIRECTORY = "./chroma_db_persist" # 存储数据库数据的目录
# --- Sample Data ---
# 文档的简单列表(此处为句子)
documents = [
"The quick brown fox jumps over the lazy dog.", # 敏捷的棕色狐狸跳过懒惰的狗。
"Artificial intelligence is transforming many industries.", # 人工智能正在改变许多行业。
"Vector databases are optimized for similarity search.", # 向量数据库为相似性搜索而优化。
"Natural language processing enables computers to understand text.", # 自然语言处理使计算机能够理解文本。
"The capital of France is Paris.", # 法国的首都是巴黎。
"Apples are a type of fruit, often red or green.", # 苹果是一种水果,通常是红色或绿色。
"Machine learning algorithms learn from data.", # 机器学习算法从数据中学习。
"Semantic search provides results based on meaning, not just keywords.", # 语义搜索提供基于含义而非仅基于关键词的结果。
]
# --- Initialization ---
print("Initializing embedding model...") # 正在初始化嵌入模型...
# 加载预训练的句子 Transformer 模型
# 这个模型将句子和段落映射到 384 维稠密向量空间
# 如果模型不存在,它将自动下载
model = SentenceTransformer(MODEL_NAME)
print("Initializing ChromaDB client...") # 正在初始化 ChromaDB 客户端...
# 初始化具有持久化功能的 ChromaDB 客户端
# 这会将数据库状态保存到指定目录
client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
print(f"Getting or creating collection: {COLLECTION_NAME}") # 正在获取或创建集合:{COLLECTION_NAME}
# 获取或创建集合。如果存在,则加载它。
# 根据我们的 SentenceTransformer 模型指定嵌入函数
collection = client.get_or_create_collection(
name=COLLECTION_NAME,
embedding_function=chromadb.utils.embedding_functions.SentenceTransformerEmbeddingFunction(model_name=MODEL_NAME)
# 如果需要,您也可以显式传递 metadata={'hnsw:space': 'cosine'},
# 但 SentenceTransformerEmbeddingFunction 通常会默认选择合适的。
)
# --- Indexing ---
print("Generating IDs and preparing data for indexing...") # 正在生成 ID 并准备数据进行索引...
# 为本示例生成简单的顺序 ID
doc_ids = [f"doc_{i}" for i in range(len(documents))]
# 检查数据是否需要索引(基于预期数量的简单检查)
# 在实际应用中,您可能需要更完善的方式来跟踪已索引数据
if collection.count() < len(documents):
print(f"Indexing {len(documents)} documents...") # 正在索引 {len(documents)} 个文档...
try:
# 将文档添加到集合
# ChromaDB 的 SentenceTransformerEmbeddingFunction 在此处自动处理嵌入生成
collection.add(
documents=documents,
ids=doc_ids
)
print("Documents indexed successfully.") # 文档索引成功。
except Exception as e:
print(f"Error indexing documents: {e}") # 索引文档时出错:{e}
else:
print("Documents seem to be already indexed.") # 文档似乎已经索引。
print(f"Collection '{COLLECTION_NAME}' now contains {collection.count()} documents.") # 集合 '{COLLECTION_NAME}' 现在包含 {collection.count()} 个文档。
print("Indexing script finished.") # 索引脚本完成。
说明:
documents 和配置参数 (parameter)。SentenceTransformer 模型。首次运行时,它将下载模型权重 (weight)。PersistentClient,并指定一个目录(./chroma_db_persist)来存储数据库文件。这使我们的索引在多次运行后仍能持久保存。client.get_or_create_collection。这很方便,因为它会在集合不存在时创建它,或在集合存在时加载现有集合。我们通过 embedding_function 将 SentenceTransformer 模型与该集合关联。当我们添加文档或执行查询时,ChromaDB 会自动使用此函数。doc_ids)。collection.add 将文档及其对应的 ID 添加到集合中。由于我们配置了 embedding_function,ChromaDB 会在内部调用模型,在存储每个文档之前获取其向量 (vector)。我们包含了一个基本检查,以避免每次都重复索引。运行此脚本一次来填充您的本地 ChromaDB:
python index_data.py
您应该会看到指示初始化和成功索引的输出,并且会创建一个 chroma_db_persist 目录。
现在,让我们创建将处理搜索请求的网页应用(main.py)。
# main.py
import chromadb
from fastapi import FastAPI, Query, HTTPException
from sentence_transformers import SentenceTransformer
import uvicorn # 用于运行应用
# --- Configuration ---
MODEL_NAME = 'all-MiniLM-L6-v2'
COLLECTION_NAME = "docs_collection"
PERSIST_DIRECTORY = "./chroma_db_persist"
N_RESULTS = 3 # 返回的搜索结果数量
# --- Application Initialization ---
app = FastAPI(
title="Simple Semantic Search API", # 简单语义搜索 API
description="An API that uses a vector database for semantic search.", # 一个使用向量数据库进行语义搜索的 API。
version="0.1.0"
)
# --- Global Variables / Resources ---
# 应用启动时一次性初始化资源
try:
print("Loading embedding model...") # 正在加载嵌入模型...
embedding_model = SentenceTransformer(MODEL_NAME)
print("Model loaded successfully.") # 模型加载成功。
print("Connecting to ChromaDB...") # 正在连接到 ChromaDB...
db_client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
collection = db_client.get_collection(name=COLLECTION_NAME)
# 验证集合中是否有项目(可选但良好实践)
if collection.count() == 0:
print(f"Warning: Collection '{COLLECTION_NAME}' is empty. Did you run index_data.py?") # 警告:集合 '{COLLECTION_NAME}' 为空。您运行过 index_data.py 吗?
print("ChromaDB connection successful.") # ChromaDB 连接成功。
except Exception as e:
print(f"Error during initialization: {e}") # 初始化时出错:{e}
# 适当处理初始化失败,可能是退出或抛出特定错误
embedding_model = None
collection = None
# --- API Endpoints ---
@app.get("/search/")
async def perform_search(
query: str = Query(..., min_length=3, description="The search query text.") # 搜索查询文本。
):
"""
对已索引文档执行语义搜索。
接收一个查询字符串,生成其嵌入,并搜索向量
数据库以查找最相似的文档。
"""
if not embedding_model or not collection:
raise HTTPException(status_code=503, detail="Search service is not available due to initialization error.") # 由于初始化错误,搜索服务不可用。
print(f"Received query: '{query}'") # 收到查询:'{query}'
try:
# 1. 为查询生成嵌入
print("Generating query embedding...") # 正在生成查询嵌入...
query_embedding = embedding_model.encode(query).tolist()
print("Query embedding generated.") # 查询嵌入已生成。
# 2. 查询向量数据库
print(f"Querying collection '{COLLECTION_NAME}'...") # 正在查询集合 '{COLLECTION_NAME}'...
results = collection.query(
query_embeddings=[query_embedding], # 注意:query_embeddings 期望一个嵌入列表
n_results=N_RESULTS,
include=['documents', 'distances'] # 请求 ChromaDB 返回文档和距离
)
print("Query executed successfully.") # 查询执行成功。
# 3. 格式化并返回结果
# 结果结构可能有点嵌套,我们来简化它
if results and results.get('ids') and results['ids'][0]:
formatted_results = []
ids = results['ids'][0]
distances = results['distances'][0]
documents = results['documents'][0]
for i in range(len(ids)):
formatted_results.append({
"id": ids[i],
"document": documents[i],
"distance": distances[i] # 距离越小意味着对于余弦/欧几里得距离越相似
})
return {"results": formatted_results}
else:
return {"results": []} # 如果未找到结果,返回空列表
except Exception as e:
print(f"Error during search for query '{query}': {e}") # 查询 '{query}' 时搜索出错:{e}
raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}") # 搜索失败:{str(e)}
@app.get("/")
async def read_root():
""" 一个简单的根端点,用于检查 API 是否正在运行。 """
return {"message": "Semantic Search API is running. Use the /search/ endpoint."} # 语义搜索 API 正在运行。请使用 /search/ 端点。
# --- Main Execution ---
# 此代码块允许直接使用 `python main.py` 运行应用
if __name__ == "__main__":
print("Starting FastAPI server...") # 正在启动 FastAPI 服务器...
uvicorn.run(app, host="0.0.0.0", port=8000)
说明:
SentenceTransformer 模型并连接到持久化的 ChromaDB 集合。这避免了在每次请求时重新加载模型或重新连接数据库,否则会非常低效。为提高稳定性,还添加了错误处理。query 参数 (parameter)(一个字符串)。query 生成嵌入 (embedding)。这对于有意义地比较向量 (vector)很重要。collection.query 查找与 query_embedding 最相似的 N_RESULTS 个文档嵌入。我们请求 ChromaDB 在响应中包含原始 documents 和 distances。/ 端点确认 API 正在运行。if __name__ == "__main__": 代码块允许您直接使用 python main.py 运行服务器。另外,您也可以使用 uvicorn main:app --reload --host 0.0.0.0 --port 8000。--reload 标志在开发期间很有用,因为它会在您保存更改时自动重启服务器。索引数据:如果您尚未操作,请运行 python index_data.py。
启动 API 服务器:运行 uvicorn main:app --reload --port 8000。
测试 API:打开您的网页浏览器或使用 curl 等工具向搜索端点发送请求:
http://localhost:8000/search/?query=what+is+AIcurl "http://localhost:8000/search/?query=information%20about%20databases"
curl "http://localhost:8000/search/?query=tell%20me%20about%20animals"
您应该会收到 JSON 响应,其中包含来自小型数据集的最相关文档,基于语义相似性,并附带它们的距离。例如,查询“数据库”应返回与向量 (vector)数据库和可能的机器学习 (machine learning)相关的结果。查询“动物”应检索到关于狐狸的句子。
该图说明了语义搜索应用的请求流程。用户向 API 端点发送查询,该端点使用嵌入 (embedding)模型将查询转换为向量 (vector)。然后,此向量用于在 ChromaDB 集合中搜索相似的文档向量。结果被格式化并返回给用户。
本示例提供了一个主要结构。您可以通过多种方式进行扩展:
query 方法中的 where 子句)来精炼搜索结果。这个动手实践说明了本课程中讨论的组件(嵌入模型、向量数据库和搜索逻辑)如何结合起来,创建能够理解用户查询背后含义的应用。
这部分内容有帮助吗?
© 2026 ApX Machine LearningAI伦理与透明度•