Chroma 教程

自定义嵌入函数

🎉摘要:详解ChromaDB自定义嵌入函数的最小实现与完整实现,包含register_embedding_function注册、get_config序列化和build_from_config反序列化,附带阿里云API调用示例代码。

如果内置的嵌入函数不能满足你的需要,你可以选择自定义嵌入函数。

这是最灵活也最常用的方式 —— 通常公司里一般有自研的 embedding 服务、或者调用第三方提供的服务,直接包一层就能用,抽取到单独的文件中,其他项目还可以复用。

前面“内置嵌入函数一览”介绍了很多内置函数,自定义最简单的方式就是参考内置实现。

最小实现

一个嵌入函数的最小实现,只需要一个类,实现 __call__,如下:

from chromadb import Documents, EmbeddingFunction, Embeddings


class MyEmbedding(EmbeddingFunction[Documents]):
    def __call__(self, input: Documents) -> Embeddings:
        # input 是 List[str]
        # 返回 List[List[float]]
        return None

下面以使用阿里云远程 API 为例,实现自定义嵌入函数:

import chromadb
import os
import tempfile

import numpy as np
import requests
from chromadb import Documents, EmbeddingFunction, Embeddings
from dotenv import load_dotenv

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

# 从 .env 读取环境变量 OPENAI_API_KEY
load_dotenv()

# 自定义嵌入函数
class MyEmbedding(EmbeddingFunction[Documents]):
    def __call__(self, input: Documents) -> Embeddings:
        # input 是 List[str]
        # 返回 List[List[float]]
        resp = requests.post("https://llm-0sb8vr1k7aav4kn3.cn-beijing.maas.aliyuncs.com/compatible-mode/v1/embeddings",
            json={
                "model": "qwen3.7-text-embedding-flash",
                "input": input,
                "encoding_format": "float",
            }, headers={
                "Authorization": f"Bearer {os.getenv("OPENAI_API_KEY", "")}",
                "Content-Type": "application/json",
            }, timeout=60)
        resp.raise_for_status()
        data = resp.json()

        return [
            np.array(item["embedding"], dtype=np.float32) for item in data["data"]
        ]


if __name__ == "__main__":
    ef = MyEmbedding()
    client = chromadb.PersistentClient(path="./chroma_data")
    col = client.get_or_create_collection(
        name="kb_zh",
        embedding_function=ef,
        metadata={"hnsw:space": "cosine"}
    )

    col.add(
        ids=["a", "b"],
        documents=["如何申请退款", "快递丢了怎么办"]
    )

    res = col.query(
        query_texts=["买的东西不想要了,钱能退回来吗"],
        n_results=1
    )
    print(res)

注意,在项目根目录创建 .env,添加如下内容:

# OpenAI API Key
OPENAI_API_KEY=sk-8e20ca1xxxxxxxxxxxxxxxxx012e8ffc3

运行示例,输出如下:

D:\xxx\chat07_7.7_demo1.py:45: DeprecationWarning: The class MyEmbedding does not implement __init__. This will be required in a future version.
  ef = MyEmbedding()
MyEmbedding 这个类没有实现 __init__ 方法。在未来的版本中,这会变成强制要求(会报错)。 ef = MyEmbedding()
{'ids': [['a']], 'embeddings': None, 'documents': [['如何申请退款']], 'uris': None, 
  'included': ['metadatas', 'documents', 'distances'], 'data': None, 
  'metadatas': [[None]], 'distances': [[0.411679744720459]]}

从运行效果来看,成功实现了嵌入函数,也能正常工作。但如果只做到这一步,有两个问题:

  1. 这个配置不会被持久化,换个进程 get_collection 就找不回来

  2. 服务端不知道你用的什么模型

完整实现(推荐)

为了解决最小实现的问题,下面给出一个完整版本。补上 name()、get_config()、build_from_config(),再加注册装饰器,模板如下:

from typing import Dict, Any
from chromadb import Documents, EmbeddingFunction, Embeddings
from chromadb.utils.embedding_functions import register_embedding_function

# 类名字根据自己业务需要进行修改,这里叫 AliyunEmbeddingFunction,因为内部调用阿里云服务实现
# @register_embedding_function 用来注册进全局表,否则反序列化时找不到类
@register_embedding_function
class AliyunEmbeddingFunction(EmbeddingFunction[Documents]):

    def __init__(self, dim: int = 64):
        #

    def __call__(self, input: Documents) -> Embeddings:
        #

    # 下面三个是"可持久化"的关键
    @staticmethod
    def name() -> str:
        # 注册名,全局唯一,存到集合配置里

    def get_config(self) -> Dict[str, Any]:
        # 要序列化的参数(必须能 JSON 化,所以别塞模型对象)

    @staticmethod
    def build_from_config(config: Dict[str, Any]) -> "AliyunEmbeddingFunction":
        # 从配置反构造实例

注意,使用 @register_embedding_function 注册之后,可以通过“"aliyum_df" in known_embedding_functions”验证一下,注册是否成功。

完整代码如下:

import chromadb
import os
import tempfile
import numpy as np
import requests

from typing import Dict, Any
from chromadb import Documents, EmbeddingFunction, Embeddings
from chromadb.utils.embedding_functions import register_embedding_function, known_embedding_functions
from dotenv import load_dotenv

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

# 从 .env 读取环境变量 OPENAI_API_KEY
load_dotenv()

@register_embedding_function
class AliyunEmbeddingFunction(EmbeddingFunction[Documents]):
    """演示用:一个确定性的哈希嵌入,不依赖任何模型"""

    def __init__(self, dim: int = 64):
        self.dim = dim

    def __call__(self, input: Documents) -> Embeddings:
        print(f"__call__ => {input}")
        resp = requests.post("https://llm-0sb8vr1k7aav4kn3.cn-beijing.maas.aliyuncs.com/compatible-mode/v1/embeddings",
             json={
                 "model": "qwen3.7-text-embedding-flash",
                 "input": input,
                 "encoding_format": "float",
             }, headers={
                 "Authorization": f"Bearer {os.getenv("OPENAI_API_KEY", "")}",
                 "Content-Type": "application/json",
             }, timeout=60)
        resp.raise_for_status()
        data = resp.json()

        return [
            np.array(item["embedding"], dtype=np.float32) for item in data["data"]
        ]

    # 下面三个是"可持久化"的关键
    @staticmethod
    def name() -> str:
        return "aliyum_df"

    def get_config(self) -> Dict[str, Any]:
        return {"dim": self.dim}

    @staticmethod
    def build_from_config(config: Dict[str, Any]) -> "AliyunEmbeddingFunction":
        return AliyunEmbeddingFunction(dim=config.get("dim", 64))


if __name__ == "__main__":
    # 查看自定义嵌入函数是否存在
    print(f"自定义嵌入函数是否存在?{"aliyum_df" in known_embedding_functions}")

    # 使用自定义的嵌入函数
    # 为了演示自定义嵌入函数可以恢复,下面将创建两个客户端
    ef = AliyunEmbeddingFunction(dim=32)

    # 客户端1,用来添加数据到集合
    client = chromadb.PersistentClient(path="./chroma_data")
    col = client.get_or_create_collection(
        name="kb_zh",
        embedding_function=ef,
        metadata={"hnsw:space": "cosine"}
    )
    col.add(
        ids=["a", "b"],
        documents=["如何申请退款", "快递丢了怎么办"]
    )

    # 客户端2,重新打开一个客户端,用来从结合中查询数据
    client2 = chromadb.PersistentClient(path="./chroma_data")
    col2 = client2.get_or_create_collection(name="kb_zh")
    res = col2.query(
        query_texts=["买的东西不想要了,钱能退回来吗"],
        n_results=1
    )
    print(res)

运行代码,输出如下:

自定义嵌入函数是否存在?True
__call__ => ['如何申请退款', '快递丢了怎么办']
__call__ => ['买的东西不想要了,钱能退回来吗']
{'ids': [['a']], 'embeddings': None, 'documents': [['如何申请退款']], 'uris': None, 
  'included': ['metadatas', 'documents', 'distances'], 'data': None, 
  'metadatas': [[None]], 'distances': [[0.404982328414917]]}

从输出可以看出,__call__ 被调用了两次,第一次是添加数据到集合,第二次是从数据查询数据。

通常,将自定义的嵌入函数单独放进一个 python 文件(比如 myapp/embeddings.py),入口文件里 import 一下,其他地方就可以直接使用了。

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