Chroma向量数据库实战指南：从基础使用到生产环境优化

1次阅读

共计 1865 个字符，预计需要花费 5 分钟才能阅读完成。

向量数据库在 AI 时代变得尤为重要，特别是在处理高维数据时。无论是推荐系统、图像搜索还是自然语言处理，都需要快速地对向量进行相似性检索。Chroma 是一个轻量级、高性能的向量数据库，特别适合中小规模的项目。它的主要特点包括：

简单易用的 API
支持多种距离度量（如余弦相似度、欧氏距离）
内存和持久化存储选项
与常见机器学习框架良好集成

首先我们需要安装 Chroma:

pip install chromadb

然后初始化客户端:

import chromadb

# 创建内存中的客户端
client = chromadb.Client()

# 或者指定持久化路径
# client = chromadb.PersistentClient(path="/path/to/store")

假设我们有一组文本数据需要向量化:

from sentence_transformers import SentenceTransformer

# 加载预训练模型
model = SentenceTransformer('all-MiniLM-L6-v2')

texts = ["机器学习很有趣", "深度学习很强大", "向量数据库很有用"]

# 生成向量
embeddings = model.encode(texts)

# 转换为列表格式
embeddings = [embedding.tolist() for embedding in embeddings]

# 创建或获取集合
collection = client.create_collection(name="my_collection")

# 添加数据
collection.add(
    embeddings=embeddings,
    documents=texts,
    ids=["id1", "id2", "id3"]
)

# 查询向量
query_embedding = model.encode("人工智能很酷").tolist()

# 执行搜索
results = collection.query(query_embeddings=[query_embedding],
    n_results=2
)

print(results["documents"])  # 打印相似文本

避免频繁的单条插入操作:

# 不好的做法
for text in texts:
    embedding = model.encode(text).tolist()
    collection.add(embeddings=[embedding], documents=[text])

# 推荐做法
batch_size = 100
for i in range(0, len(texts), batch_size):
    batch = texts[i:i+batch_size]
    embeddings = model.encode(batch)
    collection.add(embeddings=embeddings.tolist(),
        documents=batch,
        ids=[f"id_{j}" for j in range(i, i+len(batch))]
    )

Chroma 支持多种索引类型:

collection = client.create_collection(
    name="optimized_collection",
    metadata={"hnsw:space": "cosine"}  # 使用 HNSW 索引
)

# 持久化客户端
client = chromadb.PersistentClient(path="/data/chroma")

# 定期备份
import shutil
shutil.make_archive("/backups/chroma_backup", 'zip', "/data/chroma")

from threading import Lock

write_lock = Lock()

def safe_add(collection, embedding, text):
    with write_lock:
        collection.add(embeddings=[embedding], documents=[text])