Upload 3 files

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
 import logging
-import threading
 import time
 from generator.compute_metrics import get_attributes_text
 from generator.generate_metrics import generate_metrics, retrieve_and_generate_response
 from config import AppConfig, ConfigConstants
 from generator.initialize_llm import initialize_generation_llm, initialize_validation_llm
-from generator.document_utils import get_logs, initialize_logging 
+from generator.document_utils import get_logs, initialize_logging
+from retriever.load_selected_datasets import load_selected_datasets 
 
 def launch_gradio(config : AppConfig):
     """
@@ -14,6 +14,9 @@ def launch_gradio(config : AppConfig):
     """
     initialize_logging()
 
+    # **🔹 Always get the latest loaded datasets**
+    config.detect_loaded_datasets()
+
     def update_logs_periodically():
         while True:
             time.sleep(2)  # Wait for 2 seconds
@@ -21,6 +24,10 @@ def launch_gradio(config : AppConfig):
 
     def answer_question(query, state):
         try:
+            # Ensure vector store is updated before use
+            if config.vector_store is None:
+                return "Please load a dataset first.", state
+            
             # Generate response using the passed objects
             response, source_docs = retrieve_and_generate_response(config.gen_llm, config.vector_store, query)
             
@@ -29,7 +36,7 @@ def launch_gradio(config : AppConfig):
             state["response"] = response
             state["source_docs"] = source_docs
             
-            response_text = f"Response: {response}\n\n"
+            response_text = f"Response from Model : {response}\n\n"
             return response_text, state
         except Exception as e:
             logging.error(f"Error processing query: {e}")
@@ -49,7 +56,7 @@ def launch_gradio(config : AppConfig):
             
             attributes_text = get_attributes_text(attributes)
 
-            metrics_text = "Metrics:\n"
+            metrics_text = ""
             for key, value in metrics.items():
                 if key != 'response':
                     metrics_text += f"{key}: {value}\n"
@@ -70,11 +77,14 @@ def launch_gradio(config : AppConfig):
         return get_updated_model_info()
 
     def get_updated_model_info():
+        loaded_datasets_str = ", ".join(config.loaded_datasets) if config.loaded_datasets else "None"
         """Generate and return the updated model information string."""
         return (
             f"Embedding Model: {ConfigConstants.EMBEDDING_MODEL_NAME}\n"
             f"Generation LLM: {config.gen_llm.name if hasattr(config.gen_llm, 'name') else 'Unknown'}\n"
+            f"Re-ranking LLM: {ConfigConstants.RE_RANKER_MODEL_NAME}\n"
             f"Validation LLM: {config.val_llm.name if hasattr(config.val_llm, 'name') else 'Unknown'}\n"
+            f"Loaded Datasets: {loaded_datasets_str}\n"
         )
 
     # Wrappers for event listeners
@@ -84,50 +94,125 @@ def launch_gradio(config : AppConfig):
     def reinitialize_val_llm(val_llm_name):
         return reinitialize_llm("validation", val_llm_name)
     
+    # Function to update query input when a question is selected from the dropdown
+    def update_query_input(selected_question):
+        return selected_question
+        
     # Define Gradio Blocks layout
     with gr.Blocks() as interface:
         interface.title = "Real Time RAG Pipeline Q&A"
-        gr.Markdown("# Real Time RAG Pipeline Q&A")  # Heading
+        gr.Markdown("""
+            # Real Time RAG Pipeline Q&A
+            The **Retrieval-Augmented Generation (RAG) Pipeline** combines retrieval-based and generative AI models to provide accurate and context-aware answers to your questions. 
+            It retrieves relevant documents from a dataset (e.g., COVIDQA, TechQA, FinQA) and uses a generative model to synthesize a response. 
+            Metrics are computed to evaluate the quality of the response and the retrieval process.
+            """)
+        # Model Configuration
+        with gr.Accordion("System Information", open=False):
+            with gr.Accordion("DataSet", open=False):
+                with gr.Row():
+                    dataset_selector = gr.CheckboxGroup(ConfigConstants.DATA_SET_NAMES, label="Select Datasets to Load")
+                    load_button = gr.Button("Load", scale= 0)
+
+            with gr.Row():
+                # Column for Generation Model Dropdown
+                with gr.Column(scale=1):
+                    new_gen_llm_input = gr.Dropdown(
+                        label="Generation Model",
+                        choices=ConfigConstants.GENERATION_MODELS,
+                        value=ConfigConstants.GENERATION_MODELS[0] if ConfigConstants.GENERATION_MODELS else None,
+                        interactive=True,
+                        info="Select the generative model for response generation."
+                    )
+                
+                # Column for Validation Model Dropdown
+                with gr.Column(scale=1):
+                    new_val_llm_input = gr.Dropdown(
+                        label="Validation Model",
+                        choices=ConfigConstants.VALIDATION_MODELS,
+                        value=ConfigConstants.VALIDATION_MODELS[0] if ConfigConstants.VALIDATION_MODELS else None,
+                        interactive=True,
+                        info="Select the model for validating the response quality."
+                    )
+                
+                # Column for Model Information
+                with gr.Column(scale=2):
+                    model_info_display = gr.Textbox(
+                        value=get_updated_model_info(),  # Use the helper function
+                        label="Model Configuration",
+                        interactive=False,  # Read-only textbox
+                        lines=5 
+                    )
         
-        # Textbox for new generation LLM name
+        # Query Section
+        gr.Markdown("Ask a question and get a response with metrics calculated from the RAG pipeline.")
+        all_questions = [
+            "When was the first case of COVID-19 identified?",
+            "What are the ages of the patients in this study?",
+            "Is one party required to deposit its source code into escrow with a third party, which can be released to the counterparty upon the occurrence of certain events (bankruptcy,  insolvency, etc.)?",
+            "Explain the concept of blockchain.",
+            "What is the capital of France?",
+            "Do Surface Porosity and Pore Size Influence Mechanical Properties and Cellular Response to PEEK??",
+            "How does a vaccine work?",
+            "What is the difference between RNA and DNA?",
+            "What are the risk factors for heart disease?",
+            "What is the role of insulin in the body?",
+            # Add more questions as needed
+        ]
+
+        # Subset of questions to display as examples
+        example_questions = [
+            "When was the first case of COVID-19 identified?",
+            "What are the ages of the patients in this study?",
+            "What is the Hepatitis C virus?",
+            "Explain the concept of blockchain.",
+            "What is the capital of France?"
+        ]  
         with gr.Row():
-            new_gen_llm_input = gr.Dropdown(
-                label="Generation Model",
-                choices=ConfigConstants.GENERATION_MODELS,  # Directly use the list
-                value=ConfigConstants.GENERATION_MODELS[0] if ConfigConstants.GENERATION_MODELS else None,  # First value dynamically
-                interactive=True
-            )
-
-            new_val_llm_input = gr.Dropdown(
-                label="Validation Model",
-                choices=ConfigConstants.VALIDATION_MODELS,  # Directly use the list
-                value=ConfigConstants.VALIDATION_MODELS[0] if ConfigConstants.VALIDATION_MODELS else None,  # First value dynamically
-                interactive=True
-            )
-
-            model_info_display = gr.Textbox(
-                value=get_updated_model_info(),  # Use the helper function
-                label="System Information",
-                interactive=False  # Read-only textbox
-            )
+            with gr.Column():
+                with gr.Row():
+                    query_input = gr.Textbox(
+                        label="Ask a question ",
+                        placeholder="Type your query here or select from examples/dropdown",
+                        lines=2
+                    )
+                with gr.Row():
+                    submit_button = gr.Button("Submit", variant="primary", scale=0)
+                    clear_query_button = gr.Button("Clear", scale=0)
+            with gr.Column():
+                gr.Examples(
+                    examples=example_questions,  # Make sure the variable name matches
+                    inputs=query_input,
+                    label="Try these examples:"
+                )
+                question_dropdown = gr.Dropdown(
+                    label="",
+                    choices=all_questions,
+                    interactive=True,
+                    info="Choose a question from the dropdown to populate the query box."
+                )
+        
+        # Attach event listener to dropdown
+        question_dropdown.change(
+            fn=update_query_input,
+            inputs=question_dropdown,
+            outputs=query_input
+        )
 
-        # State to store response and source documents
-        state = gr.State(value={"query": "","response": "", "source_docs": {}})
-        gr.Markdown("Ask a question and get a response with metrics calculated from the RAG pipeline.")  # Description
+        # Response and Metrics
         with gr.Row():
-            query_input = gr.Textbox(label="Ask a question", placeholder="Type your query here")
-        with gr.Row():
-            submit_button = gr.Button("Submit", variant="primary", scale = 0)  # Submit button
-            clear_query_button = gr.Button("Clear", scale = 0)  # Clear button
-        with gr.Row():
-            answer_output = gr.Textbox(label="Response", placeholder="Response will appear here")
+            answer_output = gr.Textbox(label="Response", placeholder="Response will appear here", lines=2)
         
         with gr.Row():
             compute_metrics_button = gr.Button("Compute metrics", variant="primary" , scale = 0)
             attr_output = gr.Textbox(label="Attributes", placeholder="Attributes will appear here")
             metrics_output = gr.Textbox(label="Metrics", placeholder="Metrics will appear here")
        
-        #with gr.Row():
+        # State to store response and source documents
+        state = gr.State(value={"query": "","response": "", "source_docs": {}})
+
+        # Pass config to update vector store
+        load_button.click(lambda datasets: (load_selected_datasets(datasets, config), get_updated_model_info()), inputs=dataset_selector)
         # Attach event listeners to update model info on change
         new_gen_llm_input.change(reinitialize_gen_llm, inputs=new_gen_llm_input, outputs=model_info_display)
         new_val_llm_input.change(reinitialize_val_llm, inputs=new_val_llm_input, outputs=model_info_display)
@@ -146,8 +231,9 @@ def launch_gradio(config : AppConfig):
         )
         
         # Section to display logs
-        with gr.Row():
-            log_section = gr.Textbox(label="Logs", interactive=False, visible=True, lines=10 , every=2)  # Log section
+        with gr.Accordion("View Live Logs", open=False):
+            with gr.Row():
+                log_section = gr.Textbox(label="Logs", interactive=False, visible=True, lines=10 , every=2)  # Log section
 
         # Update UI when logs_state changes
         interface.queue() 

config.py

@@ -1,18 +1,36 @@
-
-class ConfigConstants:
-    # Constants related to datasets and models
-    DATA_SET_NAMES = ['covidqa', 'cuad']#, 'techqa' 'delucionqa', 'emanual', 'expertqa', 'finqa', 'hagrid', 'hotpotqa', 'msmarco', 'pubmedqa', 'tatqa']
-    EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-MiniLM-L3-v2"
-    RE_RANKER_MODEL_NAME = 'cross-encoder/ms-marco-electra-base'
-    GENERATION_MODEL_NAME = 'mixtral-8x7b-32768'
-    VALIDATION_MODEL_NAME = 'llama3-70b-8192'
-    GENERATION_MODELS = ["llama3-8b-8192", "qwen-2.5-32b", "mixtral-8x7b-32768", "gemma2-9b-it" ]
-    VALIDATION_MODELS = ["llama3-70b-8192", "deepseek-r1-distill-llama-70b" ]
-    DEFAULT_CHUNK_SIZE = 1000
-    CHUNK_OVERLAP = 200
-
-class AppConfig:
-    def __init__(self, vector_store, gen_llm, val_llm):
-        self.vector_store = vector_store
-        self.gen_llm = gen_llm
-        self.val_llm = val_llm
+
+import os
+
+class ConfigConstants:
+    # Constants related to datasets and models
+    DATA_SET_PATH= '/persistent/local_datasets'
+    DATA_SET_NAMES = ['covidqa', 'cuad', 'techqa','delucionqa', 'emanual', 'expertqa', 'finqa', 'hagrid', 'hotpotqa', 'msmarco', 'pubmedqa', 'tatqa']
+    EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-MiniLM-L3-v2"
+    RE_RANKER_MODEL_NAME = 'cross-encoder/ms-marco-electra-base'
+    GENERATION_MODEL_NAME = 'mixtral-8x7b-32768'
+    VALIDATION_MODEL_NAME = 'llama3-70b-8192'
+    GENERATION_MODELS = ["llama3-8b-8192", "qwen-2.5-32b", "mixtral-8x7b-32768", "gemma2-9b-it" ]
+    VALIDATION_MODELS = ["llama3-70b-8192", "deepseek-r1-distill-llama-70b" ]
+    DEFAULT_CHUNK_SIZE = 1000
+    CHUNK_OVERLAP = 200
+
+class AppConfig:
+    def __init__(self, vector_store, gen_llm, val_llm):
+        self.vector_store = vector_store
+        self.gen_llm = gen_llm
+        self.val_llm = val_llm
+        self.loaded_datasets = self.detect_loaded_datasets()  # Auto-detect loaded datasets
+
+    @staticmethod
+    def detect_loaded_datasets():
+        print('Calling detect_loaded_datasets')
+        """Check which datasets are already stored locally."""
+        local_path = ConfigConstants.DATA_SET_PATH
+        if not os.path.exists(local_path):
+            return set()
+        
+        dataset_files = os.listdir(local_path)
+        loaded_datasets = {
+            file.replace("_test.pkl", "") for file in dataset_files if file.endswith("_test.pkl")
+        }
+        return loaded_datasets

main.py

@@ -1,64 +1,34 @@
-import logging
-from config import AppConfig, ConfigConstants
-from data.load_dataset import load_data
-from generator.compute_rmse_auc_roc_metrics import compute_rmse_auc_roc_metrics
-from retriever.chunk_documents import chunk_documents
-from retriever.embed_documents import embed_documents
-from generator.initialize_llm import initialize_generation_llm
-from generator.initialize_llm import initialize_validation_llm
-from app import launch_gradio
-
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-def main():
-    logging.info("Starting the RAG pipeline")
-
-    # Dictionary to store chunked documents
-    all_chunked_documents = []
-    datasets = {}
-
-    # Load multiple datasets
-    for data_set_name in ConfigConstants.DATA_SET_NAMES:
-        logging.info(f"Loading dataset: {data_set_name}")
-        datasets[data_set_name] = load_data(data_set_name)
-
-        # Set chunk size based on dataset name
-        chunk_size = ConfigConstants.DEFAULT_CHUNK_SIZE
-        if data_set_name == 'cuad':
-            chunk_size = 4000  # Custom chunk size for 'cuad'
-        
-        # Chunk documents
-        chunked_documents = chunk_documents(datasets[data_set_name], chunk_size=chunk_size, chunk_overlap=ConfigConstants.CHUNK_OVERLAP)
-        all_chunked_documents.extend(chunked_documents)  # Combine all chunks
-
-    # Access individual datasets
-    #for name, dataset in datasets.items():
-        #logging.info(f"Loaded {name} with {dataset.num_rows} rows")
-    
-    # Logging final count
-    logging.info(f"Total chunked documents: {len(all_chunked_documents)}")
-    
-    # Embed the documents
-    vector_store = embed_documents(all_chunked_documents)
-    logging.info("Documents embedded")
-    
-     # Initialize the Generation LLM
-    gen_llm = initialize_generation_llm(ConfigConstants.GENERATION_MODEL_NAME)
-
-    # Initialize the Validation LLM
-    val_llm = initialize_validation_llm(ConfigConstants.VALIDATION_MODEL_NAME)
-
-    #Compute RMSE and AUC-ROC for entire dataset
-    #Enable below code for calculation
-    #data_set_name = 'covidqa'
-    #compute_rmse_auc_roc_metrics(gen_llm, val_llm, datasets[data_set_name], vector_store, 10)
-    
-    # Launch the Gradio app
-    config = AppConfig(vector_store= vector_store, gen_llm = gen_llm, val_llm = val_llm)
-    launch_gradio(config)
-
-    logging.info("Finished!!!")
-
-if __name__ == "__main__":
+import logging
+from config import AppConfig, ConfigConstants
+from generator.compute_rmse_auc_roc_metrics import compute_rmse_auc_roc_metrics
+from retriever.load_selected_datasets import load_selected_datasets
+from generator.initialize_llm import initialize_generation_llm
+from generator.initialize_llm import initialize_validation_llm
+from app import launch_gradio
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+def main():
+    logging.info("Starting the RAG pipeline")
+
+    # Initialize the Generation LLM
+    gen_llm = initialize_generation_llm(ConfigConstants.GENERATION_MODEL_NAME)
+
+    # Initialize the Validation LLM
+    val_llm = initialize_validation_llm(ConfigConstants.VALIDATION_MODEL_NAME)
+
+    #Compute RMSE and AUC-ROC for entire dataset
+    #Enable below code for calculation
+    #data_set_name = 'covidqa'
+    #compute_rmse_auc_roc_metrics(gen_llm, val_llm, datasets[data_set_name], vector_store, 10)
+    
+    # Launch the Gradio app
+    config = AppConfig(vector_store = None, gen_llm = gen_llm, val_llm = val_llm)
+    load_selected_datasets(['covidqa'], config)
+    launch_gradio(config)
+
+    logging.info("Finished!!!")
+
+if __name__ == "__main__":
     main()