Update app.py

app.py

@@ -30,7 +30,7 @@ async def initialize_gemini(file_path, question):
     else:
         return "Error: Unable to process the document. Please ensure the PDF file is valid."
 
-# Mistral Text Completion
+# Improved Mistral Text Completion
 class MistralModel:
     def __init__(self):
         self.model_path = "nvidia/Mistral-NeMo-Minitron-8B-Base"
@@ -39,10 +39,24 @@ class MistralModel:
         self.dtype = torch.bfloat16
         self.model = AutoModelForCausalLM.from_pretrained(self.model_path, torch_dtype=self.dtype, device_map=self.device)
 
-    def generate_text(self, prompt, max_length=50):
-        inputs = self.tokenizer.encode(prompt, return_tensors='pt').to(self.model.device)
-        outputs = self.model.generate(inputs, max_length=max_length)
-        return self.tokenizer.decode(outputs[0])
+    def generate_text(self, prompt, max_length=200):
+        # Improve the prompt for better context
+        enhanced_prompt = f"Question: {prompt}\n\nAnswer: Let's approach this step-by-step:\n1."
+        inputs = self.tokenizer.encode(enhanced_prompt, return_tensors='pt').to(self.model.device)
+        
+        # Generate with more nuanced parameters
+        outputs = self.model.generate(
+            inputs, 
+            max_length=max_length,
+            num_return_sequences=1,
+            no_repeat_ngram_size=3,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7,
+            do_sample=True
+        )
+        
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 mistral_model = MistralModel()
 
@@ -54,14 +68,14 @@ async def process_input(file, question):
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# PDF Question Answering and Text Completion System")
+    gr.Markdown("# Enhanced PDF Question Answering and Text Completion System")
     
-    input_file = gr.File(label="Upload PDF File")
+    input_file = gr.File(label="Upload PDF File (Optional)")
     input_question = gr.Textbox(label="Ask a question or provide a prompt")
     process_button = gr.Button("Process")
     
-    output_text_gemini = gr.Textbox(label="Answer - Gemini")
-    output_text_mistral = gr.Textbox(label="Answer - Mistral")
+    output_text_gemini = gr.Textbox(label="Answer - Gemini (PDF-based if file uploaded)")
+    output_text_mistral = gr.Textbox(label="Answer - Mistral (General knowledge)")
 
     process_button.click(
         fn=process_input,