在ChromaDB中使用OpenAI嵌入函数

在ChromaDB中使用OpenAI嵌入函数

本教程将引导您在ChromaDB中使用OpenAI嵌入函数。

❗ 在开始本教程之前，请确保您已安装所需的openai Python软件包。如果您尚未安装，可以使用pip快速安装：

pip install openai
​

安装好openai后，按照以下步骤即可在ChromaDB中使用OpenAI嵌入函数。

1. 使用OpenAI嵌入函数创建集合

首先，我们需要初始化ChromaDB客户端和OpenAI嵌入函数。对于OpenAI嵌入函数，您需要提供您的OpenAI API密钥和模型名称。在本示例中，我们使用的是“text-embedding-3-small”模型。

import chromadb
import chromadb.utils.embedding_functions as embedding_functions

# Initialize ChromaDB client
client = chromadb.Client()

# Initialize OpenAI embedding function
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key="USE YOUR OPEN-AI KEY",
    model_name="text-embedding-3-small",
)

# Create the collection with the OpenAI embedding function
matrix_collection = client.create_collection(
    name="matrix",
    embedding_function=openai_ef,
)
​

2. 添加原始文档

接下来，我们使用add方法将《黑客帝国》中的三句著名台词添加到我们的集合中。每篇文档都关联有一个唯一ID。由于我们已为集合设置了OpenAI嵌入函数，Chroma会自动使用该函数对这些文档进行嵌入处理。

matrix_collection.add(
    documents=[
        "The Matrix is all around us.",
        "What you know you can't explain, but you feel it",
        "There is a difference between knowing the path and walking the path",
    ],
    ids=["quote_1", "quote_2", "quote_3"],
)
​

3. 根据一组查询文本进行查询

我们可以使用query方法，根据一组查询文本对集合进行查询。Chroma会先用集合的嵌入函数（本例中为OpenAI嵌入函数）对每条查询文本进行嵌入，然后利用生成的嵌入向量执行查询。在本示例中，我们以文本“什么是矩阵？”作为查询，并要求返回两条结果。

results = matrix_collection.query(query_texts=["What is the Matrix?"], n_results=2)
print(results)
​

有关完整实现的Python代码，请参阅本教程中提供的完整代码。

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