在ChromaDB中使用自定义嵌入函数

在ChromaDB中使用自定义嵌入函数

本教程演示了如何在ChromaDB的集合中使用自定义嵌入函数。

1. 导入所需库

首先，让我们导入必要的库：

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

2. 定义自定义嵌入函数

我们将创建一个自定义嵌入函数，该函数实现了EmbeddingFunction协议：

class SimpleEmbeddingFunction(EmbeddingFunction):
    def __call__(self, texts: Documents) -> Embeddings:
        return [[len(text)] for text in texts]
​

此函数接收一个文本列表，并返回一个嵌入向量列表。为简单起见，我们的示例仅使用每个文本的长度作为其嵌入向量。

⚠️ 警告：这并不是适用于真实场景的有效嵌入函数。它只是一个用于演示目的的简化示例。

3. 使用自定义嵌入函数创建集合

现在，让我们初始化ChromaDB客户端并创建一个集合：

client = chromadb.Client()
matrix_quotes = client.create_collection(
    name="matrix_quotes",
    embedding_function=SimpleEmbeddingFunction()
)
​

4. 向集合中添加文档

我们将一些来自《黑客帝国》的著名台词添加到我们的集合中：

matrix_quotes.add(
    documents=[
        "I know kung fu.",
        "There is no spoon.",
        "The Matrix has you.",
        "Follow the white rabbit."
    ],
    ids=["quote_1", "quote_2", "quote_3", "quote_4"]
)
​

5. 从集合中检索条目

让我们从集合中检索条目及其嵌入向量：

items = matrix_quotes.get(include=['embeddings', 'documents'])
for item in items['ids']:
    index = items['ids'].index(item)
    print(f"ID: {item}")
    print(f"Document: {items['documents'][index]}")
    print(f"Embedding: {items['embeddings'][index]}")
    print("---")
​

6. 对集合进行查询

最后，让我们对集合执行一次查询：

results = matrix_quotes.query(
    query_texts=["What do you know?"],
    n_results=2
)
print("Query Results:")
print(results)
​

当您需要定义自己的文本转数值表示方法时，自定义嵌入函数是非常强大的工具。它们尤其适用于特定领域的应用，或者当您希望使用某种特定的自然语言处理技术时。

有关本教程中使用的完整Python代码，请参阅custom_emb_func.py文件。

免责声明：
本文档采用基于机器的 AI 翻译服务进行翻译。尽管我们力求准确，但请注意，自动翻译可能存在错误或不准确之处。应以原文语言版本的文档作为权威依据。如需获取关键信息，建议使用专业的人工翻译。对于因使用本翻译而产生的任何误解或误读，我们概不负责。