Update app.py

app.py

@@ -14,12 +14,12 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # Set environment variables for memory optimization
-os.environ['TRANSFORMERS_CACHE'] = '/home/user/.cache/huggingface/hub'  # HF Spaces cache directory
+os.environ['TRANSFORMERS_CACHE'] = '/home/user/.cache/huggingface/hub'
 os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 
 class HealthAssistant:
     def __init__(self):
-        self.model_id = "microsoft/Phi-2"  # Using smaller Phi-2 model instead of Phi-3
+        self.model_id = "microsoft/Phi-2"  # Using smaller Phi-2 model
         self.model = None
         self.tokenizer = None
         self.pipe = None
@@ -36,7 +36,6 @@ class HealthAssistant:
 
             logger.info(f"Loading model: {self.model_id}")
             
-            # Initialize tokenizer with optimizations
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_id,
                 trust_remote_code=True,
@@ -45,7 +44,6 @@ class HealthAssistant:
             )
             logger.info("Tokenizer loaded")
 
-            # Load model with memory optimizations
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_id,
                 torch_dtype=torch.float32,
@@ -56,7 +54,6 @@ class HealthAssistant:
 
             gc.collect()
             
-            # Setup pipeline
             self.pipe = pipeline(
                 "text-generation",
                 model=self.model,
@@ -74,7 +71,6 @@ class HealthAssistant:
             raise
 
     def unload_model(self):
-        """Unload model to free up memory"""
         if hasattr(self, 'model') and self.model is not None:
             del self.model
             self.model = None
@@ -93,7 +89,6 @@ class HealthAssistant:
             if not self.is_model_loaded:
                 self.initialize_model()
             
-            # Limit message length
             message = message[:200]  # Truncate long messages
             
             prompt = self._prepare_prompt(message, history[-self.max_history_length:] if history else None)
@@ -120,10 +115,6 @@ class HealthAssistant:
         except Exception as e:
             logger.error(f"Error generating response: {str(e)}")
             return "I apologize, but I encountered an error. Please try again."
-        finally:
-            # Attempt to free memory after each generation
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
 
     def _prepare_prompt(self, message: str, history: List = None) -> str:
         prompt_parts = [
@@ -132,11 +123,17 @@ class HealthAssistant:
         ]
         
         if history:
-            for prev_msg, prev_response in history:
-                prompt_parts.extend([
-                    f"Human: {prev_msg[:100]}",  # Truncate history messages
-                    f"Assistant: {prev_response[:100]}"
-                ])
+            for h in history:
+                if isinstance(h, dict):  # New message format
+                    if h['role'] == 'user':
+                        prompt_parts.append(f"Human: {h['content'][:100]}")
+                    else:
+                        prompt_parts.append(f"Assistant: {h['content'][:100]}")
+                else:  # Old format (tuple)
+                    prompt_parts.extend([
+                        f"Human: {h[0][:100]}",
+                        f"Assistant: {h[1][:100]}"
+                    ])
         
         prompt_parts.extend([
             f"Human: {message}",
@@ -162,7 +159,6 @@ class HealthAssistant:
 
     def add_metrics(self, weight: float, steps: int, sleep: float) -> bool:
         try:
-            # Keep only last 5 metrics
             if len(self.metrics) >= 5:
                 self.metrics.pop(0)
                 
@@ -178,7 +174,6 @@ class HealthAssistant:
 
     def add_medication(self, name: str, dosage: str, time: str, notes: str = "") -> bool:
         try:
-            # Keep only last 5 medications
             if len(self.medications) >= 5:
                 self.medications.pop(0)
                 
@@ -209,16 +204,20 @@ class GradioInterface:
         
         try:
             response = self.assistant.generate_response(message, history)
-            history.append([message, response])
+            # Convert to new message format
+            history.append({"role": "user", "content": message})
+            history.append({"role": "assistant", "content": response})
             
-            # Unload model periodically
             if len(history) % 3 == 0:
                 self.assistant.unload_model()
                 
             return "", history
         except Exception as e:
             logger.error(f"Error in chat response: {e}")
-            return "", history + [[message, "I apologize, but I encountered an error. Please try again."]]
+            return "", history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": "I apologize, but I encountered an error. Please try again."}
+            ]
 
     def add_health_metrics(self, weight: float, steps: int, sleep: float) -> str:
         if not all([weight is not None, steps is not None, sleep is not None]):
@@ -254,24 +253,23 @@ class GradioInterface:
             """)
             
             with gr.Tabs():
-                # Chat Interface
                 with gr.Tab("💬 Medical Consultation"):
                     chatbot = gr.Chatbot(
                         value=[],
                         height=400,
-                        show_label=False
+                        label=False,
+                        type="messages"  # Using new message format
                     )
                     with gr.Row():
                         msg = gr.Textbox(
                             placeholder="Ask your health question...",
                             lines=1,
-                            show_label=False,
+                            label=False,
                             scale=9
                         )
                         send_btn = gr.Button("Send", scale=1)
                     clear_btn = gr.Button("Clear Chat")
 
-                # Health Metrics
                 with gr.Tab("📊 Health Metrics"):
                     gr.Markdown("### Track Your Health Metrics")
                     with gr.Row():
@@ -293,7 +291,6 @@ class GradioInterface:
                     metrics_btn = gr.Button("Save Metrics")
                     metrics_status = gr.Markdown()
 
-                # Medication Manager
                 with gr.Tab("💊 Medication Manager"):
                     gr.Markdown("### Track Your Medications")
                     med_name = gr.Textbox(
@@ -316,7 +313,6 @@ 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)
@@ -339,7 +335,7 @@ class GradioInterface:
             Always consult healthcare professionals for medical decisions.
             """)
 
-            demo.queue(concurrency_count=1, max_size=5)
+            demo.queue(max_size=5)
             
         return demo
 
@@ -348,6 +344,7 @@ def main():
         interface = GradioInterface()
         demo = interface.create_interface()
         demo.launch(
+            server_name="0.0.0.0",
             show_error=True,
             share=True
         )