Chroma 教程

一个配置模板

🎉摘要:本文提供ChromaDB集合配置模板,包含make_collection、tune_for_bulk_load和tune_for_serving函数,用于创建集合、批量灌数据时调优sync_threshold,以及灌数据后恢复服务配置,并附有使用OllamaEmbeddingFunction的示例代码。

为了使用方便,下面给出了一个配置模板,包含如下函数:

  • make_collection()  创建一个集合,必须指定 client、集合名称,默认 embedding_function 为 None。经过测试,如果没有传递 embedding_function,会抛出 ValueError: You must provide an embedding function to compute embeddings. 错误(经过测试,将 hnsw 配置注释了,运行就正常)。

  • tune_for_bulk_load()  在批量灌数据前,将 sync_threshold 设置为 10 万。

  • tune_for_serving()  在灌数据后,又将 sync_threshold 恢复到 1000。

代码如下:

def make_collection(
    client,
    name: str,
    embedding_function=None,
    *,
    space: str = "cosine",
    ef_search: int = 100,
    ef_construction: int = 100,
    max_neighbors: int = 16,
    description: str = "",
):
    """带默认调优配置的建集合函数"""
    return client.get_or_create_collection(
        name=name,
        embedding_function=embedding_function,
        configuration={
            "hnsw": {
                "space": space,
                "ef_search": ef_search,
                "ef_construction": ef_construction,
                "max_neighbors": max_neighbors,
                # 灌数据阶段可以临时调大,这里保持默认
                "sync_threshold": 1000,
                "batch_size": 100,
            }
        },
        metadata={"description": description, "space": space, "ef_search": ef_search},
    )


def tune_for_bulk_load(col):
    """批量导入前调用"""
    col.modify(configuration={"hnsw": {"sync_threshold": 100_000, "batch_size": 1000}})


def tune_for_serving(col, ef_search: int = 100):
    """导入完切回服务模式"""
    col.modify(configuration={
        "hnsw": {"sync_threshold": 1000, "batch_size": 100, "ef_search": ef_search}
    })


# 使用模板创建集合和灌数据
if __name__ == "__main__":
    import chromadb
    import os
    import tempfile
    from chromadb.utils.embedding_functions import OllamaEmbeddingFunction

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

    # 创建 Ollama 本地嵌入函数
    ef = OllamaEmbeddingFunction(
        url="http://localhost:11434",
        model_name="bge-m3",  # 先在 ollama 里 pull 下来
        timeout=60,
    )

    # 创建持久化客户端和集合
    client = chromadb.PersistentClient(path="./chroma_data")
    col = make_collection(client, "kb_demo", embedding_function=ef,
                          space="cosine", ef_search=150, description="产品知识库")
    # 灌数据前调用,修改 sync_threshold 阈值
    tune_for_bulk_load(col)

    # ... 灌几十万条 ...
    # 下面仅仅写入几条,你可以灌几十万条数据
    col.add(
        ids=["r1", "r2", "r3", "r4", "r5"],
        documents=[
            "退款流程:在订单页点申请退款,3 个工作日内到账",
            "发票开具:在订单详情页申请电子发票",
            "物流查询:发货后可在物流详情里看到进度",
            "退货政策:7 天无理由退货,运费买家承担",
            "账号安全:如何修改密码和绑定手机",
        ],
        metadatas=[
            {"cat": "售后", "year": 2024, "score": 0.9, "tags": ["退款"]},
            {"cat": "财务", "year": 2023, "score": 0.5, "tags": ["发票"]},
            {"cat": "物流", "year": 2024, "score": 0.7, "tags": ["物流"]},
            {"cat": "售后", "year": 2022, "score": 0.3, "tags": ["退货"]},
            {"cat": "账号", "year": 2024, "score": 0.6, "tags": ["安全"]},
        ],
    )

    # 数据灌完,恢复 sync_threshold 阈值
    tune_for_serving(col, ef_search=150)
    print("完成了")

注意,上面示例使用 OllamaEmbeddingFunction 作为演示的嵌入函数,要求你打开 Ollama 服务,并且下载 bge-m3 模型,如下图:

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