VectorDB with Chroma

Overview

This tutorial demonstrates how to use the PaLM API to create a vector database and retrieve answers to questions from the database. Moreover, you will use ChromaDB, an open-source Python tool that creates embedding databases. ChromaDB allows you to:

Store embeddings as well as their metadata
Embed documents and queries
Search through the database of embeddings

In this tutorial, you'll use embeddings to retrieve an answer from a database of vectors created with ChromaDB.

Setup

First, download and install ChromaDB and the PaLM API Python library.

pip install -q chromadb
pip install -q google-generativeai

import google.generativeai as palm

import pandas as pd

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

Grab an API Key

To get started, you'll need to create an API key.

palm.configure(api_key='PALM_KEY')

models = [m for m in palm.list_models() if 'embedText' in m.supported_generation_methods]
model = models[0]
model

Data

Here is a small set of documents you will use to create an embedding database:

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."

Creating the embedding database with ChromaDB

You will create a custom function for performing embedding using the PaLM API. By inputting a set of documents into this custom function, you will receive vectors, or embeddings of the documents.

def embed_function(texts: Documents) -> Embeddings:
  # Embed the documents using any supported method
  return  [palm.generate_embeddings(model=model, text=text)['embedding']
           for text in texts]

Now you will create the vector database. In the create_chroma_db function, you will instantiate a Chroma client. From there, you will create a collection, which is where you store your embeddings, documents, and any metadata. Note that the embedding function from above is passed as an argument to the create_collection.

Next, you use the add method to add the documents to the collection.

def create_chroma_db(documents, name):
  chroma_client = chromadb.Client()
  db = chroma_client.create_collection(name=name, embedding_function=embed_function)
  for i,d in enumerate(documents):
    db.add(
      documents=d,
      ids=str(i)
    )
  return db

# Set up the DB
db = create_chroma_db([DOCUMENT1, DOCUMENT2, DOCUMENT3], "googlecardb")

WARNING:chromadb:Using embedded DuckDB without persistence: data will be transient

Confirm that the data was inserted by looking at the database:

pd.DataFrame(db.peek(1))

Getting the relevant document

db is a Chroma collection object. You can call query on it to perform a nearest neighbors search to find similar embeddings or documents.

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)
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.'

Now that you have found the relevant passage in your set of documents, you can use it make a prompt to pass into the PaLM 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

The answer function will generate a response based on the query you have passed in. It retrieves the relevant document, and from there calls the PaLM text generation API to generate a response to the query.

text_models = [m for m in palm.list_models() if 'generateText' in m.supported_generation_methods]
text_model = text_models[0]
text_model

def answer(model, query, db, temperature=0.5):
  passage = get_relevant_passage(query, db)
  prompt = make_prompt(query, passage)
  answer = palm.generate_text(prompt=prompt, model=model, candidate_count=3, temperature=temperature, max_output_tokens=1000)
  return answer.candidates[0]['output']

The temperature controls the randomness of the output. The larger the value, the more random the generated text will be.

temperature = 0.65
query = "How do you shift gears in the Google car?"
answer(text_model, query, db, temperature)

'To shift gears in your Google car, simply move the shift lever to the desired position.'

Next steps

You've now created your own embeddings database with ChromaDB and PaLM APIs! Try using your own data to create your own database as well. To learn more about how you can use the embeddings, check out the examples available. To learn how to use other services in the PaLM API, visit the various quickstart guides: