Update app.py

app.py

@@ -13,76 +13,106 @@ 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['TOKENIZERS_PARALLELISM'] = 'false'
+
 class HealthAssistant:
     def __init__(self):
-        self.model_id = "microsoft/Phi-3-small-128k-instruct"
+        self.model_id = "microsoft/Phi-2"  # Using smaller Phi-2 model instead of Phi-3
         self.model = None
         self.tokenizer = None
         self.pipe = None
         self.metrics = []
         self.medications = []
-        self.device = "cpu"  # Force CPU usage
-        self.initialize_model()
+        self.device = "cpu"
+        self.is_model_loaded = False
+        self.max_history_length = 2
 
     def initialize_model(self):
         try:
+            if self.is_model_loaded:
+                return True
+
             logger.info(f"Loading model: {self.model_id}")
             
-            # Initialize tokenizer
+            # Initialize tokenizer with optimizations
             self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_id,
-                trust_remote_code=True
+                trust_remote_code=True,
+                model_max_length=256,
+                padding_side="left"
             )
             logger.info("Tokenizer loaded")
 
-            # Initialize model
+            # Load model with memory optimizations
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_id,
-                torch_dtype=torch.float32,  # Use float32 for CPU
+                torch_dtype=torch.float32,
                 trust_remote_code=True,
-                device_map=None  # Ensure no device mapping is used
+                device_map=None,
+                low_cpu_mem_usage=True
             ).to(self.device)
-            logger.info("Model loaded on CPU")
 
+            gc.collect()
+            
             # Setup pipeline
             self.pipe = pipeline(
                 "text-generation",
                 model=self.model,
                 tokenizer=self.tokenizer,
-                device=self.device  # Explicitly set to CPU
+                device=self.device,
+                model_kwargs={"low_cpu_mem_usage": True}
             )
-            logger.info("Pipeline created successfully")
             
+            self.is_model_loaded = True
+            logger.info("Model initialized successfully")
             return True
             
         except Exception as e:
             logger.error(f"Error in model initialization: {str(e)}")
             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
+        if hasattr(self, 'pipe') and self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+        if hasattr(self, 'tokenizer') and self.tokenizer is not None:
+            del self.tokenizer
+            self.tokenizer = None
+        self.is_model_loaded = False
+        gc.collect()
+        logger.info("Model unloaded successfully")
+
     def generate_response(self, message: str, history: List = None) -> str:
         try:
-            # Prepare prompt
-            prompt = self._prepare_prompt(message, history)
+            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)
 
-            # Generation configuration
             generation_args = {
-                "max_new_tokens": 500,
+                "max_new_tokens": 200,
                 "return_full_text": False,
                 "temperature": 0.7,
                 "do_sample": True,
                 "top_k": 50,
                 "top_p": 0.9,
-                "repetition_penalty": 1.1
+                "repetition_penalty": 1.1,
+                "num_return_sequences": 1,
+                "batch_size": 1
             }
 
-            # Generate response
-            output = self.pipe(
-                prompt,
-                **generation_args
-            )
+            output = self.pipe(prompt, **generation_args)
             response = output[0]['generated_text']
 
-            # Basic cleanup
             gc.collect()
             
             return response.strip()
@@ -90,22 +120,22 @@ 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:
-        """Prepare prompt with context and history"""
         prompt_parts = [
-            "You are a medical AI assistant providing healthcare information and guidance.",
-            "Always be professional and include appropriate medical disclaimers.",
-            "\nCurrent Health Information:",
-            self._get_health_context(),
-            "\nConversation:"
+            "Medical AI assistant. Be professional, include disclaimers.",
+            self._get_health_context()
         ]
         
         if history:
-            for prev_msg, prev_response in history[-3:]:
+            for prev_msg, prev_response in history:
                 prompt_parts.extend([
-                    f"Human: {prev_msg}",
-                    f"Assistant: {prev_response}"
+                    f"Human: {prev_msg[:100]}",  # Truncate history messages
+                    f"Assistant: {prev_response[:100]}"
                 ])
         
         prompt_parts.extend([
@@ -116,29 +146,26 @@ class HealthAssistant:
         return "\n".join(prompt_parts)
 
     def _get_health_context(self) -> str:
-        context_parts = []
-        
+        if not self.metrics and not self.medications:
+            return "No health data"
+            
+        context = []
         if self.metrics:
             latest = self.metrics[-1]
-            context_parts.extend([
-                "Recent Health Metrics:",
-                f"- Weight: {latest.get('Weight', 'N/A')} kg",
-                f"- Steps: {latest.get('Steps', 'N/A')}",
-                f"- Sleep: {latest.get('Sleep', 'N/A')} hours"
-            ])
-
+            context.append(f"Metrics: W:{latest['Weight']}kg S:{latest['Steps']} Sl:{latest['Sleep']}h")
+        
         if self.medications:
-            context_parts.append("\nCurrent Medications:")
-            for med in self.medications:
-                med_info = f"- {med['Medication']} ({med['Dosage']}) at {med['Time']}"
-                if med.get('Notes'):
-                    med_info += f" | Note: {med['Notes']}"
-                context_parts.append(med_info)
-
-        return "\n".join(context_parts) if context_parts else "No health data recorded"
+            meds = [f"{m['Medication']}({m['Dosage']}@{m['Time']})" for m in self.medications[-2:]]
+            context.append("Meds: " + ", ".join(meds))
+            
+        return " | ".join(context)
 
     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)
+                
             self.metrics.append({
                 'Weight': weight,
                 'Steps': steps,
@@ -151,6 +178,10 @@ 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)
+                
             self.medications.append({
                 'Medication': name,
                 'Dosage': dosage,
@@ -176,9 +207,18 @@ class GradioInterface:
         if not message.strip():
             return "", history
         
-        response = self.assistant.generate_response(message, history)
-        history.append([message, response])
-        return "", history
+        try:
+            response = self.assistant.generate_response(message, history)
+            history.append([message, 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."]]
 
     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]):
@@ -210,9 +250,7 @@ class GradioInterface:
         with gr.Blocks(title="Medical Health Assistant") as demo:
             gr.Markdown("""
             # 🏥 Medical Health Assistant
-            
-            This AI assistant provides medical information and health guidance.
-            **Note**: This is not a replacement for professional medical advice.
+            This AI assistant provides general health information and guidance.
             """)
             
             with gr.Tabs():
@@ -220,13 +258,13 @@ class GradioInterface:
                 with gr.Tab("💬 Medical Consultation"):
                     chatbot = gr.Chatbot(
                         value=[],
-                        height=450,
+                        height=400,
                         show_label=False
                     )
                     with gr.Row():
                         msg = gr.Textbox(
-                            placeholder="Describe your medical concern... (Press Enter)",
-                            lines=2,
+                            placeholder="Ask your health question...",
+                            lines=1,
                             show_label=False,
                             scale=9
                         )
@@ -235,50 +273,48 @@ class GradioInterface:
 
                 # Health Metrics
                 with gr.Tab("📊 Health Metrics"):
+                    gr.Markdown("### Track Your Health Metrics")
                     with gr.Row():
-                        with gr.Column():
-                            gr.Markdown("### Enter Your Health Metrics")
-                            weight_input = gr.Number(
-                                label="Weight (kg)",
-                                minimum=0,
-                                maximum=500
-                            )
-                            steps_input = gr.Number(
-                                label="Steps",
-                                minimum=0,
-                                maximum=100000
-                            )
-                            sleep_input = gr.Number(
-                                label="Hours Slept",
-                                minimum=0,
-                                maximum=24
-                            )
-                            metrics_btn = gr.Button("Save Metrics")
-                            metrics_status = gr.Markdown()
+                        weight_input = gr.Number(
+                            label="Weight (kg)",
+                            minimum=0,
+                            maximum=500
+                        )
+                        steps_input = gr.Number(
+                            label="Steps",
+                            minimum=0,
+                            maximum=100000
+                        )
+                        sleep_input = gr.Number(
+                            label="Hours Slept",
+                            minimum=0,
+                            maximum=24
+                        )
+                    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(
+                        label="Medication Name",
+                        placeholder="Enter medication name"
+                    )
                     with gr.Row():
-                        with gr.Column():
-                            gr.Markdown("### Add Medication Details")
-                            med_name = gr.Textbox(
-                                label="Medication Name",
-                                placeholder="Enter medication name"
-                            )
-                            med_dosage = gr.Textbox(
-                                label="Dosage",
-                                placeholder="e.g., 500mg"
-                            )
-                            med_time = gr.Textbox(
-                                label="Time",
-                                placeholder="e.g., 9:00 AM"
-                            )
-                            med_notes = gr.Textbox(
-                                label="Notes (optional)",
-                                placeholder="Additional instructions or notes"
-                            )
-                            med_btn = gr.Button("Add Medication")
-                            med_status = gr.Markdown()
+                        med_dosage = gr.Textbox(
+                            label="Dosage",
+                            placeholder="e.g., 500mg"
+                        )
+                        med_time = gr.Textbox(
+                            label="Time",
+                            placeholder="e.g., 9:00 AM"
+                        )
+                    med_notes = gr.Textbox(
+                        label="Notes (optional)",
+                        placeholder="Additional instructions or notes"
+                    )
+                    med_btn = gr.Button("Add Medication")
+                    med_status = gr.Markdown()
 
             # Event handlers
             msg.submit(self.chat_response, [msg, chatbot], [msg, chatbot])
@@ -298,27 +334,22 @@ class GradioInterface:
             )
 
             gr.Markdown("""
-            ### ⚠️ Important Medical Disclaimer
-            This AI assistant provides general health information only.
-            - Not a replacement for professional medical advice
-            - Always consult healthcare professionals for medical decisions
-            - Seek immediate medical attention for emergencies
+            ### ⚠️ Medical Disclaimer
+            This AI assistant provides general health information only. Not a replacement for professional medical advice.
+            Always consult healthcare professionals for medical decisions.
             """)
 
-            demo.queue()
+            demo.queue(concurrency_count=1, max_size=5)
             
         return demo
 
 def main():
     try:
-        logger.info("Starting Medical Health Assistant...")
         interface = GradioInterface()
         demo = interface.create_interface()
-        logger.info("Launching interface...")
         demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=False
+            show_error=True,
+            share=True
         )
     except Exception as e:
         logger.error(f"Error starting application: {e}")