Update app.py

app.py

@@ -1,135 +1,150 @@
 import gradio as gr
 import torch
-from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
-from qwen_vl_utils import process_vision_info
+import transformers
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 from typing import List, Dict
 import gc
 import os
 
-# Setup logging
-logging.basicConfig(level=logging.INFO)
+# Setup logging with more detail
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
+# Set environment variables for better stability
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+transformers.logging.set_verbosity_info()
+
 class HealthAssistant:
-    def __init__(self):
-        self.model_name = "Qwen/Qwen2-VL-7B-Instruct"
+    def __init__(self, use_smaller_model=True):
+        # Use a smaller model for testing/CPU
+        if use_smaller_model:
+            self.model_name = "facebook/opt-125m"  # Much smaller model for testing
+        else:
+            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):
         try:
-            logger.info("Loading Qwen2-VL model...")
-            # Initialize model with default settings
-            self.model = Qwen2VLForConditionalGeneration.from_pretrained(
+            logger.info(f"Starting model initialization: {self.model_name}")
+            
+            # First try loading tokenizer
+            logger.info("Loading tokenizer...")
+            self.tokenizer = AutoTokenizer.from_pretrained(
                 self.model_name,
-                torch_dtype="auto",
-                device_map="auto",
                 trust_remote_code=True
             )
+            if self.tokenizer is None:
+                raise ValueError("Failed to load tokenizer")
+            logger.info("Tokenizer loaded successfully")
+
+            # Then load the model
+            logger.info("Loading model...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float32,  # Use float32 for CPU
+                low_cpu_mem_usage=True,
+                trust_remote_code=True
+            )
+            if self.model is None:
+                raise ValueError("Failed to load model")
+            
+            # Move model to CPU explicitly
+            self.model = self.model.to("cpu")
+            logger.info("Model loaded successfully and moved to CPU")
             
-            # Initialize processor
-            self.processor = AutoProcessor.from_pretrained(self.model_name)
-            logger.info("Model loaded successfully")
+            # Set padding token if needed
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                logger.info("Set padding token")
+            
+            return True
+
         except Exception as e:
-            logger.error(f"Error initializing model: {e}")
-            raise
+            logger.error(f"Error in model initialization: {str(e)}")
+            raise RuntimeError(f"Model initialization failed: {str(e)}")
+
+    def is_initialized(self):
+        """Check if model is properly initialized"""
+        return (self.model is not None and 
+                self.tokenizer is not None and 
+                hasattr(self.model, 'generate') and 
+                hasattr(self.tokenizer, 'encode'))
 
     def generate_response(self, message: str, history: List = None) -> str:
         try:
-            # Format messages for Qwen2-VL
-            messages = self._format_messages(message, history)
-            
-            # Prepare for inference using qwen_vl_utils
-            text = self.processor.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True
-            )
-            
-            # Process vision info (empty for text-only)
-            image_inputs, video_inputs = process_vision_info(messages)
+            if not self.is_initialized():
+                raise RuntimeError("Model not properly initialized")
+
+            # Prepare prompt
+            prompt = self._prepare_prompt(message, history)
             
-            # Prepare inputs
-            inputs = self.processor(
-                text=[text],
-                images=image_inputs,
-                videos=video_inputs,
+            # Tokenize
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
                 padding=True,
-                return_tensors="pt"
-            )
-            
-            # Move to appropriate device
-            inputs = inputs.to(self.model.device)
+                truncation=True,
+                max_length=512
+            ).to("cpu")  # Ensure CPU tensor
 
-            # Generate response
-            generated_ids = self.model.generate(
-                **inputs,
-                max_new_tokens=128,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9
-            )
+            # Generate
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs["input_ids"],
+                    max_new_tokens=128,
+                    num_beams=1,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
+                )
 
-            # Trim and 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]
+            # Decode
+            response = self.tokenizer.decode(
+                outputs[0][inputs["input_ids"].shape[1]:],
+                skip_special_tokens=True
+            )
 
             # Cleanup
-            del inputs, generated_ids, generated_ids_trimmed
+            del outputs, inputs
             gc.collect()
-            torch.cuda.empty_cache() if torch.cuda.is_available() else None
 
-            return output_text.strip()
+            return response.strip()
 
         except Exception as e:
-            logger.error(f"Error generating response: {e}")
+            logger.error(f"Error generating response: {str(e)}")
             return "I apologize, but I encountered an error. Please try again."
 
-    def _format_messages(self, message: str, history: List = None) -> List[Dict]:
-        """Format messages for Qwen2-VL"""
-        messages = []
-
-        # Add health context as system message
-        health_context = self._get_health_context()
-        if health_context:
-            messages.append({
-                "role": "system",
-                "content": [{"type": "text", "text": f"Health Context:\n{health_context}"}]
-            })
-
-        # Add conversation history
+    def _prepare_prompt(self, message: str, history: List = None) -> str:
+        parts = [
+            "You are a helpful healthcare assistant. Provide accurate and helpful information.",
+            self._get_health_context() or "No health data available yet."
+        ]
+        
         if history:
-            for user_msg, assistant_msg in history[-3:]:
-                messages.extend([
-                    {
-                        "role": "user",
-                        "content": [{"type": "text", "text": user_msg}]
-                    },
-                    {
-                        "role": "assistant",
-                        "content": [{"type": "text", "text": assistant_msg}]
-                    }
+            parts.append("Previous conversation:")
+            for user_msg, bot_msg in history[-3:]:
+                parts.extend([
+                    f"User: {user_msg}",
+                    f"Assistant: {bot_msg}"
                 ])
-
-        # Add current message
-        messages.append({
-            "role": "user",
-            "content": [{"type": "text", "text": message}]
-        })
-
-        return messages
+        
+        parts.extend([
+            f"User: {message}",
+            "Assistant:"
+        ])
+        
+        return "\n\n".join(parts)
 
     def _get_health_context(self) -> str:
         context_parts = []
@@ -180,7 +195,15 @@ class HealthAssistant:
 
 class GradioInterface:
     def __init__(self):
-        self.assistant = HealthAssistant()
+        try:
+            logger.info("Initializing Health Assistant...")
+            self.assistant = HealthAssistant(use_smaller_model=True)  # Use smaller model for testing
+            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}")
+            raise
 
     def chat_response(self, message: str, history: List) -> tuple:
         if not message.strip():
@@ -211,7 +234,6 @@ class GradioInterface:
             gr.Markdown("# 🏥 AI Health Assistant")
             
             with gr.Tabs():
-                # Chat Interface
                 with gr.Tab("💬 Health Chat"):
                     chatbot = gr.Chatbot(
                         height=450,
@@ -227,7 +249,6 @@ class GradioInterface:
                         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)")
@@ -236,7 +257,6 @@ class GradioInterface:
                     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")
@@ -246,7 +266,6 @@ class GradioInterface:
                     med_btn = gr.Button("Add Medication")
                     med_status = gr.Markdown()
 
-            # Event handlers
             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)
@@ -267,8 +286,10 @@ class GradioInterface:
 
 def main():
     try:
+        logger.info("Starting application...")
         interface = GradioInterface()
         demo = interface.create_interface()
+        logger.info("Launching Gradio interface...")
         demo.launch(
             share=False,
             server_name="0.0.0.0",
@@ -277,6 +298,7 @@ def main():
         )
     except Exception as e:
         logger.error(f"Error starting application: {e}")
+        raise
 
 if __name__ == "__main__":
     main()