Use Cases

• 1: Hybrid Search
• 2: Adaptive retrieval and shortening embeddings
• 3: Image Search
• 4: Multi Tenancy
• 5: Sparse Vector Search
• 6: Generative Question-Answering

1 - Hybrid Search

We will show you how to use pgvecto.rs to implement hybrid search in this section.

Vector Search and Full Text Search

Postgres provides full text search, based on the match operator @@. It returns true if a tsvector (document) matches a tsquery (query). It doesn’t matter which data type is written first:

SELECT 'a fat cat sat on a mat and ate a fat rat'::tsvector @@ 'cat & rat'::tsquery;
 ?column?
----------
 t

SELECT 'cat & rat'::tsquery @@ 'a fat cat sat on a mat and ate a fat rat'::tsvector;
 ?column?
----------
 t

You could use full text search and vector search together in a single query, with the help of pgvecto.rs. In this example, we create a table named items with two columns: content and embedding. The column content is of type text, and the column embedding is of type vector(3). We insert four items into the table:

CREATE TABLE items (
  id bigserial PRIMARY KEY,
  content text NOT NULL,
  embedding vector(3) NOT NULL -- 3 dimensions
);

INSERT INTO items (content, embedding) VALUES
  ('a fat cat sat on a mat and ate a fat rat', '[1, 2, 3]'),
  ('a fat dog sat on a mat and ate a fat rat', '[4, 5, 6]'),
  ('a thin cat sat on a mat and ate a thin rat', '[7, 8, 9]'),
  ('a thin dog sat on a mat and ate a thin rat', '[10, 11, 12]');

You should get the following table:

postgres=# select * from items;
 id |                  content                   |  embedding   
----+--------------------------------------------+--------------
  2 | a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]
  3 | a fat dog sat on a mat and ate a fat rat   | [4, 5, 6]
  4 | a thin cat sat on a mat and ate a thin rat | [7, 8, 9]
  5 | a thin dog sat on a mat and ate a thin rat | [10, 11, 12]
(4 rows)

Let’s search for the items with the word cat and rat in their content first. We will use full text search to do that. The function to_tsvector converts a text to a tsvector. The operator @@ checks if a tsvector matches a tsquery.

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery;

SELECT to_tsvector('english', 'a fat  cat sat on a mat - it ate a fat rats');
                  to_tsvector
-----------------------------------------------------
 'ate':9 'cat':3 'fat':2,11 'mat':7 'rat':12 'sat':4

We should get the following result:

                  content                   | embedding 
--------------------------------------------+-----------
 a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]
 a thin cat sat on a mat and ate a thin rat | [7, 8, 9]
(2 rows)

Then we will sort the result by the distance between the embedding of each item and the embedding of the query. We will use vector search to do that:

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
ORDER BY embedding <-> '[7, 8, 9]'::vector;

We will get the item with the embedding [7, 8, 9] first, and the item with the embedding [1, 2, 3] second:

                  content                   | embedding 
--------------------------------------------+-----------
 a thin cat sat on a mat and ate a thin rat | [7, 8, 9]
 a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]
(2 rows)

Advanced search: Merge the results of full text search and vector search

You could normalize the rank of the full text search result and the vector search result, and then merge them together:

-- The query starts with a Common Table Expression (CTE), named "semantic_search".
WITH semantic_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY embedding <-> '[7, 8, 9]') AS rank
    FROM items
    ORDER BY embedding <-> '[7, 8, 9]'::vector
    ), 
-- Another CTE, named "full_text_search", is defined.
    full_text_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY 
        ts_rank(to_tsvector('english', content), 
            'cat & rat'::tsquery) DESC) AS rank
    FROM items
    WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
    ORDER BY ts_rank(
        to_tsvector('english', content), 'cat & rat'::tsquery) DESC
    )
-- The main query selects columns from both CTEs.
    SELECT 
        COALESCE(semantic_search.id, full_text_search.id) AS id,
        COALESCE(semantic_search.content, full_text_search.content) AS content,
        COALESCE(semantic_search.embedding, full_text_search.embedding) AS embedding,
        COALESCE(1.0 / (1 + semantic_search.rank), 0.0) + 
        COALESCE(1.0 / (1 + full_text_search.rank), 0.0) AS rank
    FROM semantic_search FULL OUTER JOIN full_text_search USING (id)
    ORDER BY rank DESC;

In this query, we first calculate the rank of the full text search result and the vector search result separately. Then we merge the two results together. We will get the following result:

 id |                  content                   |  embedding   |          rank          
----+--------------------------------------------+--------------+------------------------
  4 | a thin cat sat on a mat and ate a thin rat | [7, 8, 9]    | 1.00000000000000000000
  2 | a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]    | 0.70000000000000000000
  3 | a fat dog sat on a mat and ate a fat rat   | [4, 5, 6]    | 0.33333333333333333333
  5 | a thin dog sat on a mat and ate a thin rat | [10, 11, 12] | 0.33333333333333333333
(4 rows)

2 - Adaptive retrieval and shortening embeddings

OpenAI latest embedding models support shortening embeddings, where developers could simply “remove some numbers from the end of a sequence” and still maintain a valid representation for text.

This enables very flexible usage. For example, when using a vector data store that only supports embeddings up to 1024 dimensions long, developers can now still use our best embedding model text-embedding-3-large and specify a value of 1024 for the dimensions API parameter, which will shorten the embedding down from 3072 dimensions, trading off some accuracy in exchange for the smaller vector size.

Shortening embeddings

OpenAI’s new embedding model text-embedding-3-large produces embeddings with 3072 dimensions by default.

But you could safely remove some numbers from the end of the sequence and still maintain a valid representation for text. For example, you could shorten the embeddings to 1024 dimensions.

embedding_3072_dim_from_openai = [-0.00611, -0.00269, ...]
embedding_1024_dim_from_openai = [-0.01520, -0.03839, ...]

embedding_1024_dim_truncated_from_3072_dim = [-0.00611, -0.00269, ...]
embedding_1024_dim_normalized = [-0.01520, -0.03839, ...]

There are lots of benefits to using shortening embeddings:

• Reduced memory usage: Shortening embeddings can reduce the memory usage to store the embeddings.
• Lower dimensions: Shortening embeddings can reduce the dimensions of the embeddings. Some vector data stores only support embeddings up to a certain number of dimensions (e.g. 2000).
• Faster indexing and search: Smaller embeddings can be indexed and searched faster than larger embeddings.
• Trade-off performance and cost: Shortening embeddings can trade off some accuracy in exchange for the smaller vector size.

However, the accuracy of the shortened embeddings is not as good as the full embeddings.

Trade-off performance and cost

We have performed an experiment to evaluate the effectiveness of reduced embeddings compared to the complete embeddings. The benchmark was conducted using the dbpedia-entities-openai3-text-embedding-3-large-3072-1M dataset. The experiment was carried out on a Google Cloud VM with specifications of n2-standard-8, which includes 8 vCPUs and 32 GB of memory.

We have established a table and populated it with 1 million embeddings, each consisting of 3072 dimensions.

CREATE TABLE openai3072 (
  id bigserial PRIMARY KEY,
  text_embedding_3_large_3072_embedding vector(3072),
  text_embedding_3_large_3072_embedding bvector(3072)
);

from pgvecto_rs.sdk import PGVectoRs, Record
from pgvecto_rs.psycopg import register_vector
import psycopg

conn = psycopg.connect(conninfo='postgresql://postgres:mysecretpassword@localhost:5433/postgres', autocommit=True)

import json
import os

import numpy as np
import random
from datasets import load_dataset
from datasets.exceptions import DatasetNotFoundError
from tqdm import tqdm

MODEL_NAME, DIMENSIONS = "text-embedding-3-large", 3072
DATASET_NAME = f"Qdrant/dbpedia-entities-openai3-{MODEL_NAME}-{DIMENSIONS}-1M"
collection_name = f"dbpedia-{MODEL_NAME}-{DIMENSIONS}"
embedding_column_name = f"{MODEL_NAME}-{DIMENSIONS}-embedding"

dataset = load_dataset(
        DATASET_NAME,
        streaming=False,
        split="train",
    )

bs = 1000
for i, record in tqdm(enumerate(dataset)):
    if i % bs == 0:
        points = []
    points.append({
        "embedding": record[embedding_column_name]
    })
    if i % bs == bs - 1:
        batch_points = f", ".join([f"('{p['embedding']}')" for p in points])
        conn.execute(f"INSERT INTO openai3072 (text_embedding_3_large_3072_embedding) VALUES %s" % (batch_points))
    print(f"Inserted {i} records")

Furthermore, we have created indexes for embeddings with dimensions of 256, 1024 and 3072, and performed an ANN (Approximate Nearest Neighbors) query for each dimension. There is no need to create the 256-dimensional column, as we can use the first 256 dimensions of the 3072-dimensional embeddings.

-- Create 256-dim index
CREATE INDEX openai_vector_index_256 ON public.openai3072 USING vectors (((text_embedding_3_large_3072_embedding[0:256])::vector(256)) vector_l2_ops);

-- Create 1024-dim index
CREATE INDEX openai_vector_index_1024 on openai3072 using vectors((text_embedding_3_large_3072_embedding[0:1024]::vector(1024)) vector_l2_ops);

-- Create 3072-dim index
CREATE INDEX openai_vector_index on openai3072 using vectors(text_embedding_3_large_3072_embedding vector_l2_ops);

Additionally, we have constructed a binary vector index for the 3072-dimensional embeddings and conducted an ANN query using this index as well.

CREATE INDEX openai_vector_index_bvector ON public.openai3072 USING vectors (text_embedding_3_large_3072_bvector bvector_l2_ops);

We run the KNN query with different limits (5, 10, 50, 100) for the number of neighbors to retrieve.

export PGHOST="localhost"
export PGUSER="postgres"
export PGPASSWORD="mysecretpassword"
export PGDATABASE="postgres"
export PGPORT=5433
export VECTOR_DIM=3072
TEST_TIME=300
for clients in 2 4 8 16 32; do
    pgbench -n -T "${TEST_TIME}" \
      --file=query.sql \
      -c "${clients}" -j "${clients}"
done

\set random_id random(1, 1000000)

select id from openai3072 order by text_embedding_3_large_3072_embedding <-> (select text_embedding_3_large_3072_embedding from openai3072 where id = :random_id) limit 100;

The results of the experiment are shown here:

As anticipated, the 3072-dimensional embeddings exhibit the highest accuracy, while the 1024-dimensional embeddings demonstrate a slightly lower accuracy at around 85%. The 256-dimensional embeddings yield the lowest accuracy, approximately 65%. On the other hand, the binary vector index achieves an accuracy of about 80%.

Regarding the Requests Per Second (RPS) metric, the binary vector index showcases the most efficient performance, followed by the 256-dimensional embeddings, the 1024-dimensional embeddings, and finally, the 3072-dimensional embeddings.

Details

The memory usage of the indexes is an important aspect to take into account. Indexing 1 million 3072-dimensional binary vectors requires only around 600MB of memory. In contrast, the 3072-dimensional vector index consumes approximately 18GB of memory. This represents a significant difference in memory usage, as the 3072-dimensional vector index utilizes approximately 30x more memory than the binary vector index.

Improve the accuracy via adaptive retrieval

Lower dimension and binary vector indexes trade off accuracy for advantages such as higher RPS and lower memory usage compared to full dimension indexes. Adaptive retrieval techniques can be used to combine the strengths of both approaches, achieving an optimized solution that balances accuracy, memory usage, and query speed.

The logic behind adaptive retrieval is really simple. Let’s take get top 100 candidates as an example. We can perform the following steps:

1. Query the lower dimensional or binary vector index first to retrieve 200 candidates from the 1 million embeddings. This is a fast operation.
2. Rerank the candidates using a KNN query to retrieve the top 100 candidates. K-NN is well-suited for situations where smaller sets and precise similarity search are necessary, making it an excellent choice for reranking in this context.

The reranking step is a bit slower than the initial query, but it is still much faster than querying the full dimension index. It could be done in the pgvecto.rs:

CREATE OR REPLACE FUNCTION match_documents_adaptive(
  query_embedding vector(3072),
  match_count int
)
RETURNS SETOF openai3072
LANGUAGE SQL
AS $$
WITH shortlist AS (
  SELECT *
  FROM openai3072
  ORDER BY (text_embedding_3_large_3072_embedding[0:256])::vector(256) <-> (query_embedding[0:256])::vector(256)
  LIMIT match_count * 2
)
SELECT *
FROM shortlist
ORDER BY text_embedding_3_large_3072_embedding <-> query_embedding
LIMIT match_count;
$$;

The match_documents_adaptive function accepts a query embedding and a match count as input parameters. It first retrieves match_count * 2 candidates from the 1 million embeddings using the 256-dimensional index. Then, it reranks the candidates using a KNN query to retrieve the top match_count candidates.

The function for binary vector and 1024-dimensional embeddings can be implemented in a similar manner.

We conduct the benchmark again with the adaptive retrieval technique. We tag the adaptive retrieval technique as Adaptive Retrieval in the following figure.

Adding a reranking step significantly improves the accuracy of the lower dimensional and binary vector indexes.

The binary vector index with adaptive retrieval achieves an accuracy of around 95%, with only a slight decrease in RPS compared to the binary vector index without adaptive retrieval. Remarkably, the memory usage of the binary vector index with adaptive retrieval is only around 600MB, which is the same as the binary vector index without adaptive retrieval.

Details

Conclusion

Shortening embeddings can reduce the memory usage to store the embeddings, and speed up indexing and search. However, the accuracy of the shortened embeddings is not as good as the full embeddings. That’s why we need to use adaptive retrieval techniques to improve the accuracy of the lower dimensional indexes.

It seems that binary vector indexes with adaptive retrieval are the best choice for the text-embedding-3-large model. You could also consider 1024-dimensional indexes with adaptive retrieval if you need a higher accuracy (e.g. 99%).

PS: All the code and data used in this experiment are available in the pgvecto.rs-matryoshka-embeddings-benchmark.

3 - Image Search

Image intelligent search, also known as content-based image retrieval, goes beyond traditional keyword-based search and enables users to search for images based on their visual content. For example, you could search images that contain a red mushroom (image from immich):

The feature looks like a magic because it allows you to search for images that you don’t have keywords for. The search engine understands the content of the images and videos and can find the ones that are similar to the query.

How it works

Image intelligent search leverages advanced computer vision techniques and machine learning algorithms to understand the visual characteristics and features within images, allowing users to perform freeform searches without relying on specific keywords or metadata.

The first step is to extract meaningful features from images and generate a vector representation of each image. This is typically done using deep learning models such as CLIP. The vector representation is called an embedding. It is a list of numbers that represents the image.

Then you could use vector stores like pgvecto.rs to store the embeddings and perform efficient vector search.

Example

We will leverage CLIP, a neural network that connects text and images, to generate embeddings for our images. These embeddings capture the visual concepts present in the images. Subsequently, we will utilize pgvecto.rs, a scalable vector search tool built on Postgres, to store these embeddings and conduct efficient vector searches. By combining the power of CLIP and pgvecto.rs, we can retrieve similar images based on their visual content.

Generate embeddings

You could use CLIP to generate embeddings for images:

from transformers import CLIPProcessor, CLIPModel

processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

def compute_image_embeddings(list_of_images):
    return model.get_image_features(**processor(images=list_of_images, return_tensors="pt", padding=True))

Store embeddings

Let’s say you have a table named images with three columns: id, author and embedding. The column embedding is of type vector(512). You could insert an image into the table like this:

-- Create the table
CREATE TABLE images (
    id SERIAL PRIMARY KEY,
    author TEXT NOT NULL,
    vector vector(512) NOT NULL -- 3 dimensions
);

INSERT INTO images (author, embedding) VALUES
  ('Allen', '[...]');

Search for similar images

You could search for similar images like this:

SELECT * FROM images ORDER BY embedding <-> '[...]' LIMIT 5;

Advanced search

Assuming you want to search for the 64 nearest rows and the results don’t contain the images from the author Allen.

You may want the following SQL to work:

SET vectors.hnsw_ef_search = 64;
SELECT * FROM images ORDER BY embedding <-> '[...]' WHERE author != 'Allen' LIMIT 64;

In base mode, you may only get 32 rows since the HNSW algorithm bypasses filter conditions during the search and relies on PostgreSQL to apply filters post-search.

You could use the vbase search mode to achieve this. In vbase mode, searching results become an iterator that continuously returns results. The database can continuously retrieve results from it until the user’s query conditions are met. It’s quite different from other vector search library, such as faiss. The latter can only return a fixed-size result set, and this size needs to be determined before the search. For queries with filtering conditions, it is not possible to determine the size of the result set in advance. As a result, either insufficient results can be obtained or excessive searching will affect latency. The original idea comes from VBASE: Unifying Online Vector Similarity Search and Relational Queries via Relaxed Monotonicity.

You can enable vbase by a SQL statement SET vectors.search_mode = vbase;. Then you can get the correct behavior. This mode is used in immich’s image search feature. vbase is the default search mode in 0.2.

Real-world applications

Image intelligent search is used in immich to enable users to search for images based on their visual content. The code snippets provided here are inspired by the implementation of the image search feature in immich.

SELECT a.*, s.*, e.*
FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id
WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()
ORDER BY s.embedding <=> :embedding
LIMIT :numResults;

The query retrieves data from the asset_entity, smart_search, and exif_info tables based on specified conditions and ordering. Let’s break down the query:

SELECT a.*, s.*, e.*

This part of the query specifies the columns to be selected from the tables asset_entity, smart_search, and exif_info. The aliases a, s, and e are used to refer to each table respectively.

FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id

Here, the asset_entity table is aliased as a, the smart_search table as s, and the exif_info table as e. The INNER JOIN is used to combine asset_entity and smart_search based on the condition s.assetId = a.id. The LEFT JOIN is used to include matching records from the exif_info table based on the condition e.assetId = a.id.

CREATE TABLE asset_entity (
  id UUID PRIMARY KEY,
  deviceAssetId VARCHAR NOT NULL,
  ownerId VARCHAR NOT NULL,
  libraryId VARCHAR NOT NULL,
  deviceId VARCHAR NOT NULL,
  type VARCHAR NOT NULL,
  originalPath VARCHAR NOT NULL,
  resizePath VARCHAR,
  webpPath VARCHAR DEFAULT '',
  thumbhash BYTEA,
  encodedVideoPath VARCHAR,
  createdAt TIMESTAMPTZ NOT NULL,
  updatedAt TIMESTAMPTZ NOT NULL,
  deletedAt TIMESTAMPTZ,
  fileCreatedAt TIMESTAMPTZ NOT NULL,
  localDateTime TIMESTAMPTZ NOT NULL,
  fileModifiedAt TIMESTAMPTZ NOT NULL,
  isFavorite BOOLEAN DEFAULT false,
  isArchived BOOLEAN DEFAULT false,
  isExternal BOOLEAN DEFAULT false,
  isReadOnly BOOLEAN DEFAULT false,
  isOffline BOOLEAN DEFAULT false,
  checksum BYTEA NOT NULL,
  duration VARCHAR,
  isVisible BOOLEAN DEFAULT true,
  livePhotoVideoId UUID,
  originalFileName VARCHAR NOT NULL,
  sidecarPath VARCHAR,
  stackId UUID,
  CONSTRAINT fk_asset_owner FOREIGN KEY (ownerId) REFERENCES user_entity (id) ON DELETE CASCADE ON UPDATE CASCADE,
  CONSTRAINT fk_asset_library FOREIGN KEY (libraryId) REFERENCES library_entity (id) ON DELETE CASCADE ON UPDATE CASCADE,
  CONSTRAINT fk_asset_live_photo_video FOREIGN KEY (livePhotoVideoId) REFERENCES asset_entity (id) ON DELETE SET NULL ON UPDATE CASCADE,
  CONSTRAINT fk_asset_stack FOREIGN KEY (stackId) REFERENCES asset_stack_entity (id) ON DELETE SET NULL ON UPDATE CASCADE
);

WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()

In this part, the query applies filters to the asset_entity table. It selects rows where the ownerId column is among the specified userIds, isArchived is false, isVisible is true, and fileCreatedAt is earlier than the current time (NOW()).

ORDER BY s.embedding <=> :embedding

The query orders the results based on the embedding column of the smart_search table, comparing it to the provided :embedding value. The <=> operator is used for the comparison.

LIMIT :numResults;

Finally, limit the number of rows returned by the query to the value specified by the :numResults parameter.

This is an real-world use case to show how image intelligent search is implemented in a production system.

4 - Multi Tenancy

Multi-tenancy is essential for SaaS applications, allowing a single instance to serve multiple users or groups while ensuring data privacy and security. This blog outlines the design and implementation a simple multi-tenant vector search system with pgvecto.rs.

Designing the Database Schema

The success of user-based isolation in multi-tenant systems hinges on a well-structured database schema:

Users and Documents

Each user is identified by a unique user_id, which directly links them to their documents for private access. The documents table contains each user’s documents, each with a distinct id, title, and content. Additional metadata can be included in the documents table to improve search functionality. The user_id is linked to the documents instead of chunk embeddings to replicate real-world scenarios more accurately.

CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    user_id INTEGER NOT NULL,
    title TEXT,
    content TEXT,
    -- Additional metadata
);

Chunks

Embeddings store vectorized representations of chunks of documents, and use document_id to reference its document.

CREATE TABLE chunks (
    id SERIAL PRIMARY KEY,
    document_id INTEGER REFERENCES documents(id),
    embedding vector(512) NOT NULL
);

Entity-Relation Diagram

Implementing Multi-Tenancy

Create Index

Create vector index on chunks.embedding to accelerate query process. Here we use dot product as the similarity measure.

CREATE INDEX idx_chunks_embedding ON chunks USING vectors (embedding vector_dot_ops);

Run the query with user_id

SELECT chunks.id AS chunk_id, documents.id AS document_id, documents.title
FROM chunks
INNER JOIN documents ON chunks.document_id = documents.id
WHERE documents.user_id = 'alice' /* Replace with user_id */
ORDER BY chunks.embedding <-> '[3,2,1]' /* Replace with query embedding */ LIMIT 5;

In this query:

• The SELECT statement now includes chunks.id AS chunk_id, documents.id AS document_id, and documents.title to return the chunk ID, the document ID, and the document title, respectively.
• The INNER JOIN ensures that only chunks associated with documents owned by the specified user (‘alice’) are considered.
• The ORDER BY clause calculates the distance between the chunk’s embedding and the provided vector [3,2,1], sorting the results by similarity. The closest or most similar embeddings are returned first.
• Replace [3,2,1] with the query embedding and ‘specific_user_id’ with the actual ID of the user for whom you are performing the search.

Performance Tip

The vector search will scan through additional points and evaluate the conditions individually. Consequently, if there is a large number of user IDs or the filter criteria are challenging to meet, the query speed may be impacted. To enhance performance, you may want to explore utilizing PostgreSQL’s partition table function.

5 - Sparse Vector Search

In the past, hybrid search combined two search methods: traditional keyword-based search and vector-based similarity search. However, sparse vector can act as a substitute for keyword search, unlocking the full potential of the data pipeline with pure vector search.

This post explores how to use how BGE-M3 and pgvecto.rs to exploit the potential of hybrid search using both sparse and dense vectors.

Let’s dive in.

An Overview of Sparse and Dense Vector Embedding

Sparse vectors are high-dimensional vectors that contain few non-zero values. They are suitable for traditional information retrieval use cases. For example, a vector with 32,000 dimensions but only 2 non-zero elements is a typical sparse vector.

Dense vectors are embeddings from neural networks. They are generated by text embedding models and have most or all elements non-zero. They have fewer dimensions, such as 256 or 1536, much less than sparse vectors.

Dense vectors and sparse vectors are both used to represent features of text chunks, but they have significant differences. A single dimension in a dense vector embedding does not mean anything, as it is too abstract to determine its meaning. However, when we take all the dimensions together, they provide the semantic meaning of the input text.

Sparse vectors have dimensions equal to the vocabulary size, with each element representing the weight of a word in the text.

Generally, words with higher weights indicate greater importance in the text. Sparse vectors simplify the analysis and visualization of the relationship between text and vocabulary. To visualize dense vectors, additional dimension reduction algorithms are necessary, such as:

• Principal component analysis (PCA)
• T-distributed stochastic neighbor embedding (T-SNE)
• Uniform Manifold Approximation and Projection (UMAP)

Sparse vectors are generally less effective at capturing the relationships between words, particularly for synonyms and related terms, compared to dense vectors.

For instance, both rocket and SPACE X are related to space exploration, but Dense vectors are more likely to reveal similarities and assign a closer vector distance to corresponding texts.

How to create a vector embedding?

A vector embedding is the internal representation of input data in deep learning models, also known as embedding models. Most embedding models, such as text-embedding-3-small, only output dense embeddings.

SPLADE and BGE-M3 can generate sparse embeddings, sometimes called lexical weights. In this section, we will introduce the BGE-M3 model for generating dense and sparse vectors.

BGE-M3 is a multi-functionality model, which can simultaneously perform the three common retrieval functionalities of the embedding model: dense retrieval, multi-vector retrieval, and sparse retrieval. We will learn how to generate dense and sparse embeddings from it.

# pip install -U FlagEmbedding

from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)
sentences = ["What is BM25?", "Definition of BM25"]

output = model.encode(sentences, return_dense=True, return_sparse=True)

# The dense embedding of the two texts above:
dense_embedding = output['dense_vecs']
# [[-0.02501617 -0.04525582 -0.01972856 ... -0.0099566   0.01677509
#   0.01292699]
# [-0.0106499  -0.04731942 -0.01477693 ... -0.0149255   0.01040097
#   0.00965083]

sparse_weight_id = output['lexical_weights']

# You can see the weight for each token:
# print(model.convert_id_to_token(sparse_weight_id))
# 
# [{'What': 0.10430284, 'is': 0.10090457, 'BM': 0.2635918, '25': 0.3382988, '?': 0.052101523}, 
# {'Definition': 0.14298248, 'of': 0.076763615, 'BM': 0.2577639, '25': 0.33806682}]

For all sentences and words, a text-word matrix was created:

The matrix is highly sparse, especially when all vocabularies are included. Most words only appear in a few sentences, resulting in non-zero weights.

By feeding text into sparse embedding models like BGE-M3, we can easily obtain sparse embeddings from it.

# pip install -U numpy

import numpy as np

# Sparse vector embedding of text
sparse_indices = [[int(token_id) for token_id in text] for text in sparse_weight_id]
# [[4865, 83, 90017, 2588, 32], [155455, 111, 90017, 2588]]
sparse_values = [[float(text[token_id]) for token_id in text] for text in sparse_weight_id]
# [[0.10430284, 0.10090457, 0.2635918, 0.3382988, 0.052101523], [0.14298248, 0.076763615, 0.2577639, 0.33806682]]

# Sparse indexes must be ascending for each pgvecto.rs
for i in range(len(sparse_weight_id)):
    sparse_values[i] = [x for _, x in sorted(zip(sparse_indices[i], sparse_values[i]))]
    sparse_indices[i] = sorted(sparse_indices[i])

Now we have both the sparse and dense embeddings. It’s time to insert them into the vector database.

Storing vectors in pgvecto.rs

First, a deployed pgvecto.rs instance is required for vector storage and retrieval. This can be achieved by the official pgvecto-rs Docker image:

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

Now the only necessary dependency is installed. Let’s create a table to store the vectors.

pgvecto.rs supports dense or sparse vectors with different data types:

• vector for dense vector
• svector for sparse vector

The BGE-M3 model has a dense embedding dimension of 1024. The sparse embedding dimension is 250002, which matches the size of tokenizer vocabulary.

print(len(model.tokenizer))
# 250002

In this example, we create a table named documents with four columns: an ID column(id), a text column(text) and two vector columns(sparse for sparse vector and dense for dense vector).

# pip install -U psycopg

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS documents;")
    conn.execute(f"""
        CREATE TABLE documents (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            dense vector(1024) NOT NULL,
            sparse svector({len(model.tokenizer)}) NOT NULL);
    """
    )

With the table created, we can now insert embeddings into the table and create indexes for vector columns.

# pip install -U pgvecto_rs

import psycopg
from pgvecto_rs.psycopg import register_vector
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    register_vector(conn)
    for text, dense, sparse_ind, sparse_val in zip(sentences, dense_embedding, sparse_indices, sparse_values):
        conn.execute(
            "INSERT INTO documents (text, dense, sparse) VALUES (%s, %s, to_svector(%s, %s::real[]::int[], %s::real[]));",
            (text, dense, len(model.tokenizer), sparse_ind, sparse_val),
        )
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (sparse svector_dot_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (dense vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)

For dense vectors, vector_l2_ops calculates the Squared Euclidean distance, which is the most commonly used distance metric.

For sparse vectors, vector_dot_ops calculates the dot product, which is a more efficient method.

For more information about index types in pgvecto.rs, please check out our documentation.

Retrieving and Mixing Results

With the vectors stored in the database, we can efficiently query them using the pgvecto.rs extension. To find the nearest vectors to a given input text, we need to encode the text with the same model.

This process is identical to the one described in the previous chapter:

def embedding(sentence: str) -> tuple[np.ndarray, list[list[int]], list[list[float]]]:
    """
    Create dense and sparse embeddings from text
    """
    output = model.encode(sentence, return_dense=True, return_sparse=True, return_colbert_vecs=False)
    dense_embedding = output['dense_vecs']
    sparse_weight_id = output['lexical_weights']
    sparse_indices = [int(token_id) for token_id in sparse_weight_id]
    sparse_values = [float(sparse_weight_id[token_id]) for token_id in sparse_weight_id]

    sparse_values = [x for _, x in sorted(zip(sparse_indices, sparse_values))]
    sparse_indices = sorted(sparse_indices)
    return dense_embedding, sparse_indices, sparse_values

The pgvecto.rs extension provides a simple SELECT statement to get sparse and dense vectors like this:

dense_embedding, sparse_indices, sparse_values = embedding("The text you want to search...")

with psycopg.connect(URL) as conn:
    register_vector(conn)
    cur = conn.execute(
        "SELECT text FROM documents ORDER BY sparse <#> to_svector(%s, %s::real[]::int[], %s::real[]) LIMIT 50;",
        (len(model.tokenizer), sparse_indices, sparse_values),
    )
    sparse_result = cur.fetchall()
      
    cur = conn.execute(
              "SELECT text FROM documents ORDER BY dense <-> %s LIMIT 50;",
              (dense_embedding,),
          )
    dense_result = cur.fetchall()
    
# Merge sparse and dense vector search results
mix_text = set([r[0] for r in dense_result]).union([r[0] for r in sparse_result])

To rerank the merged result, you can use popular fusion methods such as Reciprocal Ranked Fusion (RRF). However, to make better use of BGE-M3, let’s introduce the Reranker model of BGE-M3. Unlike the embedding model, Reranker uses text input and directly outputs their similarity or score.

We can initialize a Reranker model and feed it with mixed candidate texts. With this method, reranking can be as simple as sorting by the Reranker model’s estimated score.

from FlagEmbedding import FlagReranker

reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True)

# Obtain relevance scores （Higher scores indicate greater relevance）
scores = reranker.compute_score([["The text you want to search...", candidate] for candidate in mix_text])
# Rerank text using combined dense and sparse embeddings
reranked_text = [t for _, t in sorted(zip(scores, mix_text))]

Congratulations! We have now completed our first step into the world of sparse vector search.

6 - Generative Question-Answering

Vector search can also be used to provide context to generative models such as OpenAI’s GPT series, improving the quality of the outputs.

This post will explore the use of pgvecto.rs in building and retrieving from a vector knowledge base to bolster the performance of generative models in question-answering tasks. A knowledge base serves as “long-term memory” for generative models, providing them persistent, curated, and accurate information as context to draw from when generating an answer to a question. This results in answers that are more accurate, reduces the possibility of the model producing “hallucinations” (statements with no basis in reality), and increases user trust in the outputs. This technique is known as Retrieval-Augmented Generation, or RAG: see this post or the related paper for more information.

Overview of the Question-Answering task

The generative model question-answering (QA) task is a method of Information Retrieval. The model is presented with a question, and generates an answer based on some context: typically, documents or other information sources the model can look at to retrieve the information necessary to answer the question.

This use case will involve giving the OpenAI gpt-3.5-turbo-instruct model context to answer questions about movies from the vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings Huggingface dataset.

The overall workflow will look something like this:

Question-Answering without Context

First, let’s try having the model answer some questions without providing it context.

We begin with defining a function to streamline getting a response for a given prompt:

# pip install openai

import openai

openai.api_key = "YOUR_API_KEY_HERE"
def complete(prompt):
    result = openai.Completion.create(
        engine='gpt-3.5-turbo-instruct',
        prompt=prompt,
        temperature=0,
        max_tokens=400,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None
    )
    return result['choices'][0]['text'].strip()

Let’s start with asking about a relatively popular movie:

base_prompt = "What is the release year and overall plot of Skyfall?"
print(complete(base_prompt))

Accurate, if somewhat vague in the plot summary. However, what about a movie that’s much less known?

base_prompt_2 = "What is the release year and overall plot of Cat Napping?"
print(complete(base_prompt_2))

We have found the limits of what the model “knows”: it is unable to produce the release year, and invents a potential plot based on the title. We will remedy this by creating a knowledge base for the model to draw on.

There are two options for allowing the model to better answer our question:

1. Fine-tune the model on text data in the domain of interest (movies in this case).
2. Employ retrieval augmented generation, which adds an information retrieval component to the process. This will allow us to retrieve and feed relevant information to the model as a secondary source of information.

We will be taking the second option in this post.

Building the Knowledge Base

Let’s build a knowledge base of movies to retrieve relevant information from, using the vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings Huggingface dataset. This dataset contains titles, release years, casts, Wikipedia pages, plot summaries/lengths, and vector embeddings of this information for 33,155 movies. Each movie also has text data, which encompasses title, release year, cast, and plot summary in one. We will be using the text and embeddings data for this application.

Load the Data

Start by loading the data:

# pip install -U datasets
from datasets import load_dataset
dataset = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings", split='train')

texts = dataset['train']['text']
# texts[0]
# 'Alice in Wonderland \n  starring May Clark \n  released in the year1903 \n  with the following plot  \n Alice follows a large white rabbit down a "Rabbit-hole". She finds a tiny door. When she finds a bottle labeled "Drink me", she does, and shrinks, but not enough to pass through the door. She then eats something labeled "Eat me" and grows larger. She finds a fan when enables her to shrink enough to get into the "Garden" and try to get a "Dog" to play with her. She enters the "White Rabbit\'s tiny House," but suddenly resumes her normal size. In order to get out, she has to use the "magic fan."\r\nShe enters a kitchen, in which there is a cook and a woman holding a baby. She persuades the woman to give her the child and takes the infant outside after the cook starts throwing things around. The baby then turns into a pig and squirms out of her grip. "The Duchess\'s Cheshire Cat" appears and disappears a couple of times to Alice and directs her to the Mad Hatter\'s "Mad Tea-Party." After a while, she leaves.\r\nThe Queen invites Alice to join the "ROYAL PROCESSION": a parade of marching playing cards and others headed by the White Rabbit. When Alice "unintentionally offends the Queen", the latter summons the "Executioner". Alice "boxes the ears", then flees when all the playing cards come for her. Then she wakes up and realizes it was all a dream.'

Generate Embeddings

This particular dataset comes with embeddings, so you can skip this step and use the embeddings directly.

embeddings = dataset['train']['embeddings']

If you want to generate embeddings from the text data, you can use the OpenAI, Cohere or other embeddings APIs. We provide an example using the Cohere API, you can get Cohere API key from here.

# pip install cohere
import getpass
import os
import cohere

if not os.getenv("COHERE_API_KEY"):
    os.environ["COHERE_API_KEY"] = getpass.getpass("Enter your Cohere API key: ")

co = cohere.Client()
# embedding model can refer to this documentation: https://docs.cohere.com/docs/cohere-embed
embed = co.embed(
    texts=texts,
    input_type="search_query",
    model="embed-english-v3.0"
)
embeddings = embed.embeddings

Insert Data into the Database

A deployed pgvecto.rs instance is required for vector storage and retrieval. We can use the official pgvecto-rs Docker image:

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

Install the required psycopg library to interact with Postgres, and ensure the extension is running:

#pip install "psycopg[binary,pool]"
import psycopg

URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"
#Create the table
with psycopg.connect(URL) as conn:
    conn.execute("DROP EXTENSION IF EXISTS vectors;")
    conn.execute("CREATE EXTENSION vectors;")

Now, we create a table movies with a few columns: an ID column(id), a text column(texts), and a vector column(embeddings).

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

#Create the table
with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS movies;")
    conn.execute(f"""
        CREATE TABLE movies (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            embedding vector(768) NOT NULL);
    """
    )

With the table created, we can now insert the embeddings from the dataset and create indexes for vector columns:

# Optional install, but recommended:
# pip install tqdm
import psycopg
from tqdm import tqdm
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

#Insert into table
with psycopg.connect(URL) as conn:
    with conn.cursor() as cursor:    
        for text, embedding in tqdm(zip(texts, embeddings), total=len(texts)):
            cursor.execute(
                """INSERT INTO movies (text, embedding) 
                    VALUES (%s, %s::real[]);""",
                (text, embedding),
            )
        conn.commit()

Create Index

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

#Creating indexes for embedding column
with psycopg.connect(URL) as conn:
    cursor = conn.cursor()
    cursor.execute("""
        CREATE INDEX ON movies 
        USING vectors (embedding vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.commit()

$$

D_{L2} = \Sigma (x_i - y_i) ^ 2

$$

Question-Answering with Context

To make use of our knowledge base, queries must first be encoded to determine what texts to retrieve:

# pip install sentence-transformers

from sentence_transformers import SentenceTransformer

#some options: 'paraphrase-mpnet-base-v2', 'all-mpnet-base-v2'
encoder = SentenceTransformer('all-mpnet-base-v2')

def encode_query(query): #Query is a string
    #Preprocess the query
    query = query.lower()
    query = query.split()
    query = [word for word in query]
    query = ' '.join(query)
    
    #then encode the query
    vector = encoder.encode(query)
    return vector

Now, we define a function to streamline adding the retrieved context to the initial prompt:

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

def add_context(prompt):
    search_vector = encode_query(prompt).tolist()
    # Construct prompt with the retrieved contexts included
    prompt_start = (
        "Answer the question based on the following context:\n\n"
    )
    prompt_end = (
        f"\n\nQuestion: {prompt}\nAnswer:"
    )
    with psycopg.connect(URL) as conn:
        cursor = conn.cursor()
        #Tune limit in accordance with how movie texts you want to give to the model
        cursor.execute("SELECT text FROM movies ORDER BY embedding <-> %s::real[] LIMIT 3;", 
                  (search_vector,),)
        result = cursor.fetchall()
    context = ""
    for r in result:
        context += r[0] + "\n\n"
    return prompt_start + context + prompt_end

Now, let’s try the prompt again, with added context:

base_prompt_2 = "What is the release year and overall plot of Cat Napping?"

context_prompt_2 = add_context(base_prompt_2)
context_response_2 = complete(context_prompt_2)
print(context_response_2)

Now the model has given us a more accurate answer! We can even ask further questions:

base_prompt_3 = "In Cat Napping, what does Tom do to Jerry?"
print(complete(add_context(base_prompt_3)))

As we saw earlier in this post, LLMs do work very well by themselves, but can struggle with niche or specific questions, creating “hallucinations” that may not be noticed. The use of the knowledge base improves factuality and subsequently user trust in the generated outputs.

Further Applications

This method is not only for asking about movies, of course. Other applications could include extracting summaries of scientific papers (try this dataset!), or answering questions about company policies provided a database of company literature.