使用开源向量数据库ChromaDB构建语义搜索应用程序

使用ChromaDB构建语义搜索应用程序

介绍

随着人工智能应用和用例的兴起,各种工具和技术的流动增加,以促进此类人工智能应用并允许人工智能开发人员构建现实世界的应用程序。在这些工具中,今天我们将学习ChromaDB的工作原理和功能,ChromaDB是一个开源的向量数据库,用于存储来自AI模型(如GPT3.5、GPT-4或任何其他OS模型)的嵌入。嵌入是任何AI应用程序流水线的关键组件。由于计算机只处理向量,因此所有数据必须以嵌入的形式进行矢量化,以在语义搜索应用程序中使用。

让我们深入了解ChromaDB的工作原理,并通过实际代码示例进行实际操作!

本文作为数据科学博客马拉松的一部分发表。

ChromaDB基础知识和安装库

ChromaDB是一个开源的向量数据库,旨在存储向量嵌入以开发和构建大型语言模型应用程序。该数据库使得存储知识、技能和事实以供LLM应用程序使用更加简单。

ChromaDB是一个开源的向量数据库,旨在存储向量嵌入以开发和构建大型语言模型应用程序。该数据库使得存储知识、技能和事实以供LLM应用程序使用更加简单。

上图显示了将ChromaDB与任何LLM应用程序集成时的工作原理。ChromaDB提供了以下功能:

  • 存储带有ID的嵌入和元数据
  • 嵌入文档和查询
  • 搜索嵌入

ChromaDB非常简单易用,并且可以与任何LLM驱动的应用程序一起设置和使用。它旨在提高开发人员的生产效率,是一种开发人员友好的工具。

现在,让我们在Python和Javascript环境中安装ChromaDB。它也可以在Jupyter Notebook中运行,使数据科学家和机器学习工程师能够对LLM模型进行实验。

Python安装

# 在Python环境中安装chromadb
pip install chromadb

Javascript安装

# 在JS环境中安装chromadb
npm install --save chromadb # yarn add chromadb

安装完库之后,我们将在下一节中了解它的各种功能。

ChromaDB的功能和工作原理

我们可以使用像Google Colab这样的Jupyter Notebook环境进行演示。您可以在Google Colab、Kaggle或本地笔记本环境中进行以下实际操作。

创建ChromaDB集合

# 导入chromadb并创建客户端
import chromadb

client = chromadb.Client()
collection = client.create_collection("my-collection")

在上面的代码中,我们实例化了客户端对象,以在存储库文件夹中创建“my-collection”集合。

集合是存储嵌入、文档和任何其他附加元数据以供以后查询的地方。

将文档添加到集合中

# 在数据库中添加文档
collection.add(
    documents=["这是关于猫的文档", "这是关于汽车的文档",
     "这是关于自行车的文档"],
    metadatas=[{"category": "动物"}, {"category": "车辆"}, 
    {"category": "车辆"}],
    ids=["id1", "id2","id3"]
)

现在,我们已经添加了一些示例文档,并附带了元数据和ID,以便以结构化方式存储它们。

ChromaDB会自动存储文本文档,并处理标记化、矢量化和索引,而无需任何额外的命令。

查询集合数据库

# 提出查询以从数据库中检索数据
results = collection.query(
    query_texts=["车辆"],
    n_results=1
)

------------------------------[Results]-------------------------------------
{'ids': [['id2']],
 'embeddings': None,
 'documents': [['这是关于汽车的文档']],
 'metadatas': [[{'category': '车辆'}]],
 'distances': [[0.8069301247596741]]}

通过在集合数据库上简单调用“query()”函数,它将返回基于输入查询的最相似文本及其元数据和ID。在我们的示例中,查询返回包含“vehicle”元数据的相似文本。

具有示例文档的语义搜索应用程序

语义搜索是技术行业中最受欢迎的应用之一,被Google、百度等用于网络搜索。语言模型现在允许使用大量数据的嵌入来开发这样的应用程序,无论是个体层面还是企业组织层面。

我们将使用“pets”文件夹和一些示例文档来处理ChromaDB中的语义搜索应用程序。我们在本地文件夹中有以下文件:

让我们从本地文件夹导入文件并将它们存储在“file_data”中。

# 从pets文件夹导入文件并存储在VectorDB中
import os

def read_files_from_folder(folder_path):
    file_data = []

    for file_name in os.listdir(folder_path):
        if file_name.endswith(".txt"):
            with open(os.path.join(folder_path, file_name), 'r') as file:
                content = file.read()
                file_data.append({"file_name": file_name, "content": content})

    return file_data

folder_path = "/content/pets"
file_data = read_files_from_folder(folder_path)

上述代码从“pets”文件夹中获取文件,并将它们作为文件列表附加到“file_data”中。我们将使用这些文件将其存储在ChromaDB中作为查询目的的嵌入。

# 从file_data获取数据并创建ChromaDB集合
documents = []
metadatas = []
ids = []

for index, data in enumerate(file_data):
    documents.append(data['content'])
    metadatas.append({'source': data['file_name']})
    ids.append(str(index + 1))

# 创建一个宠物文件的集合
pet_collection = client.create_collection("pet_collection")

# 将文件添加到ChromaDB集合中
pet_collection.add(
    documents=documents,
    metadatas=metadatas,
    ids=ids
)

上述代码将文件和元数据从文件列表中提取出来,并将它们添加到名为“pet_collection”的ChromaDB集合中。

在这里需要注意的是,默认情况下ChromaDB使用来自句子转换器的“all-MiniLM-L6-v2”嵌入模型将文本文档转换为向量。现在,让我们查询集合以查看结果。

# 查询数据库以从向量化数据中获取答案
results = pet_collection.query(
    query_texts=["宠物动物的营养需求是什么?"],
    n_results=1
)

results

当我们查询集合时,它会自动从嵌入的文档中找到与我们的查询最相似的文档,然后生成输出结果。我们还可以在输出中看到距离度量,显示特定文档与我们的查询有多接近。

使用不同的嵌入模型

到目前为止,我们已经使用了默认的嵌入模型来对输入文本进行向量化,但是ChromaDB还允许使用句子转换器库中的其他各种模型。我们将使用“paraphrase-MiniLM-L3-v2”模型来对相同的宠物文档进行嵌入,用于我们的语义搜索应用程序。

(注意:如果尚未安装sentence_transformers库,请在执行下面的代码之前安装)

# 导入句子转换器
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('paraphrase-MiniLM-L3-v2')

documents = []
embeddings = []
metadatas = []
ids = []

# 遍历file_data以收集每个文档和元数据
for index, data in enumerate(file_data):
    documents.append(data['content'])
    embedding = model.encode(data['content']).tolist()
    embeddings.append(embedding)
    metadatas.append({'source': data['file_name']})
    ids.append(str(index + 1))

# 创建新的ChromaDB,并使用嵌入添加和查询数据
pet_collection_emb = client.create_collection("pet_collection_emb")

# 将宠物文件添加到pet_collection_emb数据库中
pet_collection_emb.add(
    documents=documents,
    embeddings=embeddings,
    metadatas=metadatas,
    ids=ids
)

上面的代码使用“paraphrase-MiniLM-L3-v2”模型对输入文件进行编码,并添加到新的集合中。

现在,我们可以再次查询数据库以获取最相似的结果。

# 编写文本查询并提交到集合
query = "人们通常拥有哪些不同种类的宠物?"
input_em = model.encode(query).tolist()

results = pet_collection_emb.query(
    query_embeddings=[input_em],
    n_results=1
)
results

ChromaDB支持的嵌入

嵌入是存储各种数据用于AI应用的本地方式。它们可以根据应用的要求表示文本、图像、音频和视频数据。

ChromaDB支持来自不同嵌入提供者的许多AI模型,例如OpenAI、Sentence transformers、Cohere和Google PaLM API。让我们在这里看一些。

Sentence Transformer嵌入

# 从sentence transformer库加载任何模型
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="all-MiniLM-L6-v2")

使用上面的代码,我们可以使用可用模型中的任何模型。您可以在这里找到模型列表。

OpenAI模型

ChromaDB提供了一个包装函数,用于使用OpenAI的任何嵌入模型API进行AI应用

# 调用OpenAI嵌入的函数
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="YOUR_API_KEY",
                model_name="text-embedding-ada-002"
            )

有关ChromaDB函数的更详细信息,请访问官方文档。

Github代码库:点击这里

结论

总之,向量数据库是生成式AI应用的关键构建块。ChromaDB是这样一个向量数据库,它在各种基于LLM的应用中越来越常用。在本博客中,我们通过代码示例了解了ChromaDb的各种功能和工作原理。

要点

  1. 我们通过代码示例学习了ChromaDB的各种功能。
  2. 我们了解了ChromaDB在语义搜索应用中的用例。
  3. 最后,我们看到了ChromaDB支持的嵌入类型,如OpenAI、Cohere和sentence transformers。

常见问题

本文中显示的媒体不归Analytics Vidhya所有,仅由作者自行决定使用。