前面介绍嵌入函数都是通过 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 来计算向量,这样就不要下载模型。