Chroma 教程

直接用嵌入函数,绕过 Chroma

🎉摘要:学习在Chroma中自己计算向量并传入embeddings,从而控制批次、并发与重试,缓存向量避免重复计算,并让客户端无需模型依赖。附Python示例,使用BAAI/bge-small-zh-v1.5模型和本地持久化。

前面介绍嵌入函数都是通过 get_or_create_collection() 函数的 embedding_function 参数将它交给 Chroma,让 Chroma 在 add() 和 query() 时帮我们计算向量。

但是,大多数时候你不需要让 Chroma 帮你算向量。自己算,然后只传 embeddings 就可以了,好处是:

  • 能控制批次、并发、重试

  • 能把向量缓存起来,重复文本不用重算

  • 能在别的地方复用(比如自己算相似度)

  • Chroma 客户端不用装模型依赖

示例代码:

import os

# 使用国内端点地址
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import tempfile
import chromadb
from sentence_transformers import SentenceTransformer

# 持久化示例写到临时目录,工作目录保持干净
os.chdir(tempfile.mkdtemp(prefix="chroma_demo_"))

# 使用的嵌入模型
model = SentenceTransformer("BAAI/bge-small-zh-v1.5")
client = chromadb.PersistentClient(path="./chroma_data")

# embedding_function=None 我自己算向量
col = client.get_or_create_collection(
    "kb_coll", embedding_function=None, configuration={"hnsw": {"space": "cosine"}}
)

def embed_docs(texts):
    return model.encode(texts, normalize_embeddings=True).tolist()

def embed_query(text):
    return model.encode(["为这个句子生成表示以用于检索相关文章:" + text],
                        normalize_embeddings=True)[0].tolist()

def add(col, ids, texts, metas=None):
    col.upsert(ids=ids, embeddings=embed_docs(texts),
               documents=texts, metadatas=metas)

def search(col, query, top_k=5, where=None):
    return col.query(query_embeddings=[embed_query(query)],
                     n_results=top_k, where=where)

# 添加数据
add(col, ["d1", "d2"], ["如何申请退款", "发票开具流程"])
# 查询数据
print(search(col, "买的东西不想要了"))

运行示例,输出如下:

Loading weights: 100%|██████████| 71/71 [00:00<00:00, 2176.98it/s]
{'ids': [['d1', 'd2']], 'embeddings': None, 'documents': [['如何申请退款', '发票开具流程']], 
  'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 
  'metadatas': [[None, None]], 'distances': [[0.6550070643424988, 0.7700721621513367]]}

注意,你还可以使用云厂商提供的已远程 HTTP 接口或者本地 Ollama 来计算向量,这样就不要下载模型。

说说我的看法
全部评论()
没有评论
关于
本网站专注于 Java、数据库(MySQL、Oracle)、Linux、软件架构及大数据等多领域技术知识分享。涵盖丰富的原创与精选技术文章,助力技术传播与交流。无论是技术新手渴望入门,还是资深开发者寻求进阶,这里都能为您提供深度见解与实用经验,让复杂编码变得轻松易懂,携手共赴技术提升新高度。如有侵权,请来信告知:hxstrive@outlook.com
其他应用
公众号