乌漆嘛黑
0%

大模型应用系列(十五) 大模型RAG项目实战：Chroma 向量数据库介绍与使用

发表于 更新于 分类于 ，

介绍Chroma的原理以及使用方式，说明使用过程中需要注意的问题。

一. Chroma 核心概念与优势

参考文档: Introduction - Chroma Docs

中文文档: Chroma向量数据库完全手册. 这里算是做一个汇总，以及对它的细节做补充。 | by Lemooljiang | Medium

1.1 什么是Chroma

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

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

二. 安装与基础配置

2.1 安装

通过pyton包管理器安装ChromaDB:

1
pip install chromadb # 完整功能
2.2 初始化客户端

  • 内存模型(开发环境):

    1
    2
    import chromadb
    client = chromadb.Client()

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

    1
    2
    import chromadb
    client = chromadb.PersistentClient(path="") # 数据保存至本地目录

三. 常见操作

3.1 创建集合(Collection)

集合是Chroma中管理数据的基本单元，类似传统数据库的表: 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 新版Chromadb不支持直接加载Embedding模型，需要设计一个类先加载模型。但如果是集成到llamaindex就不用(llamaindex转好了)
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()
collection = client.create_collection(
    name="my_collection",
    metadata={"hnsw:space": "cosine"}, # 指定余弦相似度计算
    embedding_function=embed_model) # 自定义嵌入模型
    )
3.2 添加数据

支持自动生成(调用Embedding模型)或手动指定嵌入向量:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 方式1：自动生成向量（使用集合的嵌入模型）
collection.add(
    documents=["文档内容1", "文档内容2"],
    metadatas=[{"来源": "新闻"}, {"来源": "论文"}],
    ids=["id1", "id2"]
) 
# 方式2：手动传入预计算向量
collection.add(
    embeddings=[[0.1, 0.2, ...], [0.3, 0.4, ...]],
    documents=["文本1", "文本2"],
    ids=["id3", "id4"]
)


3.3 查询数据

支持自动生成(调用Embedding模型)或手动指定嵌入向量:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
### 查询和输入最接近的若干个向量
# 文本查询（自动向量化）
results = collection.query(
    query_texts=["查询文本"],
    n_results=5,
    where={"来源": "新闻"}, # 按元数据过滤
    where_document={"$contains": "关键词"} # 按文档内容过滤
)
#向量查询（自定义输入）
results = collection.query(
    query_embeddings=[[0.5, 0.6, ...]],
    n_results=3
)
### 查询向量库中已有数据
# 查看所有
all_docs = collection.get()
print("集合中所有文档：", all_docs["documents"])
# 查看指定索引
updated_docs = collection.get(ids=["doc1"])
print("查看指定索引对应的文档：", updated_docs["documents"])
3.4 数据管理

更新

1
collection.update(ids=["id1"], documents=["新内容"])

删除

1
collection.delete(ids=["id2"])

统计

1
collection.count() # 获取条目数
3.5 示例代码：增删改查
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import chromadb
from sentence_transformers import SentenceTransformer

# 新版Chromadb不支持直接加载Embedding模型，需要设计一个类先加载模型。但如果是集成到llamaindex就不用(llamaindex转好了)
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="/root/autodl-tmp/text2vec-base-chinese-sentence",
    device="cuda"  # 无 GPU 改为 "cpu"
)

# 创建客户端和集合
client = chromadb.Client()
collection = client.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)

# collection.update(
#     ids=["doc1"],  # 使用已存在的ID
#     documents=["更新后的RAG技术内容"]
# )

# # 查看更新后的内容 - 方法1：使用get()获取特定ID的内容
# updated_docs = collection.get(ids=["doc1"])
# print("更新后的文档内容：", updated_docs["documents"])

# # 查看更新后的内容 - 方法2：查询所有文档
# all_docs = collection.get()
# print("集合中所有文档：", all_docs["documents"])

# #删除内容
# collection.delete(ids=["doc1"])

# # 查看更新后的内容 - 方法2：查询所有文档
# all_docs = collection.get()
# print("集合中所有文档：", all_docs["documents"])

# #统计条目
# print(collection.count())

输出如下:

image-20250612000620955

如果我们选择不同的Embedding模型（结果如下），其实对结果影响不大，虽然距离绝对值变了，但他们的相对大小并没有变化。选型时一般根据语言类别选型就行。

使用不同嵌入模型，匹配结果是一样的。

image-20250612000939457

四. 补充

4.1 Chroma的输出

输出的不是余弦相似度(越大越相似)，而是距离(越小越相似)。distance是指两个向量之间的夹角 $\theta$ , 夹角越小，说明两个向量越接。而余弦相似度是值两个向量之间夹角的余弦值 $cos\theta$ , 而 $cos$ 函数是一个单调递减的函数，夹角越小, $cos$ 越大。

4.2 Chroma的优化

Chromadb 在添加数据时，就会把文本转化为向量，同时越接近的数据越放在一起(相当于做了聚类)，这样就能加快检索速度。

4.3 RAG中Embedding模型的使用

真正的RAG中，问题可能是很复杂的，Embedding可能无法完全理解问题的语义，所以需要先大模型理解问题，生成需求，再去检索。

如果您读文章后有收获，可以打赏我喝咖啡哦～