Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import logging
 from typing import List, Dict
 import gc
@@ -13,232 +13,113 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Set torch threads
-torch.set_num_threads(4)
+# Set random seed for reproducibility
+torch.random.manual_seed(0)
 
 class HealthAssistant:
-    def __init__(self, use_smaller_model=True):
-        if use_smaller_model:
-            self.model_name = "Qwen/Qwen2-VL-7B-Instruct"
-        else:
-            self.model_name = "Qwen/Qwen2-VL-7B-Instruct"
-        
+    def __init__(self):
+        self.model_id = "microsoft/Phi-3-small-128k-instruct"
         self.model = None
         self.tokenizer = None
+        self.pipe = None
         self.metrics = []
         self.medications = []
         self.initialize_model()
 
     def initialize_model(self):
         try:
-            logger.info(f"Starting model initialization: {self.model_name}")
+            logger.info(f"Loading model: {self.model_id}")
             
+            # Initialize tokenizer
             self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name,
+                self.model_id,
                 trust_remote_code=True
             )
             logger.info("Tokenizer loaded")
 
+            # Initialize model
             self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True,
+                self.model_id,
+                torch_dtype="auto",
                 trust_remote_code=True
             )
+
+            # Set device
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.model = self.model.to(self.device)
+            logger.info(f"Model loaded on {self.device}")
+
+            # Setup pipeline
+            self.pipe = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                device=self.device
+            )
+            logger.info("Pipeline created successfully")
             
-            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
             
         except Exception as e:
             logger.error(f"Error in model initialization: {str(e)}")
             raise
 
-    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"
+    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:"
+        ]
         
-        symptom_keywords = ["symptom", "feel", "pain", "ache", "suffering", "experiencing"]
-        if any(keyword in message_lower for keyword in symptom_keywords):
-            return "symptom_check"
+        if history:
+            for prev_msg, prev_response in history[-3:]:
+                prompt_parts.extend([
+                    f"Human: {prev_msg}",
+                    f"Assistant: {prev_response}"
+                ])
         
-        medication_keywords = ["medicine", "drug", "pill", "prescription", "medication", "dose"]
-        if any(keyword in message_lower for keyword in medication_keywords):
-            return "medication_info"
+        prompt_parts.extend([
+            f"Human: {message}",
+            "Assistant:"
+        ])
         
-        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 _prepare_medical_prompt(self, message: str, query_type: str) -> str:
-        """Prepare medical prompt based on query type"""
-        base_context = self._get_health_context()
-        
-        prompts = {
-            "symptom_check": f"""You are a medical AI assistant. Based on the following health context and symptoms, provide a careful analysis.
-
-Current Health Context:
-{base_context}
-
-Patient's Symptoms: {message}
-
-Provide a structured response covering:
-1. Key symptoms identified
-2. Possible common causes
-3. General recommendations
-4. Warning signs to watch for
-5. When to seek medical care
-
-Remember to maintain a professional and careful tone.""",
-
-            "medication_info": f"""You are a medical AI assistant. Provide information about the medication inquiry while noting you cannot give prescription advice.
-
-Current Health Context:
-{base_context}
-
-Medication Query: {message}
-
-Provide general information about:
-1. Basic medication category/purpose
-2. General usage patterns
-3. Common considerations
-4. Important precautions
-5. When to consult a healthcare provider
-
-Remember to emphasize this is general information only.""",
-
-            "emergency_guidance": f"""You are a medical AI assistant. This appears to be an urgent situation.
-
-Current Health Context:
-{base_context}
-
-Urgent Situation: {message}
-
-Provide immediate guidance:
-1. Severity assessment
-2. Immediate actions needed
-3. Emergency warning signs
-4. Whether to call emergency services
-5. Precautions while waiting
-
-Always emphasize seeking immediate medical care for emergencies.""",
-
-            "general": f"""You are a medical AI assistant. Provide helpful health information based on the query.
-
-Current Health Context:
-{base_context}
-
-Health Query: {message}
-
-Provide a structured response covering:
-1. Understanding of the question
-2. Relevant health information
-3. General guidance
-4. Important considerations
-5. Additional recommendations"""
-        }
-
-        return prompts.get(query_type, prompts["general"])
+        return "\n".join(prompt_parts)
 
     def generate_response(self, message: str, history: List = None) -> str:
         try:
-            if not hasattr(self, 'model') or self.model is None:
-                return "System is initializing. Please try again in a moment."
-
-            # Detect query type
-            query_type = self._detect_query_type(message)
-            
             # Prepare prompt
-            prompt = self._prepare_medical_prompt(message, query_type)
-
-            # Add conversation history if available
-            if history:
-                prompt += "\n\nRecent conversation context:"
-                for prev_msg, prev_response in history[-2:]:
-                    prompt += f"\nQ: {prev_msg}\nA: {prev_response}\n"
-
-            # Tokenize
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=512
-            )
-
-            # Generate
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    inputs["input_ids"],
-                    max_new_tokens=150,
-                    num_beams=1,
-                    temperature=0.7,
-                    top_p=0.9,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
-                )
-
-            # Decode
-            response = self.tokenizer.decode(
-                outputs[0][inputs["input_ids"].shape[1]:],
-                skip_special_tokens=True
-            )
-
-            # Format response
-            response = self._format_response(response, query_type)
+            prompt = self._prepare_prompt(message, history)
+
+            # Generation configuration
+            generation_args = {
+                "max_new_tokens": 500,
+                "return_full_text": False,
+                "temperature": 0.7,
+                "do_sample": True,
+                "top_k": 50,
+                "top_p": 0.9,
+                "repetition_penalty": 1.1
+            }
+
+            # Generate response
+            output = self.pipe(prompt, **generation_args)
+            response = output[0]['generated_text']
 
             # Cleanup
-            del outputs, inputs
             gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
 
             return response.strip()
 
         except Exception as e:
             logger.error(f"Error generating response: {str(e)}")
-            return "I apologize, but I encountered an error. Please try rephrasing your question."
-
-    def _format_response(self, response: str, query_type: str) -> str:
-        """Format and clean the response"""
-        # Remove repeated headers
-        lines = [line.strip() for line in response.split('\n') if line.strip()]
-        clean_lines = []
-        seen = set()
-        
-        for line in lines:
-            if line not in seen:
-                seen.add(line)
-                clean_lines.append(line)
-        
-        # Add appropriate prefix based on query type
-        prefixes = {
-            "emergency_guidance": "🚨 URGENT: ",
-            "symptom_check": "🔍 Analysis: ",
-            "medication_info": "💊 Medication Info: ",
-            "lifestyle_advice": "💡 Health Advice: ",
-            "general": "ℹ️ "
-        }
-        
-        prefix = prefixes.get(query_type, "ℹ️ ")
-        formatted_response = prefix + "\n".join(clean_lines)
-        
-        # Add disclaimer for certain types
-        if query_type in ["emergency_guidance", "medication_info"]:
-            formatted_response += "\n\n⚠️ Note: This is general information only. Always consult healthcare professionals."
-        
-        return formatted_response
+            return "I apologize, but I encountered an error. Please try again."
 
     def _get_health_context(self) -> str:
-        """Get user's health context"""
         context_parts = []
         
         if self.metrics:
@@ -289,7 +170,7 @@ class GradioInterface:
     def __init__(self):
         try:
             logger.info("Initializing Health Assistant...")
-            self.assistant = HealthAssistant(use_smaller_model=True)
+            self.assistant = HealthAssistant()
             logger.info("Health Assistant initialized successfully")
         except Exception as e:
             logger.error(f"Failed to initialize Health Assistant: {e}")