RAG 落地必备的 1 个开源 AI 原生向量数据库 —Chroma

向量数据库大模型机器学习

一、Chroma 核心概念与优势

  1. 什么是 Chroma?

Chroma 是一款开源的向量数据库,专为高效存储和检索高维向量数据设计。其核心能力在于语义相似性搜索,支持文本、图像等嵌入向量的快速匹配,广泛应用于大模型上下文增强(RAG)、推荐系统、多模态检索等场景。与传统数据库不同,Chroma 基于向量距离(如余弦相似度、欧氏距离)衡量数据关联性,而非关键词匹配。

GitHub地址:

https://github.com/chroma-core/chroma

官方文档:

https://docs.trychroma.com/

  1. 核心优势
  • 轻量易用:以 Python/JS 包形式嵌入代码,无需独立部署,适合快速原型开发。
  • 灵活集成:支持自定义嵌入模型(如 OpenAI、HuggingFace),兼容 LangChain 等框架。
  • 高性能检索:采用 HNSW 算法优化索引,支持百万级向量毫秒级响应。
  • 多模式存储:内存模式用于开发调试,持久化模式支持生产环境数据落地。

picture.image

二、安装和基础配置

1、安装Chroma

支持windows和ubuntu操作系统,Python>=3.9

创建虚拟环境以及安装:

  
#创建虚拟环境  
conda create -n chromadb python==3.10  
#激活  
conda activate chromadb  
#安装chromadb  
pip install chromadb

注意:Chroma 默认是本地嵌入式数据库,并不原生支持远程访问像传统数据库那样(比如 PostgreSQL 那种 client-server 模式)。

当然官方也提供了客户端-服务器端模式(Client-Server Mode)。服务器端的启动方式如下:

  
#服务器端启动,默认端口号8000  
chroma run --path /db_path

2、初始化客户端

内存模式(调试,实验的场景):

  
import chromadb  
client = chromadb.Client()

持久化模式(生产环境):

在创建的时候,可以配置本地的存储路径

  
import chromadb  
# 数据保存至本地目录,path填写绝对路径  
client = chromadb.PersistentClient(path="/path/to/save") 

Client-Server模式的客户端:

前两种,都是本地模式,chroma的服务端和客户端需要位于同一台机器。CS模式可以独立部署,通过httpclient进行访问。

  
import chromadb  
chroma_client = chromadb.HttpClient(host='localhost', port=8000)

picture.image

三、增删改查操作

  1. 创建集合(Collection)

集合是 Chroma 中管理数据的基本单元,类似传统数据库的表。 集合的name名称有以下约束:

  • 名称的长度必须介于 3 到 63 个字符之间。
  • 名称必须以小写字母或数字开头和结尾,中间可以包含点、破折号和下划线。
  • 名称不得包含两个连续的点。
  • 该名称不能是有效的 IP 地址。
Chroma 集合是用一个名称和一个可选的嵌入函数创建的。

如果您提供嵌入函数,则每次获取集合时都必须提供它。 |

  
# 创建  
collection = client.create_collection(name="my_collection", embedding_function=emb_fn)  
  
  
# 获取  
collection = client.get_collection(name="my_collection", embedding_function=emb_fn)  
  
  
# 若没有则创建,若有则获取  
collection = chroma_client.get_or_create_collection(name="my_collection2")

如果不提供嵌入函数,则使用默认的嵌入函数 sentence transformer 使用的的是一个小型的模型all-MiniLM-L6-v2,该模型主要是针对英语场景。一般我们都需要自定义一个嵌入函数:

  
import chromadb  
from sentence_transformers import SentenceTransformer  
  
class SentenceTransformerEmbeddingFunction:  
    def __init__(self, model_path: str, device: str = "cuda"):  
        self.model = SentenceTransformer(model_path, device=device)  
    def __call__(self, input: list[str]) -> list[list[float]]:  
        if isinstance(input, str):  
            input = [input]  
        return self.model.encode(input, convert_to_numpy=True).tolist()  
  
# 创建/加载集合(含自定义嵌入函数)  
embed_model = SentenceTransformerEmbeddingFunction(  
    model_path=r"D:\Test\LLMTrain\testllm\llm\BAAI\bge-m3",  
    device="cuda"  # 无 GPU 改为 "cpu"  
)  
# 创建客户端和集合  
client = chromadb.Client()  
collection = client.create_collection("my_knowledge_base",   
                                      metadata={"hnsw:space": "cosine"},  
                                      embedding_function=embed_model)

创建collect时,可以配置如下参数。

  • name标识collect的名称,是必填项;
  • embedding_function,指定嵌入函数,不填为默认的嵌入模型。
  • metadata,元数据,比如索引方式等,非必填。
  
from datetime import datetime  
collection = client.create_collection(  
    name="my_collection",   
    embedding_function=emb_fn,  
    metadata={  
        "description": "my first Chroma collection",  
        "created": str(datetime.now())  
    }    
)

集合还有一些常用方法:

  • peek() - 返回集合中前 10 个项目的列表。

  • count() -返回集合中的项目数。

  • modify() -重命名集合

  
collection.peek()   
collection.count()   
collection.modify(name="new_name")

2、写入数据

写入数据时,配置以下参数:

  • document,原始的文本块。

  • metadatas,描述文本块的元数据,kv键值对。

  • ids,文本块的唯一标识,每个文档必须具有唯一关联的ID。 若添加两次相同的 ID 将导致仅存储初始值。

  • embeddings,对于已经向量化的文本块,可以直接写入结果。如果不填,则在写入时,使用指定或者默认的嵌入函数对documents进行向量化。

  
collection.add(  
    documents=["lorem ipsum...", "doc2", "doc3", ...],  
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],  
    ids=["id1", "id2", "id3", ...]  
)

或者

  
collection.add(  
    embeddings=[[1.1, 2.3, 3.2], [4.5, 6.9, 4.4], [1.1, 2.3, 3.2], ...],  
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],  
    ids=["id1", "id2", "id3", ...]  
)

3、修改数据

提供ids(文本唯一标识)。

  
collection.update(  
    ids=["doc1"],  # 使用已存在的ID  
    documents=["RAG是一种检索增强生成技术222"]  
)

4、更新插入方法

Chroma 还支持更新插入操作,更新现有项目,如果项目尚不存在则添加它们。

  
collection.upsert(  
    ids=["id1", "id2", "id3", ...],  
    embeddings=[[1.1, 2.3, 3.2], [4.5, 6.9, 4.4], [1.1, 2.3, 3.2], ...],  
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],  
    documents=["doc1", "doc2", "doc3", ...],  
)

5、删除数据

Chroma 支持通过以下方式从集合中删除项目ID使用delete。与每个项目相关的嵌入、文档和元数据将被删除。

还支持where过滤器。如果没有ID提供,它将删除集合中与where筛选。

  
# 提供ids  
collection.delete(ids=["doc1"])  
  
# where 条件删除  
collection.delete(  
    ids=["id1", "id2", "id3",...],  
where={"chapter": "20"}  
)

6、查询数据

(1)查询所有数据

  
all_docs = collection.get()  
print("集合中所有文档:", all_docs)

(2)根据ids查询

可以通过以下方式从集合中检索项目ID使用get。

  
collection.get(  
 ids=["id1", "id2", "id3", ...],  
where={"style": "style1"}  
)

(3)查询嵌入

可以通过多种方式查询 Chroma 集合,使用query方法。比如 使用query_embedding。

  
collection.query(  
    query_embeddings=[[11.1, 12.1, 13.1],[1.1, 2.3, 3.2], ...],  
    n_results=10,  
    where={"metadata_field": "is_equal_to_this"},  
    where_document={"$contains":"search_string"}  
)
  • 查询将返回n_result每个最接近的匹配查询嵌入,按顺序排列。
  • 可选where过滤字典可以通过metadata与每个文档关联。
  • 此外,where document可以提供过滤字典来根据文档内容进行过滤。

(4)查询相似文档

还可以通过一组查询文本query_texts. Chroma 将首先嵌入每个查询文本与集合的嵌入函数,然后使用生成的嵌入执行查询。

  
# 查询相似文档  
results = collection.query(  
    query_texts=["什么是RAG技术?"],  
    n_results=3  
)  
print("查询的结果",results)

查询结果配置

  • 当使用 get 或 query 时,您可以使用include参数来指定您想要返回的数据包括:embeddings, documents,metadatas;include为数组,可以传多个值。
  • 对于查询query,默认返回距离distances结果。
  • embeddings出于性能考虑,默认不返回,直接显示None ,若想返回,则include中包含embeddings即可。
  • ID始终会返回。
  • 返回值里有included参数,表明本次返回的数据有哪些类型。
  • embeddings将以二维 NumPy 数组的形式返回。
  
# Only get documents and ids  
collection.get(  
    include=["documents"]  
)  
collection.query(  
    query_embeddings=[[11.1, 12.1, 13.1],[1.1, 2.3, 3.2], ...],  
    include=["documents"]  
)

查询的结果示例

  
{'ids': [['doc1', 'doc3', 'doc2']],   
'embeddings': None,   
'documents': [['RAG是一种检索增强生成技术', '三英战吕布', '向量数据库存储文档的嵌入表示']],   
'uris': None,  
'included': ['metadatas', 'documents', 'distances'],   
'data': None,   
'metadatas': [[{'source': 'tech_doc'}, {'source': 'tutorial1'}, {'source': 'tutorial'}]],   
'distances': [[0.2373753786087036, 0.7460092902183533, 0.7651787400245667]]  
}

四、实战操作

将一批数据插入向量数据库,再根据一个问题从向量数据库中找出相似数据。

1、安装包

  
pip install sentence_transformers  
pip install modelscope

2、下载Embedding模型到本地

  
#模型下载  
from modelscope import snapshot_download  
model_dir = snapshot_download('BAAI/bge-m3',cache_dir=r"D:\Test\LLMTrain\testllm\llm")

3、核心逻辑:写入数据和查询相似度

  
import chromadb  
from sentence_transformers import SentenceTransformer  
class SentenceTransformerEmbeddingFunction:  
    def __init__(self, model_path: str, device: str = "cuda"):  
        self.model = SentenceTransformer(model_path, device=device)  
    def __call__(self, input: list[str]) -> list[list[float]]:  
        if isinstance(input, str):  
            input = [input]  
        return self.model.encode(input, convert_to_numpy=True).tolist()  
# 创建/加载集合(含自定义嵌入函数)  
embed_model = SentenceTransformerEmbeddingFunction(  
    model_path=r"D:\Test\LLMTrain\testllm\llm\BAAI\bge-m3",  
    device="cpu"  # 无 GPU 改为 "cpu",有则为cuda  
)  
# 创建客户端和集合  
client = chromadb.PersistentClient(path=r"D:\Test\LLMTrain\chromadb_test\chroma_data")  
collection = client.get_or_create_collection("my_knowledge_base",  
                                      metadata={"hnsw:space": "cosine"},  
                                      embedding_function=embed_model)  
# 添加文档  
collection.add(  
    documents=[ "向量数据库存储文档的嵌入表示", "三英战吕布","RAG是一种检索增强生成技术"],  
    metadatas=[{"source": "tech_doc"}, {"source": "tutorial"}, {"source": "tutorial1"}],  
    ids=["doc1", "doc2", "doc3"]  
)  
# 查询相似文档  
results = collection.query(  
    query_texts=["什么是RAG技术?"],  
    n_results=3  
)  
print("查询的结果",results)

执行返回结果:

  
查询的结果   
{  
'ids': [['doc3', 'doc2', 'doc1']],   
'embeddings': None,   
'documents': [['RAG是一种检索增强生成技术', '三英战吕布', '向量数据库存储文档的嵌入表示']],   
'uris': None,   
'included': ['metadatas', 'documents', 'distances'],   
'data': None,   
'metadatas': [[{'source': 'tutorial1'}, {'source': 'tutorial'}, {'source': 'tech_doc'}]],   
'distances': [[0.2373753786087036, 0.7460092902183533, 0.7651787400245667]]  
}  

查看结果,我们重点看 distances,值是从小到大排序的,所以3条数据与问题“什么是RAG技术?”的相似度情况: distances值越小越相似。因此,第1条数据与问题最相似。

picture.image

RAG进阶:Embedding Models嵌入式模型原理和选型指南

LlamaIndex入门指南和RAG原理

探秘 Hugging Face:何为 Hugging Face 及其核心组件全览

LLaMA Factory 如何对大模型进行微调、导出和量化!掌握这几步,轻松搞定!

大模型部署该选谁?Ollama、vLLM 和 LMDeploy,各有千秋!

大型语言模型如何高效微调量化?答案就是 QLoRA!

LLaMA Factory微调后的大模型在vLLM框架中对齐对话模版

AI大模型分布式训练技术原理是什么?看这篇超犀利解析!

XTuner对大模型 微调训练、转换、合并和验证 !这篇文章让你轻松拿捏!

基于Bert大模型增量微调训练:中文评价情感分析

大模型高并发请求和单卡显存不足怎么办?分布式推理和量化部署来救场!

AI 大模型评测难?OpenCompass 这个‘自动跑分神器’来破局!

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 NoSQL 的实践与探索
随着 NoSQL 的蓬勃发展越来越多的数据存储在了 NoSQL 系统中,并且 NoSQL 和 RDBMS 的界限越来越模糊,各种不同的专用 NoSQL 系统不停涌现,各具特色,形态不一。本次主要分享字节跳动内部和火山引擎 NoSQL 的实践,希望能够给大家一定的启发。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论