feat: Adding RAG demo displaying Milvus usage for RAG

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
franciscojavierarceo · Jan 30, 2025 · Jan 25, 2025 · Jan 27, 2025 · Jan 27, 2025 · Jan 8, 2025
commit 006e897489eda7e107688223db5e4b6532627473
@@ -0,0 +1,29 @@
+# Feast Quickstart
+If you haven't already, check out the quickstart guide on Feast's website (http://docs.feast.dev/quickstart), which 
+uses this repo. A quick view of what's in this repository's `feature_repo/` directory:
+
+* `data/` contains raw demo parquet data
+* `feature_repo/example_repo.py` contains demo feature definitions
+* `feature_repo/feature_store.yaml` contains a demo setup configuring where data sources are
+* `feature_repo/test_workflow.py` showcases how to run all key Feast commands, including defining, retrieving, and pushing features. 
+
+You can run the overall workflow with `python test_workflow.py`.
+
+## To move from this into a more production ready workflow:
+> See more details in [Running Feast in production](https://docs.feast.dev/how-to-guides/running-feast-in-production)
+
+1. First: you should start with a different Feast template, which delegates to a more scalable offline store. 
+   - For example, running `feast init -t gcp`
+   or `feast init -t aws` or `feast init -t snowflake`. 
+   - You can see your options if you run `feast init --help`.
+2. `feature_store.yaml` points to a local file as a registry. You'll want to setup a remote file (e.g. in S3/GCS) or a 
+SQL registry. See [registry docs](https://docs.feast.dev/getting-started/concepts/registry) for more details. 
+3. This example uses a file [offline store](https://docs.feast.dev/getting-started/components/offline-store) 
+   to generate training data. It does not scale. We recommend instead using a data warehouse such as BigQuery, 
+   Snowflake, Redshift. There is experimental support for Spark as well.
+4. Setup CI/CD + dev vs staging vs prod environments to automatically update the registry as you change Feast feature definitions. See [docs](https://docs.feast.dev/how-to-guides/running-feast-in-production#1.-automatically-deploying-changes-to-your-feature-definitions).
+5. (optional) Regularly scheduled materialization to power low latency feature retrieval (e.g. via Airflow). See [Batch data ingestion](https://docs.feast.dev/getting-started/concepts/data-ingestion#batch-data-ingestion)
+for more details.
+6. (optional) Deploy feature server instances with `feast serve` to expose endpoints to retrieve online features.
+   - See [Python feature server](https://docs.feast.dev/reference/feature-servers/python-feature-server) for details.
+   - Use cases can also directly call the Feast client to fetch features as per [Feature retrieval](https://docs.feast.dev/getting-started/concepts/feature-retrieval)
diff --git a/examples/rag/feature_repo/data/driver_stats.parquet b/examples/rag/feature_repo/data/driver_stats.parquet
@@ -0,0 +1,30 @@
+from datetime import timedelta
+
+from feast import (
+    FeatureView,
+    Field,
+    FileSource,
+)
+from feast.data_format import ParquetFormat
+from feast.types import Float32, Array
+from feast import Entity
+
+item = Entity(name="item_id")
+
+parquet_file_path = "./data/city_wikipedia_summaries_with_embeddings.parquet"
+
+source = FileSource(
+    file_format=ParquetFormat(),
+    path=parquet_file_path,
+    timestamp_field="event_timestamp",
+)
+
+city_embeddings_feature_view = FeatureView(
+    name="city_embeddings",
+    entities=[item],
+    schema=[
+        Field(name="Embeddings", dtype=Array(Float32)),
+    ],
+    source=source,
+    ttl=timedelta(hours=2),
+)
@@ -0,0 +1,20 @@
+project: rag
+provider: local
+registry: data/registry.db
+online_store:
+  type: milvus
+  vector_enabled: true
+  vector_len: 384
+  host: 127.0.0.1
+#  port: 5432
+  port: 19530
+  username: "username"
+  password: "password"
+
+
+offline_store:
+  type: file
+entity_key_serialization_version: 3
+# By default, no_auth for authentication and authorization, other possible values kubernetes and oidc. Refer the documentation for more details.
+auth:
+    type: no_auth
@@ -0,0 +1,57 @@
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from feast import FeatureStore
+from transformers import AutoTokenizer, AutoModel
+from example_repo import city_embeddings_feature_view, item
+TOKENIZER = "sentence-transformers/all-MiniLM-L6-v2"
+MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+
+
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[
+        0
+    ]  # First element of model_output contains all token embeddings
+    input_mask_expanded = (
+        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    )
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
+        input_mask_expanded.sum(1), min=1e-9
+    )
+
+
+def run_model(sentences, tokenizer, model):
+    encoded_input = tokenizer(
+        sentences, padding=True, truncation=True, return_tensors="pt"
+    )
+    # Compute token embeddings
+    with torch.no_grad():
+        model_output = model(**encoded_input)
+
+    sentence_embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
+    sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
+    return sentence_embeddings
+
+def run_demo():
+    store = FeatureStore(repo_path=".")
+    df = pd.read_parquet("./data/city_wikipedia_summaries_with_embeddings.parquet")
+    store.apply([city_embeddings_feature_view, item])
+    store.write_to_online_store_async("city_embeddings", df)
+    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
+    )
+    print("features", features.to_df())
+
+
+if __name__ == "__main__":
+    run_demo()