Multiple data set support

app.py

@@ -1,71 +1,134 @@
 import gradio as gr
 import logging
-from generator.compute_rmse_auc_roc_metrics import compute_rmse_auc_roc_metrics
+import threading
+import time
+from generator.compute_metrics import get_attributes_text
+from generator.generate_metrics import generate_metrics, retrieve_and_generate_response
+from io import StringIO
 
-def launch_gradio(vector_store, dataset, gen_llm, val_llm):
+def launch_gradio(vector_store, gen_llm, val_llm):
     """
     Launch the Gradio app with pre-initialized objects.
     """
-    def answer_question_with_metrics(query):
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+
+    # Create a list to store logs
+    logs = []
+
+    # Custom log handler to capture logs and add them to the logs list
+    class LogHandler(logging.Handler):
+        def emit(self, record):
+            log_entry = self.format(record)
+            logs.append(log_entry)
+
+    # Add custom log handler to the logger
+    log_handler = LogHandler()
+    log_handler.setFormatter(logging.Formatter('%(asctime)s - %(message)s'))
+    logger.addHandler(log_handler)
+
+    def log_updater():
+        """Background function to add logs."""
+        while True:
+            time.sleep(2)  # Update logs every 2 seconds
+            pass  # Log capture is now handled by the logging system
+
+    def get_logs():
+        """Retrieve logs for display."""
+        return "\n".join(logs[-50:])  # Only show the last 50 logs for example
+
+    # Start the logging thread
+    threading.Thread(target=log_updater, daemon=True).start()
+
+    def answer_question(query, state):
         try:
-            logging.info(f"Processing query: {query}")
+            # Generate response using the passed objects
+            response, source_docs = retrieve_and_generate_response(gen_llm, vector_store, query)
             
-            # Generate metrics using the passed objects
-            from main import generate_metrics
-            response, metrics = generate_metrics(gen_llm, val_llm, vector_store, query, 1)
+            # Update state with the response and source documents
+            state["query"] = query
+            state["response"] = response
+            state["source_docs"] = source_docs
             
             response_text = f"Response: {response}\n\n"
-            metrics_text = "Metrics:\n"
-            for key, value in metrics.items():
-                if key != 'response':
-                    metrics_text += f"{key}: {value}\n"
-            
-            return response_text, metrics_text
+            return response_text, state
         except Exception as e:
             logging.error(f"Error processing query: {e}")
-            return f"An error occurred: {e}"
+            return f"An error occurred: {e}", state
 
-    def compute_and_display_metrics():
+    def compute_metrics(state):
         try:
-            # Call the function to compute metrics
-            relevance_rmse, utilization_rmse, adherence_auc = compute_rmse_auc_roc_metrics(
-                gen_llm, val_llm, dataset, vector_store, 10
-            )
+            logging.info(f"Computing metrics")
+            
+            # Retrieve response and source documents from state
+            response = state.get("response", "")
+            source_docs = state.get("source_docs", {})
+            query = state.get("query", "")
+
+            # Generate metrics using the passed objects
+            attributes, metrics = generate_metrics(val_llm, response, source_docs, query, 1)
+            
+            attributes_text = get_attributes_text(attributes)
+
+            metrics_text = "Metrics:\n"
+            for key, value in metrics.items():
+                if key != 'response':
+                    metrics_text += f"{key}: {value}\n"
             
-            # Format the result for display
-            result = (
-                f"Relevance RMSE Score: {relevance_rmse}\n"
-                f"Utilization RMSE Score: {utilization_rmse}\n"
-                f"Overall Adherence AUC-ROC: {adherence_auc}\n"
-            )
-            return result
+            return attributes_text, metrics_text
         except Exception as e:
-            logging.error(f"Error during metrics computation: {e}")
-            return f"An error occurred: {e}"
+            logging.error(f"Error computing metrics: {e}")
+            return f"An error occurred: {e}", ""
 
     # 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("Ask a question and get a response with metrics calculated from the RAG pipeline.")  # Description
         
+        # Section to display LLM names
+        with gr.Row():
+            model_info = f"Generation LLM: {gen_llm.name if hasattr(gen_llm, 'name') else 'Unknown'}\n"
+            model_info += f"Validation LLM: {val_llm.name if hasattr(val_llm, 'name') else 'Unknown'}\n"
+            gr.Textbox(value=model_info, label="Model Information", interactive=False)  # Read-only textbox
+
+        # 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
         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")  # Submit button
             clear_query_button = gr.Button("Clear")  # Clear button
-            submit_button = gr.Button("Submit", variant="primary") # Submit button
         with gr.Row():
             answer_output = gr.Textbox(label="Response", placeholder="Response will appear here")
+        
         with gr.Row():
+            compute_metrics_button = gr.Button("Compute metrics", variant="primary")
+            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():
-            compute_rmse_button = gr.Button("Compute RMSE & AU-ROC", variant="primary")
-            rmse_output = gr.Textbox(label="RMSE & AU-ROC Score", placeholder="RMSE & AU-ROC score will appear here")
-    
-        
+       
+        #with gr.Row():
+
         # Define button actions
-        submit_button.click(fn=answer_question_with_metrics, inputs=[query_input], outputs=[answer_output, metrics_output])
-        clear_query_button.click(fn=lambda: "", outputs=[query_input])  # Clear query input 
-        compute_rmse_button.click(fn=compute_and_display_metrics, outputs=[rmse_output])
+        submit_button.click(
+            fn=answer_question,
+            inputs=[query_input, state],
+            outputs=[answer_output, state]
+        )
+        clear_query_button.click(fn=lambda: "", outputs=[query_input])  # Clear query input
+        compute_metrics_button.click(
+            fn=compute_metrics,
+            inputs=[state],
+            outputs=[attr_output, metrics_output]
+        )
+
+        # Section to display logs
+        with gr.Row():
+            start_log_button = gr.Button("Start Log Update", elem_id="start_btn")  # Button to start log updates
+        with gr.Row():
+            log_section = gr.Textbox(label="Logs", interactive=False, visible=True, lines=10)  # Log section
+
+        # Set button click to trigger log updates
+        start_log_button.click(fn=get_logs, outputs=log_section)
 
-    interface.launch()
+    interface.launch()

generator/compute_metrics.py

@@ -32,18 +32,53 @@ def compute_metrics(attributes, total_sentences):
 
 def get_metrics(attributes, total_sentences):
     if attributes.content:
-        #print(attributes)
-        result_content = attributes.content  # Access the content attribute
-        # Extract the JSON part from the result_content
-        json_start = result_content.find("{")
-        json_end = result_content.rfind("}") + 1
-        json_str = result_content[json_start:json_end]
-        
         try:
+            result_content = attributes.content  # Access the content attribute
+            # Extract the JSON part from the result_content
+            json_start = result_content.find("{")
+            json_end = result_content.rfind("}") + 1
+            json_str = result_content[json_start:json_end]
             result_json = json.loads(json_str)
             # Compute metrics using the extracted attributes
             metrics = compute_metrics(result_json, total_sentences)
             logging.info(metrics)
+            
             return metrics        
         except json.JSONDecodeError as e:
-            logging.error(f"JSONDecodeError: {e}")
+            logging.error(f"JSONDecodeError: {e}")
+
+def get_attributes_text(attributes):
+        try:
+            result_content = attributes.content  # Access the content attribute
+            # Extract the JSON part from the result_content
+            json_start = result_content.find("{")
+            json_end = result_content.rfind("}") + 1
+            json_str = result_content[json_start:json_end]
+            result_json = json.loads(json_str)
+            
+            # Extract the required fields from json
+            relevance_explanation = result_json.get("relevance_explanation", "N/A")
+            all_relevant_sentence_keys = result_json.get("all_relevant_sentence_keys", [])
+            overall_supported_explanation = result_json.get("overall_supported_explanation", "N/A")
+            overall_supported = result_json.get("overall_supported", "N/A")
+            sentence_support_information = result_json.get("sentence_support_information", [])
+            all_utilized_sentence_keys = result_json.get("all_utilized_sentence_keys", [])
+
+            # Format the metrics for display
+            attributes_text = "Attributes:\n"
+            attributes_text = f"### Relevance Explanation:\n{relevance_explanation}\n\n"
+            attributes_text += f"### All Relevant Sentence Keys:\n{', '.join(all_relevant_sentence_keys)}\n\n"
+            attributes_text += f"### Overall Supported Explanation:\n{overall_supported_explanation}\n\n"
+            attributes_text += f"### Overall Supported:\n{overall_supported}\n\n"
+            attributes_text += "### Sentence Support Information:\n"
+            for info in sentence_support_information:
+                attributes_text += f"- Response Sentence Key: {info.get('response_sentence_key', 'N/A')}\n"
+                attributes_text += f"  Explanation: {info.get('explanation', 'N/A')}\n"
+                attributes_text += f"  Supporting Sentence Keys: {', '.join(info.get('supporting_sentence_keys', []))}\n"
+                attributes_text += f"  Fully Supported: {info.get('fully_supported', 'N/A')}\n"
+            attributes_text += f"\n### All Utilized Sentence Keys:\n{', '.join(all_utilized_sentence_keys)}"
+
+            return attributes_text
+        except Exception as e:
+            logging.error(f"Error extracting attributes: {e}")
+            return f"An error occurred while extracting attributes: {e}"

generator/compute_rmse_auc_roc_metrics.py

@@ -1,6 +1,6 @@
 
 from sklearn.metrics import roc_auc_score, root_mean_squared_error
-from generator.generate_metrics import generate_metrics
+from generator.generate_metrics import generate_metrics, retrieve_and_generate_response
 import logging
 
 def compute_rmse_auc_roc_metrics(gen_llm, val_llm, dataset, vector_store, num_question):
@@ -25,7 +25,8 @@ def compute_rmse_auc_roc_metrics(gen_llm, val_llm, dataset, vector_store, num_qu
         query = document['question']
         logging.info(f'Query number: {i + 1}')
         # Call the generate_metrics for each query
-        response, metrics = generate_metrics(gen_llm, val_llm, vector_store, query, 15)
+        response, source_docs = retrieve_and_generate_response(gen_llm, vector_store, query)
+        attributes, metrics = generate_metrics(val_llm, response, source_docs, query, 25)
         
         # Extract predicted metrics (ensure these are continuous if possible)
         predicted_relevance = metrics.get('Context Relevance', 0) if metrics else 0

generator/generate_metrics.py

@@ -5,7 +5,7 @@ from retriever.retrieve_documents import retrieve_top_k_documents
 from generator.compute_metrics import get_metrics
 from generator.extract_attributes import extract_attributes
 
-def generate_metrics(gen_llm, val_llm, vector_store, query, time_to_wait):
+def retrieve_and_generate_response(gen_llm, vector_store, query):
     logging.info(f'Query: {query}')
     
     # Step 1: Retrieve relevant documents for given query
@@ -21,6 +21,10 @@ def generate_metrics(gen_llm, val_llm, vector_store, query, time_to_wait):
 
     logging.info(f"Response from LLM: {response}")
 
+    return response, source_docs
+
+def generate_metrics(val_llm, response, source_docs, query, time_to_wait):
+
     # Add a sleep interval to avoid hitting the rate limit
     time.sleep(time_to_wait)  # Adjust the sleep time as needed
     
@@ -28,8 +32,8 @@ def generate_metrics(gen_llm, val_llm, vector_store, query, time_to_wait):
     logging.info(f"Extracting attributes through validation LLM")
     attributes, total_sentences = extract_attributes(val_llm, query, source_docs, response)
     logging.info(f"Extracted attributes successfully")
-
+    
     # Step 4 : Call the get metrics calculate metrics
     metrics = get_metrics(attributes, total_sentences)
 
-    return response, metrics
+    return attributes, metrics

generator/initialize_llm.py

@@ -4,8 +4,9 @@ from langchain_groq import ChatGroq
 
 def initialize_generation_llm():
     os.environ["GROQ_API_KEY"] = "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"
-    model_name = "llama3-8b-8192"
+    model_name = "mixtral-8x7b-32768"
     llm = ChatGroq(model=model_name, temperature=0.7)
+    llm.name = model_name
     logging.info(f'Generation LLM {model_name} initialized')
     return llm
 
@@ -13,5 +14,6 @@ def initialize_validation_llm():
     os.environ["GROQ_API_KEY"] = "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"
     model_name = "llama3-70b-8192"
     llm = ChatGroq(model=model_name, temperature=0.7)
+    llm.name = model_name
     logging.info(f'Validation LLM {model_name} initialized')
     return llm

main.py

@@ -3,7 +3,6 @@ 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.generate_metrics import generate_metrics
 from generator.initialize_llm import initialize_generation_llm
 from generator.initialize_llm import initialize_validation_llm
 from app import launch_gradio
@@ -13,21 +12,43 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 
 def main():
     logging.info("Starting the RAG pipeline")
-    data_set_name = 'covidqa'
+    
+
+    # Load single dataset
+    #dataset = load_data(data_set_name)
+    #logging.info("Dataset loaded")
+    # List of datasets to load
+    data_set_names = ['covidqa', 'techqa', 'cuad']
 
-    # Load the dataset
-    dataset = load_data(data_set_name)
-    logging.info("Dataset loaded")
+    default_chunk_size = 1000
+    chunk_overlap = 200
 
-    # Chunk the dataset
-    chunk_size = 1000  # default value
-    if data_set_name == 'cuad':
-        chunk_size = 3000
-    documents = chunk_documents(dataset, chunk_size)
-    logging.info("Documents chunked")
+    # Dictionary to store chunked documents
+    all_chunked_documents = []
+    # Load multiple datasets
+    datasets = {}
+    for data_set_name in 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 = 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=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(documents)
+    vector_store = embed_documents(all_chunked_documents)
     logging.info("Documents embedded")
     
      # Initialize the Generation LLM
@@ -36,18 +57,12 @@ def main():
     # Initialize the Validation LLM
     val_llm = initialize_validation_llm()
 
-    # Sample question
-    #row_num = 30
-    #query = dataset[row_num]['question']
-
-    # Call generate_metrics for above sample question
-    #generate_metrics(gen_llm, val_llm, vector_store, query)
-    
     #Compute RMSE and AUC-ROC for entire dataset
-    #compute_rmse_auc_roc_metrics(gen_llm, val_llm, dataset, vector_store, 10)
+    data_set_name = 'covidqa'
+    #compute_rmse_auc_roc_metrics(gen_llm, val_llm, datasets[data_set_name], vector_store, 10)
     
     # Launch the Gradio app
-    launch_gradio(vector_store, dataset, gen_llm, val_llm)
+    launch_gradio(vector_store, gen_llm, val_llm)
 
     logging.info("Finished!!!")
 

retriever/retrieve_documents.py

@@ -1,2 +1,78 @@
+import numpy as np
+from transformers import pipeline
+
 def retrieve_top_k_documents(vector_store, query, top_k=5):
-    return vector_store.similarity_search(query, k=top_k)
+    documents = vector_store.similarity_search(query, k=top_k)
+    documents = rerank_documents(query, documents)
+    return documents 
+
+# Reranking: Cross-Encoder for refining top-k results
+def rerank_documents(query, documents, reranker_model_name="cross-encoder/ms-marco-electra-base"):
+    """
+    Re-rank documents using a cross-encoder model.
+
+    Parameters:
+        query (str): The user's query.
+        documents (list): List of LangChain Document objects.
+        reranker_model_name (str): Hugging Face model name for re-ranking.
+
+    Returns:
+        list: Re-ranked list of Document objects with updated scores.
+    """
+    # Initialize the cross-encoder model
+    reranker = pipeline("text-classification", model=reranker_model_name, return_all_scores=False)
+
+    # Pair the query with each document's text
+    rerank_inputs = [{"text": query, "text_pair": doc.page_content} for doc in documents]
+
+    # Get relevance scores for each query-document pair
+    scores = reranker(rerank_inputs)
+
+    # Attach the new scores to the documents
+    for doc, score in zip(documents, scores):
+        doc.metadata["rerank_score"] = score["score"]  # Add score to document metadata
+
+    # Sort documents by the rerank_score in descending order
+    documents = sorted(documents, key=lambda x: x.metadata.get("rerank_score", 0), reverse=True)
+    return documents
+
+
+# Query Handling: Retrieve top-k candidates using FAISS with IVF index not used only for learning
+def retrieve_top_k_documents_manual(vector_store, query, top_k=5):
+    """
+    Retrieve top-k documents using FAISS index and optionally rerank them.
+
+    Parameters:
+        vector_store (FAISS): The vector store containing the FAISS index and docstore.
+        query (str): The user's query string.
+        top_k (int): The number of top results to retrieve.
+        reranker_model_name (str): The Hugging Face model name for cross-encoder reranking.
+
+    Returns:
+        list: Top-k retrieved and reranked documents.
+    """
+    # Encode the query into a dense vector
+    embedding_model = vector_store.embedding_function
+    query_vector = embedding_model.embed_query(query)  # Encode the query
+    query_vector = np.array([query_vector]).astype('float32')
+    
+    # Search the FAISS index for top_k results
+    distances, indices = vector_store.index.search(query_vector, top_k)
+
+    # Retrieve documents from the docstore
+    documents = []
+    for idx in indices.flatten():
+        if idx == -1:  # FAISS can return -1 for invalid indices
+            continue
+        doc_id = vector_store.index_to_docstore_id[idx]
+
+        # Access the internal dictionary of InMemoryDocstore
+        internal_docstore = getattr(vector_store.docstore, "_dict", None)
+        if internal_docstore and doc_id in internal_docstore:  # Check if doc_id exists
+            document = internal_docstore[doc_id]
+            documents.append(document)
+
+    # Rerank the documents 
+    documents = rerank_documents(query, documents)
+    
+    return documents