-
Notifications
You must be signed in to change notification settings - Fork 1.2k
feat: Add feast rag retriver functionality #5405
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
franciscojavierarceo
merged 1 commit into
feast-dev:master
from
Fiona-Waters:ragretriever
Jun 24, 2025
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| # End-to-end RAG example using Feast and Milvus. | ||
Fiona-Waters marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| ## Introduction | ||
| This example notebook provides a step-by-step demonstration of building and using a RAG system with Feast Feature Store and the custom FeastRagRetriever. The notebook walks through: | ||
Fiona-Waters marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| 1. Data Preparation | ||
| - Loads a subset of the Wikipedia DPR dataset (1% of training data) | ||
Fiona-Waters marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| - Implements text chunking with configurable chunk size and overlap | ||
| - Processes text into manageable passages with unique IDs | ||
|
|
||
| 2. Embedding Generation | ||
| - Uses `all-MiniLM-L6-v2` sentence transformer model | ||
| - Generates 384-dimensional embeddings for text passages | ||
| - Demonstrates batch processing with GPU support | ||
|
|
||
| 3. Feature Store Setup | ||
| - Creates a Parquet file as the historical data source | ||
| - Configures Feast with the feature repository | ||
| - Demonstrates writing embeddings from data source to Milvus online store which can be used for model training later | ||
|
|
||
| 4. RAG System Implementation | ||
| - **Embedding Model**: `all-MiniLM-L6-v2` (configurable) | ||
| - **Generator Model**: `granite-3.2-2b-instruct` (configurable) | ||
| - **Vector Store**: Custom implementation with Feast integration | ||
| - **Retriever**: Custom implementation extending HuggingFace's RagRetriever | ||
|
|
||
| 5. Query Demonstration | ||
| - Perform inference with retrieved context | ||
|
|
||
| ## Requirements | ||
| - A Kubernetes cluster with: | ||
| - GPU nodes available (for model inference) | ||
| - At least 200GB of storage | ||
| - A standalone Milvus deployment. See example [here](https://github.com/milvus-io/milvus-helm/tree/master/charts/milvus). | ||
|
|
||
| ## Running the example | ||
| Clone this repository: https://github.com/feast-dev/feast.git | ||
| Navigate to the examples/rag-retriever directory. Here you will find the following files: | ||
|
|
||
| * **feature_repo/feature_store.yaml** | ||
| This is the core configuration file for the RAG project's feature store, configuring a Milvus online store on a local provider. | ||
| * In order to configure Milvus you should: | ||
| - Update `feature_store.yaml` with your Milvus connection details: | ||
| - host | ||
| - port (default: 19530) | ||
| - credentials (if required) | ||
|
|
||
| * **__feature_repo/ragproject_repo.py__** | ||
| This is the Feast feature repository configuration that defines the schema and data source for Wikipedia passage embeddings. | ||
|
|
||
| * **__rag_feast.ipynb__** | ||
| This is a notebook demonstrating the implementation of a RAG system using Feast feature store. The notebook provides: | ||
Fiona-Waters marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| - A complete end-to-end example of building a RAG system with: | ||
| - Data preparation using the Wiki DPR dataset | ||
| - Text chunking and preprocessing | ||
| - Vector embedding generation using sentence-transformers | ||
| - Integration with Milvus vector store | ||
| - Inference utilising a custom RagRetriever: FeastRagRetriever | ||
| - Uses `all-MiniLM-L6-v2` for generating embeddings | ||
| - Implements `granite-3.2-2b-instruct` as the generator model | ||
|
|
||
| Open `rag_feast.ipynb` and follow the steps in the notebook to run the example. | ||
|
|
||
| ## FeastRagRetriver Low Level Design | ||
|
|
||
| <img src="images/FeastRagRetriever.png" width="800" height="450" alt="Low level design for feast rag retriever"> | ||
|
|
||
| ## Helpful Information | ||
| - Ensure your Milvus instance is properly configured and running | ||
| - Vector dimensions and similarity metrics can be adjusted in the feature store configuration | ||
| - The example uses Wikipedia data, but the system can be adapted for other datasets | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| project: ragproject | ||
| provider: local | ||
| registry: data/registry.db | ||
| online_store: | ||
| type: milvus | ||
| host: # Insert Milvus route host | ||
| username: # Insert Milvus username if required | ||
| password: # Insert Milvus password if required | ||
| port: 19530 | ||
| vector_enabled: true | ||
| embedding_dim: 384 | ||
| index_type: FLAT | ||
| metric_type: COSINE | ||
| offline_store: | ||
| type: file | ||
| entity_key_serialization_version: 3 | ||
| auth: | ||
| type: no_auth |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| from datetime import timedelta | ||
|
|
||
| from feast import Entity, FeatureView, Field, FileSource, ValueType | ||
| from feast.data_format import ParquetFormat | ||
| from feast.types import Array, Float32, String | ||
|
|
||
| # Define your entity (primary key for feature lookup) | ||
| wiki_passage = Entity( | ||
| name="passage_id", | ||
| join_keys=["passage_id"], | ||
| value_type=ValueType.STRING, | ||
| description="Unique ID of a Wikipedia passage", | ||
| ) | ||
|
|
||
| parquet_file_path = "data/wiki_dpr.parquet" | ||
|
|
||
| # Define offline source | ||
| wiki_dpr_source = FileSource( | ||
| name="wiki_dpr_source", | ||
| file_format=ParquetFormat(), | ||
| path=parquet_file_path, | ||
| timestamp_field="event_timestamp", | ||
| ) | ||
|
|
||
| # Define the feature view for the Wikipedia passage content | ||
| wiki_passage_feature_view = FeatureView( | ||
| name="wiki_passages", | ||
| entities=[wiki_passage], | ||
| ttl=timedelta(days=1), | ||
| schema=[ | ||
| Field( | ||
| name="passage_text", | ||
| dtype=String, | ||
| description="Content of the Wikipedia passage", | ||
| ), | ||
| Field( | ||
| name="embedding", | ||
| dtype=Array(Float32), | ||
| description="vectors", | ||
| vector_index=True, | ||
| vector_length=384, | ||
| vector_search_metric="COSINE", | ||
| ), | ||
| ], | ||
| online=True, | ||
| source=wiki_dpr_source, | ||
| description="Content features of Wikipedia passages", | ||
| ) |
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.