下载模型gte-large-zhfrompathlibimportPathfromhuggingface_hubimportsnapshot_download# 当前路径是 notebook/rag因此 parents[1] 是项目根目录project_rootPath.cwd().parents[1]model_path(project_root/llm_app/embedding_models/gte-large-zh)snapshot_download(repo_idthenlper/gte-large-zh,local_dirstr(model_path))print(模型保存位置,model_path)下载 gte-large-zh → 将文档转换成向量 → 保存到 Chroma → 根据问题检索相似文档 → 按元数据过滤。frompathlibimportPathimportnumpyasnpimportchromadbfromhuggingface_hubimportsnapshot_downloadfromchromadb.utils.embedding_functionsimport(SentenceTransformerEmbeddingFunction)# # 1. 设置项目、模型和数据库路径# project_rootPath(rF:\MK-RAG\rag_full_stack_course_notebooks)model_path(project_root/llm_app/embedding_models/gte-large-zh)chroma_data_path(project_root/data/chroma_db)model_path.parent.mkdir(parentsTrue,exist_okTrue)chroma_data_path.mkdir(parentsTrue,exist_okTrue)print(项目路径,project_root)print(模型路径,model_path)print(数据库路径,chroma_data_path)# # 2. 检查并下载 gte-large-zh# weight_exists((model_path/model.safetensors).is_file()or(model_path/pytorch_model.bin).is_file())model_complete((model_path/config.json).is_file()and(model_path/modules.json).is_file()andweight_exists)ifnotmodel_complete:print(本地模型不完整开始下载 gte-large-zh……)snapshot_download(repo_idthenlper/gte-large-zh,local_dirstr(model_path))print(模型下载完成。)else:print(本地模型已经存在不需要重新下载。)# 再次检查模型weight_exists((model_path/model.safetensors).is_file()or(model_path/pytorch_model.bin).is_file())ifnot(model_path/config.json).is_file():raiseFileNotFoundError(f模型缺少 config.json{model_path})ifnotweight_exists:raiseFileNotFoundError(f模型缺少权重文件{model_path})# # 3. 创建本地 Chroma 客户端# # PersistentClient 不需要启动 Chroma 服务chroma_clientchromadb.PersistentClient(pathstr(chroma_data_path))print(Chroma 状态,chroma_client.heartbeat())# # 4. 创建 Embedding 函数# embedding_functionSentenceTransformerEmbeddingFunction(model_namestr(model_path),devicecpu,# 有合适的 CUDA 环境可改为 cudanormalize_embeddingsTrue)print(Embedding 模型加载完成。)# # 5. 创建或获取集合# collectionchroma_client.get_or_create_collection(namerag_db_gte_large_zh,embedding_functionembedding_function,metadata{hnsw:space:cosine})print(集合名称,collection.name)# # 6. 准备测试文档# documents[(在向量搜索领域我们拥有多种索引方法和向量处理技术它们使我们能够在召回率、响应时间和内存使用之间做出权衡。),(虽然单独使用特定技术如倒排文件IVF、乘积量化PQ或分层导航小世界HNSW通常能够带来满意的结果。),(GraphRAG 本质上就是 RAG只不过与一般 RAG 相比其检索路径上多了一个知识图谱。)]ids[id1,id2,id3]metadatas[{chapter:3,verse:16},{chapter:4,verse:5},{chapter:12,verse:5}]# # 7. 添加或更新文档# # upsert 可以重复运行相同 ID 会更新而不是报错collection.upsert(idsids,documentsdocuments,metadatasmetadatas)print(文档添加完成。)print(文档总数,collection.count())# # 8. 查看集合中的数据# peek_resultcollection.peek(limit1)print(\n集合中的第一条数据)print(ID,peek_result[ids])print(文档,peek_result[documents])print(元数据,peek_result[metadatas])# # 9. 根据 ID 获取文档和向量# id_resultcollection.get(ids[id2],include[documents,embeddings,metadatas])print(\n根据 ID 查询)print(文档,id_result[documents])print(元数据,id_result[metadatas])ifid_result[embeddings]isnotNone:embedding_arraynp.asarray(id_result[embeddings])print(向量形状,embedding_array.shape)# # 10. 相似度检索# query索引技术有哪些query_resultcollection.query(query_texts[query],n_results2,include[documents,metadatas,distances])print(\n相似度检索结果)forindex,(document,metadata,distance)inenumerate(zip(query_result[documents][0],query_result[metadatas][0],query_result[distances][0]),start1):print(f\n第{index}条)print(文档,document)print(元数据,metadata)print(余弦距离,distance)# # 11. 根据元数据过滤后检索# filter_resultcollection.query(query_texts[query],n_results2,include[documents,metadatas,distances],where{verse:5})print(\nverse5 的过滤检索结果)forindex,(document,metadata,distance)inenumerate(zip(filter_result[documents][0],filter_result[metadatas][0],filter_result[distances][0]),start1):print(f\n第{index}条)print(文档,document)print(元数据,metadata)print(余弦距离,distance)