Update app.py

app.py

@@ -1,28 +1,25 @@
 import gradio as gr
 import torch
-import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 from typing import List, Dict
 import gc
 import os
 
-# Setup logging with more detail
+# Setup logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-# Set environment variables for better stability
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-transformers.logging.set_verbosity_info()
+# Force CPU usage and set memory optimizations
+torch.set_num_threads(4)
 
 class HealthAssistant:
     def __init__(self, use_smaller_model=True):
-        # Use a smaller model for testing/CPU
         if use_smaller_model:
-            self.model_name = "facebook/opt-125m"  # Much smaller model for testing
+            self.model_name = "facebook/opt-125m"
         else:
             self.model_name = "Qwen/Qwen2-VL-7B-Instruct"
         
@@ -36,53 +33,40 @@ class HealthAssistant:
         try:
             logger.info(f"Starting model initialization: {self.model_name}")
             
-            # First try loading tokenizer
-            logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_name,
                 trust_remote_code=True
             )
-            if self.tokenizer is None:
-                raise ValueError("Failed to load tokenizer")
-            logger.info("Tokenizer loaded successfully")
+            logger.info("Tokenizer loaded")
 
-            # Then load the model
-            logger.info("Loading model...")
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_name,
-                torch_dtype=torch.float32,  # Use float32 for CPU
+                torch_dtype=torch.float32,
                 low_cpu_mem_usage=True,
                 trust_remote_code=True
             )
-            if self.model is None:
-                raise ValueError("Failed to load model")
             
-            # Move model to CPU explicitly
             self.model = self.model.to("cpu")
-            logger.info("Model loaded successfully and moved to CPU")
             
-            # Set padding token if needed
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
-                logger.info("Set padding token")
             
+            logger.info("Model loaded successfully")
             return True
-
+            
         except Exception as e:
             logger.error(f"Error in model initialization: {str(e)}")
-            raise RuntimeError(f"Model initialization failed: {str(e)}")
+            raise
 
     def is_initialized(self):
-        """Check if model is properly initialized"""
         return (self.model is not None and 
                 self.tokenizer is not None and 
-                hasattr(self.model, 'generate') and 
-                hasattr(self.tokenizer, 'encode'))
+                hasattr(self.model, 'generate'))
 
     def generate_response(self, message: str, history: List = None) -> str:
         try:
             if not self.is_initialized():
-                raise RuntimeError("Model not properly initialized")
+                return "System is still initializing. Please try again in a moment."
 
             # Prepare prompt
             prompt = self._prepare_prompt(message, history)
@@ -94,7 +78,7 @@ class HealthAssistant:
                 padding=True,
                 truncation=True,
                 max_length=512
-            ).to("cpu")  # Ensure CPU tensor
+            )
 
             # Generate
             with torch.no_grad():
@@ -127,16 +111,16 @@ class HealthAssistant:
 
     def _prepare_prompt(self, message: str, history: List = None) -> str:
         parts = [
-            "You are a helpful healthcare assistant. Provide accurate and helpful information.",
+            "You are a helpful healthcare assistant providing accurate and helpful medical information.",
             self._get_health_context() or "No health data available yet."
         ]
         
         if history:
             parts.append("Previous conversation:")
-            for user_msg, bot_msg in history[-3:]:
+            for h in history[-3:]:
                 parts.extend([
-                    f"User: {user_msg}",
-                    f"Assistant: {bot_msg}"
+                    f"User: {h[0]}",
+                    f"Assistant: {h[1]}"
                 ])
         
         parts.extend([
@@ -197,7 +181,7 @@ class GradioInterface:
     def __init__(self):
         try:
             logger.info("Initializing Health Assistant...")
-            self.assistant = HealthAssistant(use_smaller_model=True)  # Use smaller model for testing
+            self.assistant = HealthAssistant(use_smaller_model=True)
             if not self.assistant.is_initialized():
                 raise RuntimeError("Health Assistant failed to initialize properly")
             logger.info("Health Assistant initialized successfully")
@@ -230,14 +214,15 @@ class GradioInterface:
         return "❌ Error adding medication."
 
     def create_interface(self):
-        with gr.Blocks(title="Health Assistant", theme=gr.themes.Soft()) as demo:
+        with gr.Blocks(title="Health Assistant") as demo:
             gr.Markdown("# 🏥 AI Health Assistant")
             
             with gr.Tabs():
+                # Chat Interface
                 with gr.Tab("💬 Health Chat"):
                     chatbot = gr.Chatbot(
-                        height=450,
-                        show_label=False,
+                        value=[],
+                        height=450
                     )
                     with gr.Row():
                         msg = gr.Textbox(
@@ -249,6 +234,7 @@ class GradioInterface:
                         send_btn = gr.Button("Send", scale=1)
                     clear_btn = gr.Button("Clear Chat")
 
+                # Health Metrics
                 with gr.Tab("📊 Health Metrics"):
                     with gr.Row():
                         weight_input = gr.Number(label="Weight (kg)")
@@ -257,6 +243,7 @@ class GradioInterface:
                     metrics_btn = gr.Button("Save Metrics")
                     metrics_status = gr.Markdown()
 
+                # Medication Manager
                 with gr.Tab("💊 Medication Manager"):
                     with gr.Row():
                         med_name = gr.Textbox(label="Medication Name")
@@ -266,6 +253,7 @@ class GradioInterface:
                     med_btn = gr.Button("Add Medication")
                     med_status = gr.Markdown()
 
+            # Event handlers
             msg.submit(self.chat_response, [msg, chatbot], [msg, chatbot])
             send_btn.click(self.chat_response, [msg, chatbot], [msg, chatbot])
             clear_btn.click(lambda: [], None, chatbot)
@@ -282,6 +270,8 @@ class GradioInterface:
                 outputs=[med_status]
             )
 
+            demo.queue()
+            
         return demo
 
 def main():
@@ -291,10 +281,9 @@ def main():
         demo = interface.create_interface()
         logger.info("Launching Gradio interface...")
         demo.launch(
-            share=False,
             server_name="0.0.0.0",
             server_port=7860,
-            enable_queue=True
+            share=False
         )
     except Exception as e:
         logger.error(f"Error starting application: {e}")