Update app.py

app.py

@@ -16,10 +16,26 @@ model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
 mistral_tokenizer = AutoTokenizer.from_pretrained(model_path)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 dtype = torch.bfloat16
-mistral_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
+
+# Improved model loading with error handling
+try:
+    mistral_model = AutoModelForCausalLM.from_pretrained(
+        model_path, 
+        torch_dtype=dtype, 
+        device_map=device
+    )
+    print(f"Mistral model loaded successfully on {device}")
+except Exception as e:
+    print(f"Error loading Mistral model: {str(e)}")
+    mistral_model = None
 
 def initialize(file_path, question):
     try:
+        # Check if API key is set
+        api_key = os.getenv("GOOGLE_API_KEY")
+        if not api_key:
+            return "Error: GOOGLE_API_KEY environment variable is not set."
+            
         model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.3)
         prompt_template = """Answer the question as precise as possible using the provided context. If the answer is
                               not contained in the context, say "answer not available in context" \n\n
@@ -30,43 +46,86 @@ def initialize(file_path, question):
         prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
         
         if os.path.exists(file_path):
+            # Load and process PDF
             pdf_loader = PyPDFLoader(file_path)
             pages = pdf_loader.load_and_split()
+            
+            if not pages:
+                return "Error: The PDF file appears to be empty or could not be processed."
+                
             context = "\n".join(str(page.page_content) for page in pages[:30])
+            
+            # Generate Gemini answer
             stuff_chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
-            stuff_answer = stuff_chain({"input_documents": pages, "question": question, "context": context}, return_only_outputs=True)
+            stuff_answer = stuff_chain(
+                {"input_documents": pages, "question": question, "context": context}, 
+                return_only_outputs=True
+            )
             gemini_answer = stuff_answer['output_text']
             
             # Use Mistral model for additional text generation
-            mistral_prompt = f"Based on this answer: {gemini_answer}\nGenerate a follow-up question:"
-            mistral_inputs = mistral_tokenizer.encode(mistral_prompt, return_tensors='pt').to(device)
-            with torch.no_grad():
-                mistral_outputs = mistral_model.generate(mistral_inputs, max_length=50)
-            mistral_output = mistral_tokenizer.decode(mistral_outputs[0], skip_special_tokens=True)
-            
-            combined_output = f"Gemini Answer: {gemini_answer}\n\nMistral Follow-up: {mistral_output}"
+            if mistral_model is not None:
+                mistral_prompt = f"Based on this answer: {gemini_answer}\nGenerate a follow-up question:"
+                mistral_inputs = mistral_tokenizer.encode(mistral_prompt, return_tensors='pt').to(device)
+                
+                with torch.no_grad():
+                    mistral_outputs = mistral_model.generate(
+                        mistral_inputs, 
+                        max_length=200,  # Increased max length
+                        min_length=20,   # Set min length
+                        do_sample=True,  # Enable sampling
+                        top_p=0.95,      # Top-p sampling
+                        temperature=0.7  # Temperature for creativity
+                    )
+                mistral_output = mistral_tokenizer.decode(mistral_outputs[0], skip_special_tokens=True)
+                # Clean up the output to get just the follow-up question
+                if "Generate a follow-up question:" in mistral_output:
+                    mistral_output = mistral_output.split("Generate a follow-up question:")[1].strip()
+                
+                combined_output = f"Gemini Answer: {gemini_answer}\n\nMistral Follow-up: {mistral_output}"
+            else:
+                combined_output = f"Gemini Answer: {gemini_answer}\n\n(Mistral model unavailable)"
+                
             return combined_output
         else:
-            return "Error: Unable to process the document. Please ensure the PDF file is valid."
+            return f"Error: File not found at path '{file_path}'. Please ensure the PDF file is valid."
     except Exception as e:
-        return f"An error occurred: {str(e)}"
-
-# Define Gradio Interface
-input_file = gr.File(label="Upload PDF File")
-input_question = gr.Textbox(label="Ask about the document")
-output_text = gr.Textbox(label="Answer - Combined Gemini and Mistral")
+        import traceback
+        error_details = traceback.format_exc()
+        return f"An error occurred: {str(e)}\n\nDetails: {error_details}"
 
+# Define Gradio Interface with improved error handling
 def pdf_qa(file, question):
     if file is None:
         return "Please upload a PDF file first."
-    return initialize(file.name, question)
+    
+    if not question or question.strip() == "":
+        return "Please enter a question about the document."
+    
+    try:
+        return initialize(file.name, question)
+    except Exception as e:
+        import traceback
+        error_details = traceback.format_exc()
+        return f"Error processing request: {str(e)}\n\nDetails: {error_details}"
 
-# Create Gradio Interface
-gr.Interface(
+# Create Gradio Interface with additional options
+demo = gr.Interface(
     fn=pdf_qa,
-    inputs=[input_file, input_question],
-    outputs=output_text,
+    inputs=[
+        gr.File(label="Upload PDF File", file_types=[".pdf"]),
+        gr.Textbox(label="Ask about the document", placeholder="What is the main topic of this document?")
+    ],
+    outputs=gr.Textbox(label="Answer - Combined Gemini and Mistral"),
     title="RAG Knowledge Retrieval using Gemini API and Mistral Model",
-    description="Upload a PDF file and ask questions about the content."
-).launch()
-#
+    description="Upload a PDF file and ask questions about the content. The system uses Gemini for answering and Mistral for generating follow-up questions.",
+    examples=[
+        [None, "What are the main findings in this document?"],
+        [None, "Summarize the key points discussed in this paper."]
+    ],
+    allow_flagging="never"
+)
+
+# Launch the app with additional parameters
+if __name__ == "__main__":
+    demo.launch(share=True, debug=True)