This is a demo to show how you can use Feast to do RAG

This demo assumes you have Pyenv (2.3.10) and Poetry (1.4.1) installed on your machine as well as Python 3.9.

pyenv local 3.9
poetry shell
poetry install

To fetch the data simply run

python pull_states.py

Which will output a file called city_wikipedia_summaries.csv.

Then run

python batch_score_documents.py

Which will output data to data/city_wikipedia_summaries_with_embeddings.parquet

Next we'll need to do some Feast work and move the data into a repo created by Feast.

To get started, make sure to have Feast installed and PostGreSQL.

First run

cp ./data feature_repo/

And then open the module_4.ipynb notebook and follow those instructions.

It will walk you through a trivial tutorial to retrieve the top k most similar documents using PGVector.

The overview is relatively simple, the goal is to define an architecture to support the following:

flowchart TD;
    A[Pull Data] --> B[Batch Score Embeddings];
    B[Batch Score Embeddings] --> C[Materialize Online];
    C[Materialize Online] --> D[Retrieval Augmented Generation];

The simple demo shows the code below with the retrieved data shown.

import pandas as pd

from feast import FeatureStore
from batch_score_documents import run_model, TOKENIZER, MODEL
from transformers import AutoTokenizer, AutoModel

df = pd.read_parquet("./feature_repo/data/city_wikipedia_summaries_with_embeddings.parquet")

store = FeatureStore(repo_path=".")

# Prepare a query vector
question = "the most populous city in the U.S. state of Texas?"

tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
model = AutoModel.from_pretrained(MODEL)
query_embedding = run_model(question, tokenizer, model)
query = query_embedding.detach().cpu().numpy().tolist()[0]

# Retrieve top k documents
features = store.retrieve_online_documents(
    feature="city_embeddings:Embeddings",
    query=query,
    top_k=3
)

And running features_df.head() will show:

features_df.head()
        Embeddings	                                        distance
0	[0.11749928444623947, -0.04684492573142052, 0....	0.935567
1	[0.10329511761665344, -0.07897591590881348, 0....	0.939936
2	[0.11634305864572525, -0.10321836173534393, -0...	0.983343