让我们把之前学习的知识点结合起来,构建一个小型、实用的语义搜索应用。本练习整合一个嵌入模型、一个向量数据库客户端和一个简单的网页框架,体现一个完整但简易的搜索流程,从用户查询到相关结果。我们将使用在本地容易配置的组件,让您可以专注于它们之间的配合。在本次实践中,我们将使用:嵌入模型:使用 sentence-transformers 库,搭配 all-MiniLM-L6-v2 等预训练模型。这个模型高效,能为句子和短段落提供高质量的嵌入。向量数据库:ChromaDB。我们将使用其 Python 客户端进行本地持久化存储,这简化了本示例的配置。网页框架:FastAPI。一个现代的 Python 网页框架,易于使用并能自动生成交互式 API 文档。1. 配置与依赖首先,请确保您已安装所需的库。您可以使用 pip 进行安装:pip install sentence-transformers chromadb fastapi uvicorn[standard] python-multipart Jinja2sentence-transformers:用于加载嵌入模型和生成向量。chromadb:用于与 Chroma 向量数据库交互的客户端库。fastapi:用于创建我们 API 端点的网页框架。uvicorn:一个用于运行我们的 FastAPI 应用的 ASGI 服务器。python-multipart:FastAPI 处理表单数据所需(尽管我们可能使用 JSON)。Jinja2:FastAPI 用于可选 HTML 模板,如果需要,通常包含在 FastAPI 的依赖项中。2. 数据准备与索引脚本让我们创建一个脚本(index_data.py)来准备一些示例数据,生成嵌入,并将它们索引到 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 和配置参数。我们加载 SentenceTransformer 模型。首次运行时,它将下载模型权重。我们为 ChromaDB 初始化一个 PersistentClient,并指定一个目录(./chroma_db_persist)来存储数据库文件。这使我们的索引在多次运行后仍能持久保存。我们使用 client.get_or_create_collection。这很方便,因为它会在集合不存在时创建它,或在集合存在时加载现有集合。我们通过 embedding_function 将 SentenceTransformer 模型与该集合关联。当我们添加文档或执行查询时,ChromaDB 会自动使用此函数。我们为每个文档生成简单的唯一 ID(doc_ids)。我们使用 collection.add 将文档及其对应的 ID 添加到集合中。由于我们配置了 embedding_function,ChromaDB 会在内部调用模型,在存储每个文档之前获取其向量。我们包含了一个基本检查,以避免每次都重复索引。运行此脚本一次来填充您的本地 ChromaDB:python index_data.py您应该会看到指示初始化和成功索引的输出,并且会创建一个 chroma_db_persist 目录。3. 使用 FastAPI 构建搜索 API现在,让我们创建将处理搜索请求的网页应用(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)说明:我们初始化 FastAPI。全局资源:我们在应用启动时只进行一次加载 SentenceTransformer 模型并连接到持久化的 ChromaDB 集合。这避免了在每次请求时重新加载模型或重新连接数据库,否则会非常低效。为提高稳定性,还添加了错误处理。/search/ 端点:它接受一个 query 参数(一个字符串)。它使用与索引时相同的模型为输入的 query 生成嵌入。这对于有意义地比较向量很重要。它使用 collection.query 查找与 query_embedding 最相似的 N_RESULTS 个文档嵌入。我们请求 ChromaDB 在响应中包含原始 documents 和 distances。它将 ChromaDB 的结果格式化为更整洁的字典列表,并以 JSON 形式返回。包含错误处理以应对嵌入生成或数据库查询期间可能出现的问题。根端点:一个简单的 / 端点确认 API 正在运行。运行应用:if __name__ == "__main__": 代码块允许您直接使用 python main.py 运行服务器。另外,您也可以使用 uvicorn main:app --reload --host 0.0.0.0 --port 8000。--reload 标志在开发期间很有用,因为它会在您保存更改时自动重启服务器。4. 运行与测试索引数据:如果您尚未操作,请运行 python index_data.py。启动 API 服务器:运行 uvicorn main:app --reload --port 8000。测试 API:打开您的网页浏览器或使用 curl 等工具向搜索端点发送请求:浏览器:访问 http://localhost:8000/search/?query=what+is+AIcurl:curl "http://localhost:8000/search/?query=information%20about%20databases"curl "http://localhost:8000/search/?query=tell%20me%20about%20animals"您应该会收到 JSON 响应,其中包含来自小型数据集的最相关文档,基于语义相似性,并附带它们的距离。例如,查询“数据库”应返回与向量数据库和可能的机器学习相关的结果。查询“动物”应检索到关于狐狸的句子。搜索应用流程图digraph SemanticSearchFlow { rankdir=LR; node [shape=box, style=rounded, fontname="sans-serif", color="#495057", fontcolor="#495057"]; edge [fontname="sans-serif", color="#adb5bd", fontcolor="#495057"]; subgraph cluster_api { label = "FastAPI 应用"; bgcolor="#e9ecef"; style=filled; color="#ced4da"; api_endpoint [label="/search 端点", shape=ellipse, style=filled, fillcolor="#a5d8ff"]; query_embed [label="生成查询\n嵌入"]; db_query [label="查询向量数据库"]; format_results [label="格式化结果"]; } user [label="用户 / 客户端", shape=circle, style=filled, fillcolor="#b2f2bb"]; model [label="Sentence Transformer\n(all-MiniLM-L6-v2)", style=filled, fillcolor="#ffec99"]; vector_db [label="ChromaDB 集合\n(docs_collection)", shape=cylinder, style=filled, fillcolor="#fcc2d7"]; user -> api_endpoint [label="1. 发送查询字符串"]; api_endpoint -> query_embed [label="2. 传递查询"]; query_embed -> model [label="3. 获取嵌入"]; model -> query_embed [label="4. 返回向量"]; query_embed -> db_query [label="5. 传递向量"]; db_query -> vector_db [label="6. 执行 ANN 搜索"]; vector_db -> db_query [label="7. 返回结果 (ID, 距离, 文档)"]; db_query -> format_results [label="8. 传递结果"]; format_results -> api_endpoint [label="9. 返回格式化 JSON"]; api_endpoint -> user [label="10. 发送 JSON 响应"]; }该图说明了语义搜索应用的请求流程。用户向 API 端点发送查询,该端点使用嵌入模型将查询转换为向量。然后,此向量用于在 ChromaDB 集合中搜索相似的文档向量。结果被格式化并返回给用户。进一步思考本示例提供了一个主要结构。您可以通过多种方式进行扩展:更大规模的数据集:索引一个更大规模的数据集(例如,文章、产品描述)。请记住之前讨论过的索引策略以提高效率。元数据过滤:在索引期间添加元数据(例如,类别、时间戳),并使用 ChromaDB 的过滤功能(query 方法中的 where 子句)来精炼搜索结果。不同数据库/模型:通过根据 Pinecone、Weaviate 或 Milvus 各自的 Python 客户端修改客户端初始化和查询逻辑,将 ChromaDB 替换为它们。尝试不同的嵌入模型。用户界面:使用 Jinja2 模板与 FastAPI 或独立的网页框架(如 React、Vue)构建一个简单的 HTML 前端,使其与此 API 交互。混合搜索:将关键词搜索(例如,使用 Whoosh 或 Elasticsearch)与向量搜索结合,以可能提高相关性。评估:使用带有标签的数据集实现评估指标(如 Recall@K),以衡量搜索结果的质量。这个动手实践说明了本课程中讨论的组件(嵌入模型、向量数据库和搜索逻辑)如何结合起来,创建能够理解用户查询背后含义的应用。