Update app.py

app.py

@@ -13,13 +13,10 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Force CPU usage and set memory optimizations
-torch.set_num_threads(4)
-
 class HealthAssistant:
     def __init__(self, use_smaller_model=True):
         if use_smaller_model:
-            self.model_name = "facebook/opt-125m"
+            self.model_name = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
         else:
             self.model_name = "Qwen/Qwen2-VL-7B-Instruct"
         
@@ -46,11 +43,10 @@ class HealthAssistant:
                 trust_remote_code=True
             )
             
-            self.model = self.model.to("cpu")
-            
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             
+            self.model = self.model.to("cpu")
             logger.info("Model loaded successfully")
             return True
             
@@ -58,19 +54,116 @@ class HealthAssistant:
             logger.error(f"Error in model initialization: {str(e)}")
             raise
 
-    def is_initialized(self):
-        return (self.model is not None and 
-                self.tokenizer is not None and 
-                hasattr(self.model, 'generate'))
+    def _detect_query_type(self, message: str) -> str:
+        """Detect type of medical query"""
+        message_lower = message.lower()
+        
+        emergency_keywords = ["emergency", "severe pain", "chest pain", "can't breathe", 
+                            "unconscious", "stroke", "heart attack"]
+        if any(keyword in message_lower for keyword in emergency_keywords):
+            return "emergency_guidance"
+        
+        symptom_keywords = ["symptom", "feel", "pain", "ache", "suffering", "experiencing"]
+        if any(keyword in message_lower for keyword in symptom_keywords):
+            return "symptom_check"
+        
+        medication_keywords = ["medicine", "drug", "pill", "prescription", "medication", "dose"]
+        if any(keyword in message_lower for keyword in medication_keywords):
+            return "medication_info"
+        
+        lifestyle_keywords = ["exercise", "diet", "sleep", "stress", "healthy", "lifestyle"]
+        if any(keyword in message_lower for keyword in lifestyle_keywords):
+            return "lifestyle_advice"
+        
+        return "general"
+
+    def _get_specialized_prompt(self, query_type: str, message: str) -> str:
+        """Get specialized medical prompts"""
+        base_context = f"""Current Health Context:
+{self._get_health_context()}
+
+"""
+        
+        prompts = {
+            "symptom_check": f"""{base_context}
+Analyze these symptoms professionally:
+Patient's Description: {message}
+
+Provide structured analysis:
+1. Symptoms identified
+2. Possible causes
+3. Severity assessment
+4. Recommended actions
+5. When to seek medical care
+""",
+
+            "medication_info": f"""{base_context}
+Regarding medication inquiry:
+Query: {message}
+
+Provide structured information:
+1. General medication information
+2. Common usage guidelines
+3. Important considerations
+4. General precautions
+5. When to consult healthcare provider
+""",
+
+            "lifestyle_advice": f"""{base_context}
+Health and Lifestyle Query:
+Question: {message}
+
+Provide structured guidance:
+1. Evidence-based recommendations
+2. Implementation steps
+3. Expected benefits
+4. Safety considerations
+5. Progress monitoring tips
+""",
+
+            "emergency_guidance": f"""{base_context}
+URGENT Health Situation:
+Condition: {message}
+
+Critical guidance:
+1. Severity assessment
+2. Immediate actions needed
+3. Emergency signs
+4. When to call emergency services
+5. Precautions while waiting
+⚠️ SEEK IMMEDIATE MEDICAL ATTENTION FOR EMERGENCIES
+""",
+
+            "general": f"""{base_context}
+Medical Query:
+{message}
+
+Provide structured response:
+1. Understanding of the query
+2. Relevant medical information
+3. Professional guidance
+4. Important considerations
+5. Additional recommendations
+"""
+        }
+        
+        return prompts.get(query_type, prompts["general"])
 
     def generate_response(self, message: str, history: List = None) -> str:
         try:
-            if not self.is_initialized():
-                return "System is still initializing. Please try again in a moment."
+            if not hasattr(self, 'model') or self.model is None:
+                return "System is initializing. Please try again in a moment."
 
-            # Prepare prompt
-            prompt = self._prepare_prompt(message, history)
+            # Detect query type and get appropriate prompt
+            query_type = self._detect_query_type(message)
+            prompt = self._get_specialized_prompt(query_type, message)
             
+            # Add conversation history if available
+            if history:
+                prompt += "\n\nPrevious conversation context:"
+                for prev_msg, prev_response in history[-2:]:  # Last 2 exchanges
+                    prompt += f"\nQ: {prev_msg}\nA: {prev_response}\n"
+
             # Tokenize
             inputs = self.tokenizer(
                 prompt,
@@ -84,11 +177,11 @@ class HealthAssistant:
             with torch.no_grad():
                 outputs = self.model.generate(
                     inputs["input_ids"],
-                    max_new_tokens=128,
-                    num_beams=1,
-                    do_sample=True,
-                    temperature=0.7,
+                    max_new_tokens=150,
+                    num_beams=2,
+                    temperature=0.3,
                     top_p=0.9,
+                    no_repeat_ngram_size=3,
                     pad_token_id=self.tokenizer.pad_token_id,
                     eos_token_id=self.tokenizer.eos_token_id
                 )
@@ -99,6 +192,9 @@ class HealthAssistant:
                 skip_special_tokens=True
             )
 
+            # Clean and format response
+            response = self._format_response(response, query_type)
+
             # Cleanup
             del outputs, inputs
             gc.collect()
@@ -107,30 +203,45 @@ 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."
+            return "I apologize, but I encountered an error. Please try rephrasing your question."
 
-    def _prepare_prompt(self, message: str, history: List = None) -> str:
-        parts = [
-            "You are a helpful healthcare assistant providing accurate and helpful medical information.",
-            self._get_health_context() or "No health data available yet."
-        ]
+    def _format_response(self, response: str, query_type: str) -> str:
+        """Format and clean the response"""
+        # Remove repeated headers and clean up
+        lines = [line.strip() for line in response.split('\n') if line.strip()]
+        clean_lines = []
+        seen = set()
+        
+        for line in lines:
+            # Skip common headers and duplicates
+            if any(header in line for header in ["Location:", "Date:", "M.D.", "Medical"]):
+                continue
+            if line not in seen:
+                seen.add(line)
+                clean_lines.append(line)
         
-        if history:
-            parts.append("Previous conversation:")
-            for h in history[-3:]:
-                parts.extend([
-                    f"User: {h[0]}",
-                    f"Assistant: {h[1]}"
-                ])
+        # Add appropriate emoji based on query type
+        emoji_map = {
+            "emergency_guidance": "🚨",
+            "symptom_check": "🔍",
+            "medication_info": "💊",
+            "lifestyle_advice": "💡",
+            "general": "ℹ️"
+        }
         
-        parts.extend([
-            f"User: {message}",
-            "Assistant:"
-        ])
+        emoji = emoji_map.get(query_type, "ℹ️")
         
-        return "\n\n".join(parts)
+        # Combine and format
+        formatted_response = f"{emoji} " + "\n".join(clean_lines)
+        
+        # Add disclaimer if needed
+        if query_type in ["emergency_guidance", "medication_info"]:
+            formatted_response += "\n\n⚠️ This is general information only. Always consult healthcare professionals for medical advice."
+        
+        return formatted_response
 
     def _get_health_context(self) -> str:
+        """Get user's health context"""
         context_parts = []
         
         if self.metrics:
@@ -150,8 +261,7 @@ class HealthAssistant:
                     med_info += f" | Note: {med['Notes']}"
                 context_parts.append(med_info)
 
-        return "\n".join(context_parts) if context_parts else ""
-
+        return "\n".join(context_parts) if context_parts else "No health data available"
     def add_metrics(self, weight: float, steps: int, sleep: float) -> bool:
         try:
             self.metrics.append({
@@ -182,8 +292,6 @@ class GradioInterface:
         try:
             logger.info("Initializing Health Assistant...")
             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")
         except Exception as e:
             logger.error(f"Failed to initialize Health Assistant: {e}")
@@ -193,16 +301,27 @@ class GradioInterface:
         if not message.strip():
             return "", history
         
+        # Generate response
         response = self.assistant.generate_response(message, history)
+        
+        # Update history
         history.append([message, response])
+        
+        # Clear input and return updated history
         return "", history
 
     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]):
             return "⚠️ Please fill in all metrics."
         
+        if weight <= 0 or steps < 0 or sleep < 0:
+            return "⚠️ Please enter valid positive numbers."
+        
         if self.assistant.add_metrics(weight, steps, sleep):
-            return "✅ Health metrics saved successfully!"
+            return f"""✅ Health metrics saved successfully!
+• Weight: {weight} kg
+• Steps: {steps}
+• Sleep: {sleep} hours"""
         return "❌ Error saving metrics."
 
     def add_medication_info(self, name: str, dosage: str, time: str, notes: str) -> str:
@@ -210,48 +329,86 @@ class GradioInterface:
             return "⚠️ Please fill in all required fields."
         
         if self.assistant.add_medication(name, dosage, time, notes):
-            return "✅ Medication added successfully!"
+            return f"""✅ Medication added successfully!
+• Medication: {name}
+• Dosage: {dosage}
+• Time: {time}
+• Notes: {notes if notes else 'None'}"""
         return "❌ Error adding medication."
 
     def create_interface(self):
-        with gr.Blocks(title="Health Assistant") as demo:
-            gr.Markdown("# 🏥 AI Health Assistant")
+        with gr.Blocks(title="Medical Health Assistant", theme=gr.themes.Soft()) 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.
+            """)
             
             with gr.Tabs():
                 # Chat Interface
-                with gr.Tab("💬 Health Chat"):
+                with gr.Tab("💬 Medical Consultation"):
                     chatbot = gr.Chatbot(
                         value=[],
-                        height=450
+                        height=450,
+                        bubble=True
                     )
                     with gr.Row():
                         msg = gr.Textbox(
-                            placeholder="Ask your health question... (Press Enter)",
+                            placeholder="Describe your medical concern... (Press Enter)",
                             lines=2,
                             show_label=False,
                             scale=9
                         )
-                        send_btn = gr.Button("Send", scale=1)
-                    clear_btn = gr.Button("Clear Chat")
+                        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)")
-                        steps_input = gr.Number(label="Steps")
-                        sleep_input = gr.Number(label="Hours Slept")
-                    metrics_btn = gr.Button("Save Metrics")
-                    metrics_status = gr.Markdown()
+                        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()
 
                 # Medication Manager
                 with gr.Tab("💊 Medication Manager"):
                     with gr.Row():
-                        med_name = gr.Textbox(label="Medication Name")
-                        med_dosage = gr.Textbox(label="Dosage")
-                        med_time = gr.Textbox(label="Time (e.g., 9:00 AM)")
-                        med_notes = gr.Textbox(label="Notes (optional)")
-                    med_btn = gr.Button("Add Medication")
-                    med_status = gr.Markdown()
+                        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()
 
             # Event handlers
             msg.submit(self.chat_response, [msg, chatbot], [msg, chatbot])
@@ -270,16 +427,26 @@ class GradioInterface:
                 outputs=[med_status]
             )
 
+            # Add helpful information
+            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
+            """)
+
+            # Enable queuing for better performance
             demo.queue()
             
         return demo
 
 def main():
     try:
-        logger.info("Starting application...")
+        logger.info("Starting Medical Health Assistant...")
         interface = GradioInterface()
         demo = interface.create_interface()
-        logger.info("Launching Gradio interface...")
+        logger.info("Launching interface...")
         demo.launch(
             server_name="0.0.0.0",
             server_port=7860,