Storing embeddings in a vector database

Store embeddings in a vector database for efficient retrieval.

The process of generating the embeddings doesn't take too long, especially if the machine on which the step is running has a GPU, but it's still not something we want to do every time we need to retrieve a document. Instead, we can store the embeddings in a vector database, which allows us to quickly retrieve the most relevant chunks based on their similarity to the query.

For the purposes of this guide, we'll use PostgreSQL as our vector database. This is a popular choice for storing embeddings, as it provides a scalable and efficient way to store and retrieve high-dimensional vectors. However, you can use any vector database that supports high-dimensional vectors. If you want to explore a list of possible options, this is a good website to compare different options.

For more information on how to set up a PostgreSQL database to follow along with this guide, please see the instructions in the repository which show how to set up a PostgreSQL database using Supabase.

Since PostgreSQL is a well-known and battle-tested database, we can use known and minimal packages to connect and to interact with it. We can use the psycopg2 package to connect and then raw SQL statements to interact with the database.

The code for the step is fairly simple:

from zenml import step

@step

def index_generator(

documents: List[Document],

) -> None:

try:

conn = get_db_conn()

with conn.cursor() as cur:

# Install pgvector if not already installed

cur.execute("CREATE EXTENSION IF NOT EXISTS vector")

conn.commit()

# Create the embeddings table if it doesn't exist

table_create_command = f"""

CREATE TABLE IF NOT EXISTS embeddings (

id SERIAL PRIMARY KEY,

content TEXT,

token_count INTEGER,

embedding VECTOR({EMBEDDING_DIMENSIONALITY}),

filename TEXT,

parent_section TEXT,

url TEXT

);

"""

cur.execute(table_create_command)

conn.commit()

register_vector(conn)