Update app.py

app.py

@@ -1,238 +1,211 @@
 import gradio as gr
-import pandas as pd
-from datetime import datetime
 import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-import gc
-from typing import List, Dict
-import os
-import time
+from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from qwen_vl_utils import process_vision_info
 import logging
+from typing import List, Dict
+import gc
 
 # 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:
+class HealthAssistant:
     def __init__(self):
-        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.model_name = "Qwen/Qwen2-VL-7B-Instruct"
+        self.model = None
+        self.tokenizer = None
+        self.processor = None
+        self.metrics = []
+        self.medications = []
         self.initialize_model()
 
     def initialize_model(self):
-        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 = 256) -> str:
-        if not self.initialized:
-            return "Model initialization failed. Using basic responses."
-            
         try:
-            clean_prompt = prompt.strip()
-            if len(clean_prompt) == 0:
-                return "Please provide a valid question."
+            logger.info("Loading Qwen2-VL model...")
+            self.model = Qwen2VLForConditionalGeneration.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.bfloat16,
+                attn_implementation="flash_attention_2",
+                device_map="auto"
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            self.processor = AutoProcessor.from_pretrained(
+                self.model_name,
+                min_pixels=256*28*28,
+                max_pixels=1280*28*28
+            )
+            logger.info("Model loaded successfully")
+        except Exception as e:
+            logger.error(f"Error loading model: {e}")
+            raise
 
-            inputs = self.tokenizer(
-                clean_prompt,
-                max_length=512,
-                truncation=True,
+    def generate_response(self, message: str, history: List = None) -> str:
+        try:
+            # Format conversation with health context
+            messages = self._format_messages(message, history)
+            
+            # Prepare for inference
+            text = self.processor.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+            
+            # Since we're not using images in this case
+            image_inputs, video_inputs = [], []
+            
+            # Process inputs
+            inputs = self.processor(
+                text=[text],
+                images=image_inputs,
+                videos=video_inputs,
                 padding=True,
                 return_tensors="pt"
             )
+            inputs = inputs.to(self.model.device)
+
+            # Generate response
+            generated_ids = self.model.generate(
+                **inputs,
+                max_new_tokens=256,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9
+            )
 
-            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,
-                )
-
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            del outputs, inputs
+            # Decode response
+            generated_ids_trimmed = [
+                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+            ]
+            output_text = self.processor.batch_decode(
+                generated_ids_trimmed,
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=False
+            )[0]
+
+            # Cleanup
+            del inputs, generated_ids, generated_ids_trimmed
             gc.collect()
-            
-            return response if response else "Could not generate a response."
-            
+            torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+            return output_text.strip()
+
         except Exception as e:
-            logger.error(f"Generation error: {str(e)}")
-            return self.get_fallback_response(prompt)
+            logger.error(f"Error generating response: {e}")
+            return "I apologize, but I encountered an error. Please try again."
 
-    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."
-        }
+    def _format_messages(self, message: str, history: List = None) -> List[Dict]:
+        """Format messages for the Qwen2-VL model"""
+        # Add system context
+        messages = []
         
-        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."
+        # Add health context
+        health_context = self._get_health_context()
+        if health_context:
+            messages.append({
+                "role": "system",
+                "content": [{"type": "text", "text": f"Current health information:\n{health_context}"}]
+            })
 
-class HealthData:
-    def __init__(self):
-        self.metrics = []
-        self.medications = []
+        # Add conversation history
+        if history:
+            for user_msg, assistant_msg in history[-3:]:  # Last 3 exchanges
+                messages.extend([
+                    {"role": "user", "content": [{"type": "text", "text": user_msg}]},
+                    {"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]}
+                ])
 
-    def add_metrics(self, metrics: Dict) -> bool:
-        try:
-            self.metrics.append({
-                'Date': datetime.now().strftime('%Y-%m-%d'),
-                **metrics
-            })
-            return True
-        except Exception as e:
-            logger.error(f"Error adding metrics: {str(e)}")
-            return False
+        # Add current message
+        messages.append({
+            "role": "user",
+            "content": [{"type": "text", "text": message}]
+        })
 
-    def add_medication(self, medication: Dict) -> bool:
-        try:
-            self.medications.append(medication)
-            return True
-        except Exception as e:
-            logger.error(f"Error adding medication: {str(e)}")
-            return False
+        return messages
 
-    def get_health_context(self) -> str:
+    def _get_health_context(self) -> str:
+        """Get health metrics and medications context"""
         context_parts = []
         
         if self.metrics:
             latest = self.metrics[-1]
             context_parts.extend([
-                f"Recent Health Metrics (Date: {latest['Date']}):",
-                f"- Weight: {latest['Weight']} kg",
-                f"- Steps: {latest['Steps']}",
-                f"- Sleep: {latest['Sleep']} hours"
+                "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"
             ])
 
         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['Notes']:
+                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 available."
-
-class HealthAssistant:
-    def __init__(self):
-        self.model = ModelHandler()
-        self.data = HealthData()
-        self.request_count = 0
+        return "\n".join(context_parts) if context_parts else ""
 
-    def get_response(self, message: str, history: List = None) -> str:
+    def add_metrics(self, weight: float, steps: int, sleep: float) -> bool:
         try:
-            self.request_count += 1
-            
-            # Prepare context
-            context = self.data.get_health_context()
-            
-            # Format prompt
-            prompt = (
-                f"Context: {context}\n\n"
-                f"Question: {message}\n\n"
-                "Provide a helpful and accurate health-related response."
-            )
+            self.metrics.append({
+                'Weight': weight,
+                'Steps': steps,
+                'Sleep': sleep
+            })
+            return True
+        except Exception as e:
+            logger.error(f"Error adding metrics: {e}")
+            return False
 
-            # Get response
-            response = self.model.generate_response(prompt)
-            
-            # Periodic cleanup
-            if self.request_count % 5 == 0:
-                gc.collect()
-            
-            return response
-            
+    def add_medication(self, name: str, dosage: str, time: str, notes: str = "") -> bool:
+        try:
+            self.medications.append({
+                'Medication': name,
+                'Dosage': dosage,
+                'Time': time,
+                'Notes': notes
+            })
+            return True
         except Exception as e:
-            logger.error(f"Error in get_response: {str(e)}")
-            return self.model.get_fallback_response(message)
+            logger.error(f"Error adding medication: {e}")
+            return False
 
-class HealthAssistantUI:
+class GradioInterface:
     def __init__(self):
         self.assistant = HealthAssistant()
 
-    def user_chat(self, message: str, history: List) -> tuple:
-        if message.strip() == "":
+    def chat_response(self, message: str, history: List) -> tuple:
+        if not message.strip():
             return "", history
         
-        bot_message = self.assistant.get_response(message)
-        history.append([message, bot_message])
+        response = self.assistant.generate_response(message, history)
+        history.append([message, response])
         return "", history
 
-    def save_metrics(self, weight: float, steps: int, sleep: float) -> tuple:
-        if not all([weight is not None, steps is not None, sleep is not None]):
-            return "⚠️ Please fill in all metrics.", None
+    def add_health_metrics(self, weight: float, steps: int, sleep: float) -> str:
+        if not all([weight, steps, sleep]):
+            return "⚠️ Please fill in all metrics."
         
-        metrics = {'Weight': weight, 'Steps': steps, 'Sleep': sleep}
-        if self.assistant.data.add_metrics(metrics):
-            df = pd.DataFrame(self.assistant.data.metrics)
-            return "✅ Metrics saved successfully!", df
-        return "❌ Error saving metrics", None
+        if self.assistant.add_metrics(weight, steps, sleep):
+            return "✅ Health metrics saved successfully!"
+        return "❌ Error saving metrics."
 
-    def save_medication(self, name: str, dosage: str, time: str, notes: str) -> tuple:
+    def add_medication_info(self, name: str, dosage: str, time: str, notes: str) -> str:
         if not all([name, dosage, time]):
-            return "⚠️ Please fill in all required fields.", None
+            return "⚠️ Please fill in all required fields."
         
-        medication = {
-            'Medication': name,
-            'Dosage': dosage,
-            'Time': time,
-            'Notes': notes or ''
-        }
-        if self.assistant.data.add_medication(medication):
-            df = pd.DataFrame(self.assistant.data.medications)
-            return "✅ Medication added successfully!", df
-        return "❌ Error adding medication", None
+        if self.assistant.add_medication(name, dosage, time, notes):
+            return "✅ Medication added successfully!"
+        return "❌ Error adding medication."
 
     def create_interface(self):
         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.
+                # 🏥 AI Health Assistant
+                Powered by Qwen2-VL for intelligent health guidance and monitoring.
                 """
             )
             
@@ -241,11 +214,11 @@ class HealthAssistantUI:
                 with gr.Tab("💬 Health Chat"):
                     chatbot = gr.Chatbot(
                         height=450,
-                        show_label=False,
+                        show_label=False
                     )
                     with gr.Row():
                         msg = gr.Textbox(
-                            placeholder="Type your health question... (Press Enter)",
+                            placeholder="Ask your health question... (Press Enter)",
                             lines=2,
                             show_label=False,
                             scale=9
@@ -253,86 +226,63 @@ class HealthAssistantUI:
                         send_btn = gr.Button("Send", scale=1)
                     clear_btn = gr.Button("Clear Chat")
 
-                # Health Metrics Tab
+                # Health Metrics
                 with gr.Tab("📊 Health Metrics"):
                     with gr.Row():
-                        with gr.Column():
-                            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():
-                            metrics_display = gr.Dataframe(
-                                headers=["Date", "Weight", "Steps", "Sleep"]
-                            )
+                        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()
 
-                # Medication Manager Tab
+                # Medication Manager
                 with gr.Tab("💊 Medication Manager"):
                     with gr.Row():
-                        with gr.Column():
-                            med_name = gr.Textbox(label="Medication Name")
-                            med_dosage = gr.Textbox(label="Dosage")
-                            med_time = gr.Textbox(label="Time")
-                            med_notes = gr.Textbox(label="Notes (optional)")
-                            med_btn = gr.Button("Add Medication")
-                            med_status = gr.Markdown()
-                        
-                        with gr.Column():
-                            meds_display = gr.Dataframe(
-                                headers=["Medication", "Dosage", "Time", "Notes"]
-                            )
+                        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()
 
             # Event handlers
-            msg.submit(self.user_chat, [msg, chatbot], [msg, chatbot])
-            send_btn.click(self.user_chat, [msg, chatbot], [msg, chatbot])
+            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)
             
             metrics_btn.click(
-                self.save_metrics,
+                self.add_health_metrics,
                 inputs=[weight_input, steps_input, sleep_input],
-                outputs=[metrics_status, metrics_display]
+                outputs=[metrics_status]
             )
             
             med_btn.click(
-                self.save_medication,
+                self.add_medication_info,
                 inputs=[med_name, med_dosage, med_time, med_notes],
-                outputs=[med_status, meds_display]
+                outputs=[med_status]
             )
 
             gr.Markdown(
                 """
                 ### ⚠️ Important Note
-                This is an AI assistant for general health information only.
+                This AI assistant provides 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():
     try:
-        logger.info("Starting Health Assistant")
-        ui = HealthAssistantUI()
-        demo = ui.create_interface()
-        
-        # Register cleanup
-        demo.load(cleanup)
-        
-        # Launch app
+        interface = GradioInterface()
+        demo = interface.create_interface()
         demo.launch(
             share=False,
             enable_queue=True,
             max_threads=4
         )
     except Exception as e:
-        logger.error(f"Error starting app: {str(e)}")
+        logger.error(f"Error starting application: {e}")
 
 if __name__ == "__main__":
     main()