Update app.py

app.py

@@ -11,8 +11,8 @@ from langchain.text_splitter import (
 from langchain_community.vectorstores import FAISS, Chroma, Qdrant
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.chains import ConversationalRetrievalChain
-from langchain_community.embeddings import HuggingFaceEmbeddings 
-from langchain_community.llms import HuggingFaceEndpoint
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEndpoint
 from langchain.memory import ConversationBufferMemory
 from sentence_transformers import SentenceTransformer, util
 import torch
@@ -48,53 +48,46 @@ class RAGEvaluator:
         self.test_samples = []
         
     def load_dataset(self, dataset_name: str, num_samples: int = 10):
-        """Load a smaller subset of questions with proper error handling"""
+        """Load dataset with proper error handling"""
         try:
             if dataset_name == "squad":
                 dataset = load_dataset("squad_v2", split="validation")
-                # Select diverse questions
                 samples = dataset.select(range(0, 1000, 100))[:num_samples]
                 
                 self.test_samples = []
                 for sample in samples:
-                    # Check if answers exist and are not empty
-                    if sample.get("answers") and isinstance(sample["answers"], dict) and sample["answers"].get("text"):
+                    # Handle SQuAD format
+                    answers = sample["answers"]
+                    if answers["text"]:  # Check if there are answers
                         self.test_samples.append({
                             "question": sample["question"],
-                            "ground_truth": sample["answers"]["text"][0],
+                            "ground_truth": answers["text"][0],
                             "context": sample["context"]
                         })
                 
             elif dataset_name == "msmarco":
-                dataset = load_dataset("ms_marco", "v2.1", split="dev")
+                dataset = load_dataset("ms_marco", "v2.1", split="test")  # Changed from dev to test
                 samples = dataset.select(range(0, 1000, 100))[:num_samples]
                 
                 self.test_samples = []
                 for sample in samples:
-                    # Check for valid answers
-                    if sample.get("answers") and sample["answers"]:
+                    if sample["answers"]:  # Check if answers exist
                         self.test_samples.append({
                             "question": sample["query"],
                             "ground_truth": sample["answers"][0],
-                            "context": sample["passages"][0]["passage_text"] 
-                                     if isinstance(sample["passages"], list) 
-                                     else sample["passages"]["passage_text"][0]
+                            "context": sample["passages"]["passage_text"][0]
                         })
             
             self.current_dataset = dataset_name
-            
-            # Return dataset info
             return {
                 "dataset": dataset_name,
-                "num_samples": len(self.test_samples),
-                "sample_questions": [s["question"] for s in self.test_samples[:3]],
-                "status": "success"
+                "samples_loaded": len(self.test_samples),
+                "example_questions": [s["question"] for s in self.test_samples[:3]]
             }
             
         except Exception as e:
             print(f"Error loading dataset: {str(e)}")
             return {
-                "dataset": dataset_name,
                 "error": str(e),
                 "status": "failed"
             }
@@ -205,36 +198,58 @@ def create_db(splits, db_choice: str = "faiss"):
     return db_creators[db_choice]()
 
 def initialize_database(list_file_obj, splitting_strategy, chunk_size, db_choice, progress=gr.Progress()):
-    list_file_path = [x.name for x in list_file_obj if x is not None]
-    doc_splits = load_doc(list_file_path, splitting_strategy, chunk_size)
-    vector_db = create_db(doc_splits, db_choice)
-    return vector_db, f"Database created using {splitting_strategy} splitting and {db_choice} vector database!"
+    """Initialize vector database with error handling"""
+    try:
+        if not list_file_obj:
+            return None, "No files uploaded. Please upload PDF documents first."
+            
+        list_file_path = [x.name for x in list_file_obj if x is not None]
+        if not list_file_path:
+            return None, "No valid files found. Please upload PDF documents."
+            
+        doc_splits = load_doc(list_file_path, splitting_strategy, chunk_size)
+        if not doc_splits:
+            return None, "No content extracted from documents."
+            
+        vector_db = create_db(doc_splits, db_choice)
+        return vector_db, f"Database created successfully using {splitting_strategy} splitting and {db_choice} vector database!"
+        
+    except Exception as e:
+        return None, f"Error creating database: {str(e)}"
 
 def initialize_llmchain(llm_choice, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    llm_model = list_llm[llm_choice]
-    
-    llm = HuggingFaceEndpoint(
-        repo_id=llm_model,
-        huggingfacehub_api_token=api_token,
-        temperature=temperature,
-        max_new_tokens=max_tokens,
-        top_k=top_k
-    )
-    
-    memory = ConversationBufferMemory(
-        memory_key="chat_history",
-        output_key='answer',
-        return_messages=True
-    )
+    """Initialize LLM chain with error handling"""
+    try:
+        if vector_db is None:
+            return None, "Please create vector database first."
+            
+        llm_model = list_llm[llm_choice]
+        
+        llm = HuggingFaceEndpoint(
+            repo_id=llm_model,
+            huggingfacehub_api_token=api_token,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
+            top_k=top_k
+        )
+        
+        memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            output_key='answer',
+            return_messages=True
+        )
 
-    retriever = vector_db.as_retriever()
-    qa_chain = ConversationalRetrievalChain.from_llm(
-        llm,
-        retriever=retriever,
-        memory=memory,
-        return_source_documents=True
-    )
-    return qa_chain, "LLM initialized successfully!"
+        retriever = vector_db.as_retriever()
+        qa_chain = ConversationalRetrievalChain.from_llm(
+            llm,
+            retriever=retriever,
+            memory=memory,
+            return_source_documents=True
+        )
+        return qa_chain, "LLM initialized successfully!"
+        
+    except Exception as e:
+        return None, f"Error initializing LLM: {str(e)}"
 
 def conversation(qa_chain, message, history):
     """Fixed conversation function returning all required outputs"""
@@ -424,12 +439,26 @@ def demo():
             initialize_database,
             inputs=[document, splitting_strategy, chunk_size, db_choice],
             outputs=[vector_db, db_progress]
+        ).then(
+            lambda x: gr.update(interactive=True) if x[0] is not None else gr.update(interactive=False),
+            inputs=[vector_db],
+            outputs=[init_llm_btn]
         )
         
         init_llm_btn.click(
             initialize_llmchain,
             inputs=[llm_choice, temperature, max_tokens, top_k, vector_db],
             outputs=[qa_chain, llm_progress]
+        ).then(
+            lambda x: gr.update(interactive=True) if x[0] is not None else gr.update(interactive=False),
+            inputs=[qa_chain],
+            outputs=[msg]
+        )
+    
+        load_dataset_btn.click(
+            lambda x: evaluator.load_dataset(x),
+            inputs=[dataset_choice],
+            outputs=[dataset_info]
         )
         
         msg.submit(