Quickstart Guide to Vector Databases using Chroma

Vector Databases (or Vector Stores) are one of the hot topics in the AI space right now. A lot of traditional database providers like MongoDB, Cassandra, etc., have started supporting vector search, and a whole lot of new companies are providing commercial vector databases like PineCone. On the other hand, there’s a sharp rise in the number of open-source vector stores coming to the limelight like Weaviate and ChromaDB. The reason behind the explosion of this technology is due to its usage in Generative AI applications. In this article, we will discuss how to get started using ChromaDB and see the basic working along with code that can be run on Google Colab.

How do Vector Databases work?

As the name suggests, a Vector database is a type of database that stores information in the form of Vectors (arrays of real numbers). Almost any type of data (text, images, audio) can be encoded in the form of vectors. However, Generative AI and Large Language Models predominantly work with text data which can be converted to embeddings/vectors. All these data points are represented in a high-dimensional space. Unlike traditional databases, where the query results in data that satisfies specified criteria, here the query output is a set of data points that are closest to the input query. They can be used in applications involving similarity search, nearest neighbors in a high dimensional space, recommendation systems, etc.

Image by KDNuggets

Chroma is an open-source vector store with support for popular languages like Python and JavaScript. Another advantage is that it can run in memory. So, it can be run directly on Colab with minimal setup. In just 4 steps, we can get started with a vector database in action. We will explore Chroma using Python Client.

import chromadb

# Create a Client Connection
# To load/persist db use db location as argument in Client method
client = chromadb.Client()

# Create/Fetch a collection 
collection = client.get_or_create_collection(name="students")

Adding data to the database. It can handle the input of documents or embeddings. If just documents are passed to add method without corresponding embeddings, it uses the default embedding model sentence transformers to convert docs to vectors.

collection.add(
  documents = ["i am a graduate student", 
               "i work as software engineer", 
               "he is the richest guy in the class"],
  ids = ["1", "2", "3"]
)

Now, let’s query a simple text against this database, to get most similar documents.

results = collection.query(query_texts=['that person is a millionaire', 
                                        'he is the class topper'], 
                           n_results=2)

Image generated by Author

As you can see, the above code gave the top 2 closest documents to each of the query strings. The distance indicates how close/similar the two documents are. By altering the n_results parameter, we can fetch most similar n documents that are similar to our input.

Check out the entire Colab Notebook by Chroma where larger data is stored & queried here. This is a quick introduction to how Vector databases work. Have a look at the Chroma Usage Guide for more features and methods.