Chroma 向量数据库查询指南（自定义嵌入模型）

本指南使用1024 维的自定义嵌入模型BAAI/bge-large-zh-v1.5所有示例代码均在本地 Chroma 中运行无需 API Key。 推荐版本组合包名推荐版本说明chromadb0.13.2或更新支持SenteceTransformerEmbeddingFunction兼容自定义模型sentence-transformers2.2.2或更新支持BAAI/bge-large-zh-v1.5等模型torch(自动依赖)1.13.0由 sentence-transformers 自动安装用于深度学习计算Python3.8~3.11确保兼容所有库一、Chroma 简介Chroma 是一个开源的向量数据库专为 AI 应用尤其是大模型应用设计。它的核心目标是将知识knowledge、事实facts和技能skills等现实世界的文档高效地整合进大模型应用中。Chroma 提供三类核心工具存储存储文档数据及其元数据嵌入对文档和查询进行向量化搜索基于向量相似度进行搜索Chroma 的设计优先考虑三个方向足够简单并提升开发者效率搜索之上的分析能力追求高性能本文以 Python 为例使用自定义嵌入模型进行讲解。二、环境准备与安装2.1 安装依赖pipinstallchromadb sentence-transformers说明sentence-transformers用于加载我们指定的 1024 维嵌入模型BAAI/bge-large-zh-v1.5。首次运行时会自动下载模型文件约 1.3 GB。2.2 创建自定义嵌入函数importchromadbfromchromadb.utils.embedding_functionsimportSentenceTransformerEmbeddingFunction# 创建使用 1024 维模型的自定义嵌入函数embedding_fnSentenceTransformerEmbeddingFunction(model_nameBAAI/bge-large-zh-v1.5# 输出 1024 维向量)# 创建 Chroma 客户端chroma_clientchromadb.Client()# 创建集合并指定嵌入函数collectionchroma_client.create_collection(nameacademic_papers,embedding_functionembedding_fn)print(f集合创建成功使用的嵌入模型维度为 1024。)三、准备示例数据我们添加一批学术论文让 Chroma 使用自定义模型自动生成 1024 维向量collection.add(ids[paper_001,paper_002,paper_003,paper_004,paper_005],documents[视觉Transformer在目标检测中的应用与改进,基于图神经网络的分子属性预测方法研究,大规模语言模型的指令微调技术综述,计算机视觉中的自监督学习最新进展,Transformer在医学图像分割中的探索],metadatas[{author:Alex Chen,category:AI,year:2023,citations:45,keywords:[transformer,object detection]},{author:Maria Rodriguez,category:Chemistry,year:2021,citations:12,keywords:[gnn,molecule]},{author:Wei Zhang,category:AI,year:2023,citations:89,keywords:[llm,instruction tuning]},{author:Sarah Lee,category:Computer Vision,year:2022,citations:67,keywords:[self-supervised,vision]},{author:James Brown,category:Medical Imaging,year:2020,citations:23,keywords:[transformer,segmentation]}])print(f已添加{collection.count()}篇论文所有向量均为 1024 维。)四、基础查询query()4.1 使用文本查询自动向量化resultscollection.query(query_texts[有哪些关于深度学习模型的最新研究],n_results2)print(查询结果)fori,docinenumerate(results[documents][0]):print(f{i1}.{doc})print(f 元数据:{results[metadatas][0][i]})4.2 使用预计算的 1024 维查询向量如果您想手动传入向量必须确保维度为 1024importrandom# 生成两个 1024 维随机向量仅用于演示query_vectors[[random.random()for_inrange(1024)],[random.random()for_inrange(1024)]]resultscollection.query(query_embeddingsquery_vectors,n_results2)print(手动向量查询完成。)提示实际使用时应使用与集合相同的嵌入模型生成查询向量否则语义将不匹配。4.3 返回结果结构query()返回列式格式{ids:[[paper_001,paper_003]],embeddings:[[[...],[...]]],# 默认不返回documents:[[[doc1],[doc2]]],metadatas:[[[meta1],[meta2]]],distances:[[[0.3],[0.7]]]}五、元数据过滤wherewhere参数用于根据元数据如年份、类别、引用量筛选结果。5.1 基础比较操作符# 只查 2023 年的论文resultscollection.query(query_texts[大模型微调],where{year:2023})# 引用量大于 50resultscollection.query(query_texts[神经网络],where{citations:{$gt:50}})支持$eq,$ne,$gt,$gte,$lt,$lte。5.2 逻辑组合$and/$or# AND2020-2022 年且引用量 30resultscollection.query(query_texts[机器学习],where{$and:[{year:{$gte:2020}},{year:{$lte:2022}},{citations:{$gt:30}}]})# ORAI 或 Computer Visionresultscollection.get(where{$or:[{category:AI},{category:Computer Vision}]})关键where顶层只能有一个键要么是$and/$or要么是字段名。5.3 包含操作符$in/$ninresultscollection.get(where{author:{$in:[Alex Chen,Wei Zhang]}})5.4 数组元数据过滤$contains# 查找关键词包含 transformer 的论文resultscollection.query(query_texts[注意力机制],where{keywords:{$contains:transformer}})六、文档内容过滤where_document# 文档包含 Transformerresultscollection.get(where_document{$contains:Transformer})# 正则匹配 自监督 或 无监督resultscollection.get(where_document{$regex:(自监督|无监督)})# 组合2023 年且包含 指令微调resultscollection.query(query_texts[大模型],where{year:2023},where_document{$contains:指令微调})注意$contains和$regex均大小写敏感。七、获取数据get()# 按 ID 获取resultscollection.get(ids[paper_001,paper_003])# 分页每页 2 条跳过 2 条resultscollection.get(limit2,offset2)# 结合元数据过滤resultscollection.get(where{category:AI})八、高级查询组合过滤与本地混合搜索8.1 多重过滤的语义搜索函数defadvanced_search(collection,query_text,min_yearNone,max_yearNone,categoryNone,min_citationsNone,n_results10):conditions[]ifmin_yearisnotNone:conditions.append({year:{$gte:min_year}})ifmax_yearisnotNone:conditions.append({year:{$lte:max_year}})ifcategoryisnotNone:conditions.append({category:category})ifmin_citationsisnotNone:conditions.append({citations:{$gte:min_citations}})iflen(conditions)0:where_clauseNoneeliflen(conditions)1:where_clauseconditions[0]else:where_clause{$and:conditions}returncollection.query(query_texts[query_text],n_resultsn_results,wherewhere_clause,include[documents,metadatas,distances])# 示例2022-2023 年 AI 方向引用 ≥ 50resultsadvanced_search(collection,深度学习效率,min_year2022,max_year2023,categoryAI,min_citations50,n_results5)8.2 手动分页defpaginated_query(collection,query_text,page0,page_size5,**where_filters):total_needed(page1)*page_size resultscollection.query(query_texts[query_text],n_resultstotal_needed,wherewhere_filtersifwhere_filterselseNone)startpage*page_size endstartpage_sizereturn{ids:[results[ids][0][start:end]],documents:[results[documents][0][start:end]],metadatas:[results[metadatas][0][start:end]],distances:[results[distances][0][start:end]]}page2paginated_query(collection,神经网络,page1,page_size3,year2023)8.3 本地混合搜索语义 关键词defhybrid_search(collection,query_text,keyword,n_results5):# 语义搜索sem_resultscollection.query(query_texts[query_text],n_resultsn_results*2)sem_idssem_results[ids][0]sem_docssem_results[documents][0]sem_metassem_results[metadatas][0]sem_distsem_results[distances][0]sem_dict{id_:{doc:doc,meta:meta,dist:dist,source:semantic}forid_,doc,meta,distinzip(sem_ids,sem_docs,sem_metas,sem_dist)}# 关键词搜索try:kw_resultscollection.get(where_document{$contains:keyword},limitn_results*2)kw_dict{id_:{doc:doc,meta:meta,dist:None,source:keyword}forid_,doc,metainzip(kw_results[ids],kw_results[documents],kw_results[metadatas])}except:kw_dict{}# 合并merged{**sem_dict,**kw_dict}forid_inkw_dict:ifid_insem_dict:merged[id_][source]both# 排序both 语义距离 关键词both[dfordinmerged.values()ifd[source]both]sem_onlysorted([dfordinmerged.values()ifd[source]semantic],keylambdax:x[dist])kw_only[dfordinmerged.values()ifd[source]keyword]final(bothsem_onlykw_only)[:n_results]# 重构结果final_ids,final_docs,final_metas[],[],[]foriteminfinal:forid_,datainmerged.items():ifdataisitem:final_ids.append(id_)final_docs.append(data[doc])final_metas.append(data[meta])breakreturn{ids:[final_ids],documents:[final_docs],metadatas:[final_metas],distances:[[]]}mixedhybrid_search(collection,深度学习前沿,Transformer,n_results3)九、最佳实践与常见错误9.1 性能优化使用n_results控制返回数量。利用ids缩小搜索范围。高频查询可预先计算嵌入向量。9.2 元数据设计将常用过滤字段单独存储如year,category。使用数组元数据配合$contains实现多标签过滤。9.3 调试技巧collection.count()检查数据量。先无过滤查询逐步增加条件。注意字段名和值的大小写、类型整数/字符串。9.4 维度相关的错误错误原因解决InvalidArgumentError: expecting embedding with dimension of 1024, got X手动向量维度不一致使用自动向量化或生成 1024 维向量查询结果为空过滤条件过于严格或无匹配检查元数据值放宽条件十、总结查询方式适用场景关键参数/API向量相似度搜索语义检索query()元数据过滤按年份、类别、引用量筛选where全文搜索精确关键词匹配where_document$contains正则匹配模式匹配where_document$regex按 ID 获取精确获取已知记录get()混合搜索语义 关键词query()get()组合本指南所有示例均使用1024 维自定义嵌入模型代码可直接运行。您可以根据需要更换其他模型只需修改SentenceTransformerEmbeddingFunction中的model_name参数即可。温馨提示如果模型下载速度慢可提前设置 HuggingFace 镜像或手动下载模型文件。