Update app.py

app.py

@@ -13,108 +13,67 @@ import nltk
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize OpenAI API key
-openai.api_key = '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'  # Replace with your API key
+# Load the ragbench datasets
+ragbench = {}
+for dataset in ['covidqa', 'cuad', 'delucionqa', 'emanual', 'expertqa', 'finqa', 'hagrid', 'hotpotqa', 'msmarco', 'pubmedqa', 'tatqa', 'techqa']:
+    ragbench[dataset] = load_dataset("rungalileo/ragbench", dataset)
+    logger.info(f"Loaded {dataset}")
 
-# Download NLTK data
-nltk.download('punkt')
-
-# Initialize models and configurations
-model_name = 'intfloat/e5-small'
+# Initialize with a stronger model for better semantic understanding
+model_name = 'sentence-transformers/all-mpnet-base-v2'
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 embedding_model = HuggingFaceEmbeddings(model_name=model_name)
 embedding_model.client.to(device)
 
-# Initialize Chroma with existing database
-vectordb = Chroma(
-    persist_directory='./docs/chroma/',
-    embedding_function=embedding_model
-)
-
-def process_query(query):
+def process_query(query, dataset_choice):
     try:
-        logger.info(f"Processing query: {query}")
+        logger.info(f"Processing query for {dataset_choice}: {query}")
         
-        # Get relevant documents
-        relevant_docs = vectordb.similarity_search(query, k=30)
-        context = " ".join([doc.page_content for doc in relevant_docs])
+        # Get relevant documents specific to the chosen dataset
+        relevant_docs = vectordb.max_marginal_relevance_search(
+            query, 
+            k=5,  # Top 5 most relevant documents
+            fetch_k=10  # Fetch top 10 then select most diverse 5
+        )
         
-        # Add delay to respect API rate limits
-        time.sleep(1)
+        context = " ".join([doc.page_content for doc in relevant_docs])
         
-        # Generate response using OpenAI
         response = openai.chat.completions.create(
             model="gpt-4",
             messages=[
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": f"Given the document: {context}\n\nGenerate a response to the query: {query}"}
+                {"role": "system", "content": "You are a specialized assistant for the RagBench dataset. Provide precise answers based solely on the given context."},
+                {"role": "user", "content": f"Dataset: {dataset_choice}\nContext: {context}\nQuestion: {query}\n\nProvide a detailed answer using only the information from the context above."}
             ],
             max_tokens=300,
             temperature=0.7,
         )
         
-        answer = response.choices[0].message.content.strip()
-        logger.info("Successfully generated response")
-        
-        # Extract and display metrics
-        metrics = extract_metrics(query, answer, relevant_docs)
-        
-        return answer, metrics
+        return response.choices[0].message.content.strip()
         
     except Exception as e:
         logger.error(f"Error processing query: {str(e)}")
-        return f"Error: {str(e)}", "Metrics unavailable"
+        return f"Error: {str(e)}"
 
-def extract_metrics(query, response, relevant_docs):
-    try:
-        context = " ".join([doc.page_content for doc in relevant_docs])
-        metrics_prompt = f"""
-        Question: {query}
-        Context: {context}
-        Response: {response}
-        
-        Extract metrics for:
-        - Context Relevance
-        - Context Utilization
-        - Completeness
-        - Response Quality
-        """
-        
-        metrics_response = openai.chat.completions.create(
-            model="gpt-4",
-            messages=[{"role": "user", "content": metrics_prompt}],
-            max_tokens=150,
-            temperature=0.7,
-        )
-        
-        return metrics_response.choices[0].message.content.strip()
-    except Exception as e:
-        return "Metrics calculation failed"
-
-# Create Gradio interface
+# Create Gradio interface with dataset selection
 demo = gr.Interface(
     fn=process_query,
     inputs=[
-        gr.Textbox(
-            label="Enter your question",
-            placeholder="Type your question here...",
-            lines=2
+        gr.Textbox(label="Question", placeholder="Type your question here...", lines=2),
+        gr.Dropdown(
+            choices=list(ragbench.keys()),
+            label="Select Dataset",
+            value="hotpotqa"
         )
     ],
-    outputs=[
-        gr.Textbox(label="Answer", lines=5),
-        gr.Textbox(label="Metrics", lines=4)
-    ],
-    title="RAG-Powered Question Answering System",
-    description="Ask questions and get answers based on the embedded document knowledge.",
+    outputs=gr.Textbox(label="Answer", lines=5),
+    title="RagBench Question Answering System",
+    description="Ask questions across different RagBench datasets",
     examples=[
-        ["What role does T-cell count play in severe human adenovirus type 55 (HAdV-55) infection?"],
-        ["In what school district is Governor John R. Rogers High School located?"],
-        ["Is there a functional neural correlate of individual differences in cardiovascular reactivity?"],
-        ["How do I select Natural mode?"]
+        ["What role does T-cell count play in severe human adenovirus type 55 (HAdV-55) infection?", "covidqa"],
+        ["In what school district is Governor John R. Rogers High School located?", "hotpotqa"],
+        ["Is there a functional neural correlate of individual differences in cardiovascular reactivity?", "pubmedqa"]
     ]
 )
 
-# Launch with debugging enabled
 if __name__ == "__main__":
     demo.launch(debug=True)