使用 ChromaDB 记录问答#

概述#

本教程演示如何使用 Gemini API 创建矢量数据库并从数据库中检索问题的答案。此外，您将使用 ChromaDB，这是一个用于创建嵌入数据库的开源 Python 工具。 ChromaDB 允许您：

  • 存储嵌入及其元数据

  • 嵌入文档和查询

  • 搜索嵌入数据库

在本教程中，您将使用嵌入从使用 ChromaDB 创建的向量数据库中检索答案。

前提条件#

您可以在 Google Colab 中运行此快速入门。

要在您自己的开发环境中完成本快速入门，请确保您的环境满足以下要求：

  • Python 3.9+

  • 安装 jupyter 以运行笔记本

安装#

首先，下载并安装 Gemini API Python 库。

#!pip install -q google.generativeai

E1220 16:11:32.073512301  374504 backup_poller.cc:127]                 Run client channel backup poller: UNKNOWN:pollset_work {created_time:"2023-12-20T16:11:32.073469114+08:00", children:[UNKNOWN:Bad file descriptor {created_time:"2023-12-20T16:11:32.073447673+08:00", errno:9, os_error:"Bad file descriptor", syscall:"epoll_wait"}]}

#!pip install -q chromadb

然后导入您将在本教程中使用的模块。

import textwrap
import chromadb
import numpy as np
import pandas as pd

import google.generativeai as genai
import google.ai.generativelanguage as glm

# Used to securely store your API key
# from google.colab import userdata

from IPython.display import Markdown
from chromadb import Documents, EmbeddingFunction, Embeddings

获取 API 密钥#

在使用 Gemini API 之前，您必须先获取 API 密钥。如果您还没有密钥，请在 Google AI Studio 中一键创建密钥。

在 Colab 中，将密钥添加到左侧面板“🔑”下的秘密管理器中。将其命名为 API_KEY。 获得 API 密钥后，将其传递给 SDK。您可以通过两种方式执行此操作：

  • 将密钥放入 GOOGLE_API_KEY 环境变量中（SDK 将自动从那里获取它）。

  • 将密钥传递给 genai.configure(api_key=…)

# Or use `os.getenv('API_KEY')` to fetch an environment variable.
# API_KEY=userdata.get('API_KEY')
GOOGLE_API_KEY = "YOUR-API-KEY"
genai.configure(api_key=GOOGLE_API_KEY)

Tip

要点：接下来，您将选择一个模型。任何嵌入模型都适用于本教程，但对于实际应用程序，选择特定模型并坚持使用非常重要。不同型号的输出互不兼容。

Warning

注意：目前，Gemini API 仅在某些区域可用。

for m in genai.list_models():
  if 'embedContent' in m.supported_generation_methods:
    print(m.name)

models/embedding-001

数据#

以下是您将用于创建嵌入数据库的一小部分文档：

DOCUMENT1 = "Operating the Climate Control System  Your Googlecar has a climate control system that allows you to adjust the temperature and airflow in the car. To operate the climate control system, use the buttons and knobs located on the center console.  Temperature: The temperature knob controls the temperature inside the car. Turn the knob clockwise to increase the temperature or counterclockwise to decrease the temperature. Airflow: The airflow knob controls the amount of airflow inside the car. Turn the knob clockwise to increase the airflow or counterclockwise to decrease the airflow. Fan speed: The fan speed knob controls the speed of the fan. Turn the knob clockwise to increase the fan speed or counterclockwise to decrease the fan speed. Mode: The mode button allows you to select the desired mode. The available modes are: Auto: The car will automatically adjust the temperature and airflow to maintain a comfortable level. Cool: The car will blow cool air into the car. Heat: The car will blow warm air into the car. Defrost: The car will blow warm air onto the windshield to defrost it."
DOCUMENT2 = "Your Googlecar has a large touchscreen display that provides access to a variety of features, including navigation, entertainment, and climate control. To use the touchscreen display, simply touch the desired icon.  For example, you can touch the \"Navigation\" icon to get directions to your destination or touch the \"Music\" icon to play your favorite songs."
DOCUMENT3 = "Shifting Gears Your Googlecar has an automatic transmission. To shift gears, simply move the shift lever to the desired position.  Park: This position is used when you are parked. The wheels are locked and the car cannot move. Reverse: This position is used to back up. Neutral: This position is used when you are stopped at a light or in traffic. The car is not in gear and will not move unless you press the gas pedal. Drive: This position is used to drive forward. Low: This position is used for driving in snow or other slippery conditions."

documents = [DOCUMENT1, DOCUMENT2, DOCUMENT3]

使用 ChromaDB 创建嵌入数据库#

您将创建一个自定义函数来使用 Gemini API 执行嵌入。通过将一组文档输入到此自定义函数中，您将收到向量或文档的嵌入。

使用模型 embedding-001 对嵌入进行 API 更改#

对于新的嵌入模型 embedding-001，有一个新的任务类型参数和可选标题（仅在 task_type=RETRIEVAL_DOCUMENT 时有效）。
这些新参数仅适用于最新的嵌入模型。任务类型为：

任务类型

描述

RETRIEVAL_QUERY

指定给定文本是搜索/检索设置中的查询。

RETRIEVAL_DOCUMENT

指定给定文本是搜索/检索设置中的文档。

SEMANTIC_SIMILARITY

指定给定文本将用于语义文本相似性 (STS)。

CLASSIFICATION

指定嵌入将用于分类。

CLUSTERING

指定嵌入将用于聚类。

Tip

注意：为RETRIEVAL_DOCUMENT指定标题可以为检索提供更好质量的嵌入，但此处省略，因为 ChromaDB 的EmbeddingFunction仅支持内容字符串的输入。

class GeminiEmbeddingFunction(EmbeddingFunction):
  def __call__(self, input: Documents) -> Embeddings:
    model = 'models/embedding-001'
    title = "Custom query"
    return genai.embed_content(model=model,
                                content=input,
                                task_type="retrieval_document",
                                title=title)["embedding"]

现在您将创建矢量数据库。在create_chroma_db函数中，您将实例化一个Chroma 客户端。从那里，您将创建一个集合，在其中存储嵌入、文档和任何元数据。请注意，上面的嵌入函数作为参数传递给 create_collection。

接下来，使用 add 方法将文档添加到集合中。

def create_chroma_db(documents, name):
  chroma_client = chromadb.Client()
  db = chroma_client.create_collection(name=name, embedding_function=GeminiEmbeddingFunction())

  for i, d in enumerate(documents):
    db.add(
      documents=[d],
      ids=[str(i)]
    )
  return db

### 
# Set up the DB
db = create_chroma_db(documents, "googlecarsdatabase")

通过查看数据库确认数据已插入：

pd.DataFrame(db.peek(3))
ids embeddings metadatas documents uris data
0 0 [-0.020994942635297775, -0.03876612335443497, ... None Operating the Climate Control System Your Goo... None None
1 1 [0.017410801723599434, -0.04757162556052208, -... None Your Googlecar has a large touchscreen display... None None
2 2 [-0.03194405511021614, -0.023281503468751907, ... None Shifting Gears Your Googlecar has an automatic... None None

获取相关文件#

db是 Chroma 集合对象。您可以对其调用query来执行最近邻搜索以查找相似的嵌入或文档。

def get_relevant_passage(query, db):
  passage = db.query(query_texts=[query], n_results=1)['documents'][0][0]
  return passage

# Perform embedding search
passage = get_relevant_passage("touch screen features", db)
Markdown(passage)

Your Googlecar has a large touchscreen display that provides access to a variety of features, including navigation, entertainment, and climate control. To use the touchscreen display, simply touch the desired icon. For example, you can touch the “Navigation” icon to get directions to your destination or touch the “Music” icon to play your favorite songs.

现在您已经在文档集中找到了相关段落，您可以使用它进行提示以传递到 Gemini API。

def make_prompt(query, relevant_passage):
  escaped = relevant_passage.replace("'", "").replace('"', "").replace("\n", " ")
  prompt = ("""You are a helpful and informative bot that answers questions using text from the reference passage included below. \
  Be sure to respond in a complete sentence, being comprehensive, including all relevant background information. \
  However, you are talking to a non-technical audience, so be sure to break down complicated concepts and \
  strike a friendly and converstional tone. \
  If the passage is irrelevant to the answer, you may ignore it.
  QUESTION: '{query}'
  PASSAGE: '{relevant_passage}'

    ANSWER:
  """).format(query=query, relevant_passage=escaped)

  return prompt

将查询传递给提示：

query = "How do you use the touchscreen in the Google car?"
prompt = make_prompt(query, passage)
Markdown(prompt)

You are a helpful and informative bot that answers questions using text from the reference passage included below. Be sure to respond in a complete sentence, being comprehensive, including all relevant background information. However, you are talking to a non-technical audience, so be sure to break down complicated concepts and strike a friendly and converstional tone. If the passage is irrelevant to the answer, you may ignore it. QUESTION: ‘How do you use the touchscreen in the Google car?’ PASSAGE: ‘Your Googlecar has a large touchscreen display that provides access to a variety of features, including navigation, entertainment, and climate control. To use the touchscreen display, simply touch the desired icon. For example, you can touch the Navigation icon to get directions to your destination or touch the Music icon to play your favorite songs.’

ANSWER:

现在使用generate_content方法从模型生成响应。

model = genai.GenerativeModel('gemini-pro')
answer = model.generate_content(prompt)
Markdown(answer.text)

In your Google car, you can interact with the large touchscreen display by simply touching the desired icon. For instance, to navigate to a specific destination, tap the Navigation icon, or to listen to your favorite tunes, tap the Music icon.

下一步#

要了解有关如何使用嵌入的更多信息，请查看可用的示例。要了解如何使用 Gemini API 中的其他服务，请访问Python 快速入门。