Update app.py

app.py

@@ -6,65 +6,110 @@ from transformers import T5Tokenizer, T5ForConditionalGeneration
 import gc
 from typing import List, Dict
 import os
+import time
+import logging
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Disable gradient computation and set memory efficient settings
+torch.set_grad_enabled(False)
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+# Create cache directory
+os.makedirs("model_cache", exist_ok=True)
 
 class ModelHandler:
     def __init__(self):
-        self.model_name = "google/flan-t5-base"
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {self.device}")
+        self.model_name = "google/flan-t5-small"  # Small model for Spaces
+        self.device = "cpu"
+        self.initialized = False
+        self.load_attempts = 0
+        self.max_attempts = 3
         self.initialize_model()
 
     def initialize_model(self):
-        try:
-            print(f"Loading model: {self.model_name}")
-            self.tokenizer = T5Tokenizer.from_pretrained(self.model_name)
-            self.model = T5ForConditionalGeneration.from_pretrained(
-                self.model_name,
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True
-            )
-            self.model.to(self.device)
-            print("Model loaded successfully")
-            return True
-        except Exception as e:
-            print(f"Error initializing model: {str(e)}")
-            return False
+        while not self.initialized and self.load_attempts < self.max_attempts:
+            try:
+                logger.info(f"Loading model attempt {self.load_attempts + 1}")
+                self.tokenizer = T5Tokenizer.from_pretrained(
+                    self.model_name,
+                    model_max_length=512,
+                    cache_dir="model_cache"
+                )
+                self.model = T5ForConditionalGeneration.from_pretrained(
+                    self.model_name,
+                    low_cpu_mem_usage=True,
+                    cache_dir="model_cache"
+                )
+                self.initialized = True
+                logger.info("Model loaded successfully")
+                return True
+            except Exception as e:
+                logger.error(f"Loading attempt failed: {str(e)}")
+                self.load_attempts += 1
+                time.sleep(1)
+        return False
 
-    def generate_response(self, prompt: str, max_length: int = 512) -> str:
-        try:
-            # Format prompt for T5
-            formatted_prompt = f"Answer the health question: {prompt}"
+    def generate_response(self, prompt: str, max_length: int = 256) -> str:
+        if not self.initialized:
+            return "Model initialization failed. Using basic responses."
             
-            # Generate response
-            input_ids = self.tokenizer(
-                formatted_prompt,
-                return_tensors="pt",
+        try:
+            clean_prompt = prompt.strip()
+            if len(clean_prompt) == 0:
+                return "Please provide a valid question."
+
+            inputs = self.tokenizer(
+                clean_prompt,
+                max_length=512,
                 truncation=True,
-                max_length=512
-            ).input_ids.to(self.device)
-
-            outputs = self.model.generate(
-                input_ids,
-                max_length=max_length,
-                min_length=20,
-                num_beams=2,
-                temperature=0.7,
-                do_sample=True
+                padding=True,
+                return_tensors="pt"
             )
 
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    input_ids=inputs["input_ids"],
+                    max_length=max_length,
+                    min_length=10,
+                    num_beams=1,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_k=50,
+                    top_p=0.95,
+                )
 
-            # Memory cleanup
-            del outputs, input_ids
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            del outputs, inputs
             gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-            return response
-
+            
+            return response if response else "Could not generate a response."
+            
         except Exception as e:
-            print(f"Error in generate_response: {str(e)}")
-            return "I apologize, but I encountered an error processing your request."
+            logger.error(f"Generation error: {str(e)}")
+            return self.get_fallback_response(prompt)
+
+    def get_fallback_response(self, query: str) -> str:
+        responses = {
+            "hello": "Hello! I'm your health assistant.",
+            "help": "I can help with health information and tracking.",
+            "health": "I provide general health information.",
+            "sleep": "Aim for 7-9 hours of sleep daily.",
+            "exercise": "Regular exercise is important for health.",
+            "diet": "Eat a balanced diet with plenty of vegetables.",
+            "medication": "Always follow prescribed medication schedules.",
+            "water": "Stay hydrated by drinking plenty of water daily.",
+            "stress": "Managing stress is important for overall health."
+        }
+        
+        query = query.lower()
+        for key, response in responses.items():
+            if key in query:
+                return response
+        return "I understand you have a health question. Please try rephrasing it simply."
 
 class HealthData:
     def __init__(self):
@@ -78,14 +123,16 @@ class HealthData:
                 **metrics
             })
             return True
-        except:
+        except Exception as e:
+            logger.error(f"Error adding metrics: {str(e)}")
             return False
 
     def add_medication(self, medication: Dict) -> bool:
         try:
             self.medications.append(medication)
             return True
-        except:
+        except Exception as e:
+            logger.error(f"Error adding medication: {str(e)}")
             return False
 
     def get_health_context(self) -> str:
@@ -93,10 +140,12 @@ class HealthData:
         
         if self.metrics:
             latest = self.metrics[-1]
-            context_parts.append(f"Recent Health Metrics (Date: {latest['Date']}):")
-            context_parts.append(f"- Weight: {latest['Weight']} kg")
-            context_parts.append(f"- Steps: {latest['Steps']}")
-            context_parts.append(f"- Sleep: {latest['Sleep']} hours")
+            context_parts.extend([
+                f"Recent Health Metrics (Date: {latest['Date']}):",
+                f"- Weight: {latest['Weight']} kg",
+                f"- Steps: {latest['Steps']}",
+                f"- Sleep: {latest['Sleep']} hours"
+            ])
 
         if self.medications:
             context_parts.append("\nCurrent Medications:")
@@ -116,34 +165,30 @@ class HealthAssistant:
 
     def get_response(self, message: str, history: List = None) -> str:
         try:
+            self.request_count += 1
+            
             # Prepare context
             context = self.data.get_health_context()
             
-            # Format prompt with context and history
-            prompt = "Given the following context:\n"
-            prompt += f"{context}\n\n"
-            
-            if history:
-                prompt += "Previous conversation:\n"
-                for user_msg, bot_msg in history[-3:]:  # Last 3 exchanges
-                    prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n"
-            
-            prompt += f"Current question: {message}"
-            
+            # Format prompt
+            prompt = (
+                f"Context: {context}\n\n"
+                f"Question: {message}\n\n"
+                "Provide a helpful and accurate health-related response."
+            )
+
             # Get response
             response = self.model.generate_response(prompt)
             
-            # Memory management
+            # Periodic cleanup
             if self.request_count % 5 == 0:
                 gc.collect()
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-                    
+            
             return response
             
         except Exception as e:
-            print(f"Error in get_response: {str(e)}")
-            return "I apologize, but I encountered an error. Please try again."
+            logger.error(f"Error in get_response: {str(e)}")
+            return self.model.get_fallback_response(message)
 
 class HealthAssistantUI:
     def __init__(self):
@@ -153,7 +198,7 @@ class HealthAssistantUI:
         if message.strip() == "":
             return "", history
         
-        bot_message = self.assistant.get_response(message, history)
+        bot_message = self.assistant.get_response(message)
         history.append([message, bot_message])
         return "", history
 
@@ -183,8 +228,13 @@ class HealthAssistantUI:
         return "❌ Error adding medication", None
 
     def create_interface(self):
-        with gr.Blocks(title="Virtual Health Assistant", theme=gr.themes.Soft()) as demo:
-            gr.Markdown("# 🏥 Virtual Health Assistant")
+        with gr.Blocks(title="Health Assistant", theme=gr.themes.Soft()) as demo:
+            gr.Markdown(
+                """
+                # 🏥 Health Assistant
+                Your AI-powered health companion. Track metrics, manage medications, and get health information.
+                """
+            )
             
             with gr.Tabs():
                 # Chat Interface
@@ -251,12 +301,38 @@ class HealthAssistantUI:
                 outputs=[med_status, meds_display]
             )
 
+            gr.Markdown(
+                """
+                ### ⚠️ Important Note
+                This is an AI assistant for general health information only.
+                Always consult healthcare professionals for medical advice.
+                """
+            )
+
         return demo
 
+def cleanup():
+    """Cleanup function for memory management"""
+    gc.collect()
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
 def main():
-    ui = HealthAssistantUI()
-    demo = ui.create_interface()
-    demo.launch()
+    try:
+        logger.info("Starting Health Assistant")
+        ui = HealthAssistantUI()
+        demo = ui.create_interface()
+        
+        # Register cleanup
+        demo.load(cleanup)
+        
+        # Launch app
+        demo.launch(
+            share=False,
+            enable_queue=True,
+            max_threads=4
+        )
+    except Exception as e:
+        logger.error(f"Error starting app: {str(e)}")
 
 if __name__ == "__main__":
     main()