Update app.py

app.py

@@ -38,33 +38,38 @@ def initialize_models():
     
     try:
         # Initialize embedding model (CPU to save GPU memory)
-        print("Loading embedding model...")
-        embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
-        print("✅ Embedding model loaded successfully")
+        if embedding_model is None:
+            print("Loading embedding model...")
+            embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
+            print("✅ Embedding model loaded successfully")
         
         # Initialize chatbot model
-        hf_token = os.getenv('HF_TOKEN')
-        if not hf_token:
-            return False, "HF_TOKEN not found in environment"
-        
-        print("Loading Gemma 3n model...")
-        chatbot_model = Gemma3nForConditionalGeneration.from_pretrained(
-            "google/gemma-3n-e4b-it",
-            device_map="auto",
-            torch_dtype=torch.bfloat16,
-            token=hf_token
-        ).eval()
-        
-        chatbot_processor = AutoProcessor.from_pretrained(
-            "google/gemma-3n-e4b-it",
-            token=hf_token
-        )
+        if chatbot_model is None or chatbot_processor is None:
+            hf_token = os.getenv('HF_TOKEN')
+            if not hf_token:
+                return False, "HF_TOKEN not found in environment"
+            
+            print("Loading Gemma 3n model...")
+            chatbot_model = Gemma3nForConditionalGeneration.from_pretrained(
+                "google/gemma-3n-e4b-it",
+                device_map="auto",
+                torch_dtype=torch.bfloat16,
+                token=hf_token
+            ).eval()
+            
+            chatbot_processor = AutoProcessor.from_pretrained(
+                "google/gemma-3n-e4b-it",
+                token=hf_token
+            )
+            
+            print("✅ Gemma 3n model loaded successfully")
         
-        print("✅ Gemma 3n model loaded successfully")
         return True, "All models loaded successfully"
         
     except Exception as e:
         print(f"Error loading models: {e}")
+        import traceback
+        traceback.print_exc()
         return False, f"Error: {str(e)}"
 
 def extract_text_from_pdf(pdf_file):
@@ -169,14 +174,20 @@ def process_pdf(pdf_file, progress=gr.Progress()):
 
 def chat_with_pdf(message, history):
     """Generate response using RAG"""
+    global chatbot_model, chatbot_processor
+    
     if not message.strip():
         return history
     
     if not processed_text:
         return history + [[message, "❌ Please upload and process a PDF first"]]
     
+    # Check if models are loaded
     if chatbot_model is None or chatbot_processor is None:
-        return history + [[message, "❌ Chatbot model not loaded"]]
+        print("Models not loaded, attempting to reload...")
+        success, error_msg = initialize_models()
+        if not success:
+            return history + [[message, f"❌ Failed to load models: {error_msg}"]]
     
     try:
         # Retrieve relevant chunks
@@ -242,6 +253,24 @@ def clear_chat():
     
     return [], "Ready to process a new PDF"
 
+def get_model_status():
+    """Get current model loading status"""
+    global chatbot_model, chatbot_processor, embedding_model
+    
+    statuses = []
+    
+    if embedding_model is not None:
+        statuses.append("✅ Embedding model loaded")
+    else:
+        statuses.append("❌ Embedding model not loaded")
+    
+    if chatbot_model is not None and chatbot_processor is not None:
+        statuses.append("✅ Chatbot model loaded")
+    else:
+        statuses.append("❌ Chatbot model not loaded")
+    
+    return " | ".join(statuses)
+
 # Initialize models on startup
 model_status = "⏳ Initializing models..."
 if RAG_AVAILABLE:
@@ -265,7 +294,18 @@ with gr.Blocks(
     gr.Markdown("### Upload a PDF and ask questions about it using Retrieval-Augmented Generation")
     
     with gr.Row():
-        gr.Markdown(f"**Status:** {model_status}")
+        status_display = gr.Markdown(f"**Status:** {model_status}")
+        
+        # Add refresh button for status
+        refresh_btn = gr.Button("♾️ Refresh Status", size="sm")
+        
+        def update_status():
+            return get_model_status()
+        
+        refresh_btn.click(
+            fn=update_status,
+            outputs=[status_display]
+        )
     
     with gr.Row():
         # Left column - PDF upload