RAG System Architecture
My RAG system architecture relies on a few major components:
- LlamaIndex: Orchestrates querying engine with accompanying prompts, system prompt templates, RAG index querying, intermediate embedding model requests, vector index store, and LLM querying.
- FIASS: FAISS is the vector search engine that powers the actual similarity search and retrieval of documents from our index. It integrates into LlamaIndex using the FaissVectorStore class, then stored using StorageContext, and finally used in our searchable index using the VectorStoreIndex class. FAISS is used to test the following 3 vector search methods:
- Flat: A flat index performs brute‑force search by computing the distance between the query embedding and every vector in the index.
- IVF + PQ: By first using IVF to narrow down to a few clusters and then applying PQ within those clusters, we achieve both high throughput and low memory usage without severely impacting recall. The composite IVFPQ index delivers speedups compared to non‑quantized brute‑force search
- HNSW32 + Flat: This is built on top of a Flat base storage index and starts at the top layer’s entry point, greedily navigates to the neighbor closest to the query embedding until no improvement is possible, then drops down one layer and repeats the greedy search, and continue until layer 0, where the final nearest neighbors are returned. It typically achieves polylogarithmic performance.
- Embedding models: 3 embedding models are used as part of our experiments each integrating with LlamaIndex through setting the ‘Settings.embed_model’ setting to an instance of the HuggingFaceEmbedding or GoogleGenAIEmbedding classes which are part of the LlamaIndex library. The three embedding models tested across the 3 previously mentioned vector search methods are:
- SentenceTransformers/all-MiniLM-L6-v2
- Embedding Size: 384 dimensions
- Model Size: 22M parameters (~90MB file size)
- MTEB: 56.09
- Efficiency: ~200MB VRAM on GPU with FP16
- BAAI/bge-large-en-v1.5:
- Embedding Size: 1024 dimensions
- Model Size: 335M parameters (~1.3GB FP32, ~639MB FP16)
- MTEB: 64.23
- Requires ~1.3GB VRAM (FP32) / ~640MB (FP16)
- Google/text-embedding-004:
- Embedding Size: 768 dimensions
- MTEB: 69.50
- LLM: Using my previously fine-tuned Llama 3.2 3B model, it integrates with LlamaIndex through the HuggingFaceLLM package to perform inference and generation of answers to the user’s queries to the RAG chain. The package’s backend uses accelerate.utils.modeling library for running the actual model. The specific configuration of the model:
- Context window: 128,000
- Max new tokens: 128
- Temperature: 0.75
- Repetition penalty: 1.15
Inference Performance
To keep this experiment simple, I used the following list of 10 generic and broadly related queries to benchmark both the performance of the vanilla, non-RAG, fine-tuned model as well as the performance of the fine-tuned model + each of the 3 embedding models + each of the 3 vector search methods to produce the results below. The queries used:
- What is the impact of climate change on agriculture?
- How does climate change affect biodiversity?
- What are the main causes of climate change?
- What are the potential solutions to climate change?
- How does climate change affect human health?
- What are the economic impacts of climate change?
- How does climate change affect water resources?
- What are the social impacts of climate change?
- How does climate change affect ecosystems?
- What are the political implications of climate change?
Notes:
- The indexes all contain the full 830 documents in our climate dataset in .txt format.
- The results below include performance in seconds and my human-rated accuracy for each query’s response as a score from 0-10, with 0 being the worst and 10 being the best accuracy scores respectively, as the measurement of how well these methods worked.
- The non-RAG run that only includes results for the fine-tuned model is run using the same LlamaIndex query engine system configs with a modified system prompt as it normalizes the LLM performance across RAG and non-RAG tests.
Non-RAG:
Performance:
- Average query processing time: 4.59 seconds
- Total time for all queries: 45.92 seconds
Accuracy: 4
RAG:
| Index Type | SentenceTransformers/all-MiniLM-L6-v2 | BAAI/bge-large-en-v1.5 | Google/text-embedding-004 |
|---|
| Flat | Performance: Average query processing time: 4.26 seconds Total time: 42.64 seconds Accuracy: 4 | Performance: Average query processing time: 4.04 seconds Total time: 40.37 seconds Accuracy: 7 | Performance: Average query processing time: 4.43 seconds Total time for all queries: 44.33 seconds Accuracy: 7 |
| IVF256,PQ32 | Performance: Average query processing time: 2.59 seconds Total time: 25.94 seconds Accuracy: 6 | Performance: Average query processing time: 3.79 seconds Total time: 37.86 seconds Accuracy: 8 | Performance: Average query processing time: 4.60 seconds Total time for all queries: 46.04 seconds Accuracy: 6 |
| HNSW32,Flat | Performance: Average query processing time: 3.87 seconds Total time: 38.73 seconds Accuracy: 5 | Performance: Average query processing time: 4.00 seconds Total time: 39.97 seconds Accuracy: 7 | Performance: Average query processing time: 4.82 seconds Total time for all queries: 48.20 seconds Accuracy: 9 |
How to run code
After installing the required dependencies based on the environment.yml file, run the following:
Note:
- For running NON-RAG experiment: use --non_rag option which runs non-RAG fine-tuned model alone and skips RAG results
- Provide a Google Gemini API key if testing the Google embedding model
- The script looks for an index folder in your ‘./’ current directory with the name format: {last name of selected model}_storage_faiss_{simplified vector search algorithm name}. If it finds this folder, it will load the index from it, otherwise if it does NOT find it, then it must build it from scratch. Examples of all folder names:
- all-MiniLM-L6-v2_storage_faiss_flat
- all-MiniLM-L6-v2_storage_faiss_hnsw
- all-MiniLM-L6-v2_storage_faiss_ivfpq
- bge-large-en-v1.5_storage_faiss_flat
- bge-large-en-v1.5_storage_faiss_hnsw
- bge-large-en-v1.5_storage_faiss_ivfpq
- text-embedding-004_storage_faiss_flat
- text-embedding-004_storage_faiss_hnsw
- text-embedding-004_storage_faiss_ivfpq
The default embed_model_type is sentence-transformers/all-MiniLM-L6-v2 and the default index_type is Flat. Below is the full list of all optional arguments that allow you to configure this script.
Here is the list of optional arguments:
-h, --help show this help message and exit
--model_path MODEL_PATH
Path to the HuggingFace LLM model directory.
--embed_model_type {1,2,3}
Type of embedding model to use: 1 = sentence-transformers/all-MiniLM-L6-v2, 2 = text-embedding-004 (Google),
3 = BAAI/bge-large-en-v1.5
--google_api_key GOOGLE_API_KEY
Google API Key for text-embedding-004. If not provided, attempts to read from GOOGLE_API_KEY environment
variable.
--data_dir DATA_DIR Directory containing the text data files.
--index_type {1,2,3} FAISS index type: 1 = Flat, 2 = IVF256,PQ32, 3 = HNSW32,Flat
--chunk_size CHUNK_SIZE
Size of text chunks for processing.
--chunk_overlap CHUNK_OVERLAP
Overlap between text chunks.
--embed_batch_size EMBED_BATCH_SIZE
Batch size for embedding generation (used for HuggingFace models).
--persist_dir_prefix PERSIST_DIR_PREFIX
Prefix for the persistence directory path.
--top_k TOP_K Number of top similar documents to retrieve for context.
--temperature TEMPERATURE
Sampling temperature for LLM generation.
--repetition_penalty REPETITION_PENALTY
Repetition penalty for LLM generation.
--max_new_tokens MAX_NEW_TOKENS
Maximum number of new tokens for the LLM to generate.
--non_rag If set, runs queries directly against the LLM without RAG context.