ChromaDB调用BGE模型的两种实践方式

前言

在语义搜索、知识库构建等场景中，文本向量化（Embedding）是核心技术环节。作为一款开源的向量数据库，ChromaDB允许开发者通过自定义嵌入函数灵活对接各类模型。本文将详细介绍两种基于BGE模型的实现方案：​​远程API调用​​与​​本地模型部署​​，并解析它们的应用场景与实现细节。

1.chromadb调用BGE模型api

此api接口是Ollama接口方式：

关键点解析：

​​API服务对接​​：通过HTTP POST请求调用部署在9.1.47.89:11434的Ollama服务
​​模型指定​​：使用bge-m3:latest模型的最新版本
​​超时控制​​：设置30秒超时避免长期阻塞
​​异常处理​​：非200状态码时抛出详细错误信息

import requests
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings

class MyEmbeddingFunction(EmbeddingFunction):
    def __call__(self, texts: Documents) -> Embeddings:
        # 调用远程Ollama服务的BGE-M3模型
        response = requests.post(
            "http://9.1.47.89:11434/v1/embeddings",
            json={
                "model": "bge-m3:latest",
                "input": texts
            },
            timeout=30  # 增加超时设置
        )
        
        if response.status_code == 200:
            return [vec['embedding'] for vec in response.json()['data']]
        else:
            raise Exception(f"Embedding API调用失败: {response.text}")

# 初始化自定义嵌入函数
ef = MyEmbeddingFunction()

注意不同版本之间的访问方式可能不一致。“http://9.1.47.89:11434/v1/embeddings”,可修改为"http://9.1.47.89:11434/api/embeddings"。

[vec[‘embedding’] for vec in response.json()[‘data’]]中的"data"可修改为“embeddings”。 都可进行尝试。

2.调用本地模型

from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
from sentence_transformers import SentenceTransformer
model_path = "emmodel/bge-large-zh-v1.5"
model = SentenceTransformer(model_name_or_path=model_path)

class MyEmbeddingFunction(EmbeddingFunction):
    def __call__(self, texts: Documents) -> Embeddings:
        embeddings = [model.encode(x).tolist() for x in texts]
        return embeddings

ef = MyEmbeddingFunction()

关键点解析：
​​本地模型加载​​：使用sentence-transformers库加载预训练模型
​​路径指定​​：从emmodel/目录加载bge-large-zh-v1.5模型文件
​​批量编码​​：对输入文本列表进行并行向量化

模型准备：

# 下载官方模型
git clone https://www.modelscope.cn/company/BAAI/bge-large-zh-v1.5.git

# 或使用huggingface-hub
from huggingface_hub import snapshot_download
snapshot_download(repo_id="BAAI/bge-large-zh-v1.5")