Update app.py

app.py

@@ -42,10 +42,12 @@ from langchain.vectorstores import FAISS
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.document_loaders import TextLoader
-from google.colab import output
 import IPython.display as display
 from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
-
+import warnings
+import re
+from dotenv import load_dotenv
+warnings.filterwarnings('ignore')
 
 # Ensure Hugging Face login
 try:
@@ -67,34 +69,85 @@ os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+class ModelManager:
+    """Handles model loading and resource management"""
+    
+    @staticmethod
+    def verify_model_path(checkpoint_path: str) -> str:
+        """Verify and return valid model path"""
+        if os.path.exists(checkpoint_path):
+            return checkpoint_path
+            
+        alternate_paths = [
+            f"{os.getcwd()}/checkpoint-500.zip",
+            "./checkpoint-500.zip",
+            "../checkpoint-500.zip"
+        ]
+        
+        for path in alternate_paths:
+            if os.path.exists(path):
+                return path
+                
+        raise FileNotFoundError(
+            f"Model checkpoint not found in any of these locations: "
+            f"{[checkpoint_path] + alternate_paths}"
+        )
+    
+    @staticmethod
+    def clear_gpu_memory():
+        """Clear GPU memory cache"""
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            gc.collect()
+
 class PearlyBot:
-    def __init__(self):
-        self.setup_model()
+    def __init__(self, model_path: str = "./checkpoint-500.zip"):
+        self.setup_model(model_path)
         self.setup_rag()
         self.conversation_history = []
+        self.last_interaction_time = time.time()
+        self.interaction_cooldown = 1.0  # seconds
         
-    def setup_model(self):
+    def setup_model(self, model_path: str):
+        """Initialize the model with proper error handling"""
         try:
-            logger.info("Loading local checkpoint...")
-            # Load from the local checkpoint in your Space
-            checkpoint_path = "checkpoint-500.zip"  # Path to your uploaded checkpoint
-            base_model_id = "google/gemma-2b"  # Your base model
+            logger.info("Starting model initialization...")
+            ModelManager.clear_gpu_memory()
             
-            # Load tokenizer from base model
-            self.tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+            # Verify model path
+            verified_path = ModelManager.verify_model_path(model_path)
+            logger.info(f"Using model checkpoint from: {verified_path}")
             
-            # Load model with checkpoint
-            self.model = AutoModelForCausalLM.from_pretrained(
-                checkpoint_path,
-                device_map="auto",
-                load_in_8bit=True,
-                torch_dtype=torch.float16,
-                low_cpu_mem_usage=True
-            )
-            self.model.eval()
-            logger.info("Model loaded successfully")
+            # Base model configuration
+            base_model_id = "google/gemma-2b"
+            logger.info(f"Loading base model: {base_model_id}")
+            
+            # Load tokenizer
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                logger.info("Tokenizer loaded successfully")
+            except Exception as e:
+                logger.error(f"Failed to load tokenizer: {str(e)}")
+                raise
+            
+            # Load model
+            try:
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    verified_path,
+                    device_map="auto",
+                    load_in_8bit=True,
+                    torch_dtype=torch.float16,
+                    low_cpu_mem_usage=True
+                )
+                self.model.eval()
+                logger.info("Model loaded successfully")
+            except Exception as e:
+                logger.error(f"Failed to load model: {str(e)}")
+                raise
+                
         except Exception as e:
-            logger.error(f"Error loading model: {str(e)}")
+            logger.error(f"Error in model setup: {str(e)}")
             raise
 
     def setup_rag(self):
@@ -413,7 +466,16 @@ PROFESSIONAL CONDUCT:
 
     @torch.inference_mode()
     def generate_response(self, message: str, history: list) -> str:
+        """Generate response using both fine-tuned model and RAG"""
         try:
+            # Rate limiting
+            current_time = time.time()
+            if current_time - self.last_interaction_time < self.interaction_cooldown:
+                time.sleep(self.interaction_cooldown)
+            
+            # Clear GPU memory before generation
+            ModelManager.clear_gpu_memory()
+            
             # Get RAG context
             context = self.get_relevant_context(message)
             
@@ -445,107 +507,143 @@ Guidelines:
 <start_of_turn>assistant"""
 
             # Generate response
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                truncation=True,
-                max_length=512
-            ).to(self.model.device)
-
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=256,
-                min_new_tokens=20,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
-                repetition_penalty=1.2
-            )
-            
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            response = response.split("<start_of_turn>assistant")[-1].strip()
-            if "<end_of_turn>" in response:
-                response = response.split("<end_of_turn>")[0].strip()
+            try:
+                inputs = self.tokenizer(
+                    prompt,
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=512
+                ).to(self.model.device)
+
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=256,
+                    min_new_tokens=20,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9,
+                    repetition_penalty=1.2,
+                    no_repeat_ngram_size=3
+                )
+                
+                response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+                response = response.split("<start_of_turn>assistant")[-1].strip()
+                if "<end_of_turn>" in response:
+                    response = response.split("<end_of_turn>")[0].strip()
+                
+                self.last_interaction_time = time.time()
+                return response
+                
+            except torch.cuda.OutOfMemoryError:
+                ModelManager.clear_gpu_memory()
+                logger.error("GPU out of memory, cleared cache and retrying...")
+                return "I apologize, but I'm experiencing technical difficulties. Please try again."
                 
-            return response
-            
         except Exception as e:
             logger.error(f"Error generating response: {str(e)}")
             return "I apologize, but I encountered an error. Please try again."
+            
+    def handle_feedback(self, message: str, response: str, feedback: int):
+        """Handle user feedback for responses"""
+        try:
+            timestamp = datetime.now().isoformat()
+            feedback_data = {
+                "message": message,
+                "response": response,
+                "feedback": feedback,
+                "timestamp": timestamp
+            }
+            
+            # Log feedback
+            logger.info(f"Feedback received: {feedback_data}")
+            
+            # Here you could:
+            # 1. Store feedback in a database
+            # 2. Send to monitoring system
+            # 3. Use for model improvements
+            
+            return True
+        except Exception as e:
+            logger.error(f"Error handling feedback: {e}")
+            return False
+
+    def __del__(self):
+        """Cleanup resources"""
+        try:
+            if hasattr(self, 'model'):
+                del self.model
+            ModelManager.clear_gpu_memory()
+        except Exception as e:
+            logger.error(f"Error in cleanup: {e}")
 
+    def process_feedback(positive: bool, comment: str, history: List[Dict[str, str]]):
+        try:
+            if not history or len(history) < 2:
+                return gr.update(value="")
+                
+            last_user_msg = history[-2]["content"] if isinstance(history[-2], dict) else history[-2][0]
+            last_bot_msg = history[-1]["content"] if isinstance(history[-1], dict) else history[-1][1]
+            
+            bot.handle_feedback(
+                message=last_user_msg,
+                response=last_bot_msg,
+                feedback=1 if positive else -1
+            )
+            
+            return gr.update(value="")
+            
+        except Exception as e:
+            logger.error(f"Error processing feedback: {e}")
+            return gr.update(value="")
 
 def create_demo():
     """Set up Gradio interface for the chatbot with enhanced styling and functionality."""
     try:
-        # Health check
-        @gr.routes.get("/health")
-        def health_check():
-            return {"status": "healthy"}
-        bot = PearlyBot()  # ✅
+        # Initialize bot
+        bot = PearlyBot()
 
-        def chat(message: str, history: list):  # ✅
+        def chat(message: str, history: list):
+            """Handle chat interactions"""
             try:
                 if not message.strip():
                     return history
-                    
-                bot_response = bot.generate_response(message)
                 
-                # Add user message
+                # Generate response
+                response = bot.generate_response(message, history)
+                
+                # Update history with proper formatting
                 history.append({
                     "role": "user",
                     "content": message
                 })
-                
-                # Add bot response
                 history.append({
                     "role": "assistant",
-                    "content": bot_response['response']
+                    "content": response
                 })
-                
                 return history
                 
             except Exception as e:
                 logger.error(f"Chat error: {e}")
-                history.append({
-                    "role": "user",
-                    "content": message
-                })
-                history.append({
+                return history + [{
                     "role": "assistant",
                     "content": "I apologize, but I'm experiencing technical difficulties. For emergencies, please call 999."
-                })
-                return history
-
-        def process_feedback(is_positive: bool, comment: str, history: List[Dict[str, str]]):
-            try:
-                if not history:
-                    return
-                    
-                last_interaction = history[-2:]  # Get last user message and bot response
-                if len(last_interaction) == 2:
-                    user_msg = last_interaction[0]["content"]
-                    bot_msg = last_interaction[1]["content"]
-                    
-                    feedback_data = {
-                        "user_message": user_msg,
-                        "bot_response": bot_msg,
-                        "feedback": 1 if is_positive else -1,
-                        "comment": comment,
-                        "timestamp": datetime.now().isoformat()
-                    }
-                    
-                    bot.handle_feedback(
-                        message=user_msg,
-                        response=bot_msg,
-                        feedback=1 if is_positive else -1
-                    )
-                    
-                    # Clear feedback inputs
-                    return gr.update(value="")
-                    
-            except Exception as e:
-                logger.error(f"Error processing feedback: {e}")
-
+                }]
+
+        # Add these event handlers
+    feedback_positive.click(
+        lambda history: process_feedback(True, feedback_text.value, history),
+        inputs=[chatbot],
+        outputs=[feedback_text]
+    )
+    
+    feedback_negative.click(
+        lambda history: process_feedback(False, feedback_text.value, history),
+        inputs=[chatbot],
+        outputs=[feedback_text]
+    )
+    
+    # Add clear chat functionality
+    clear.click(lambda: None, None, chatbot)
         # Create enhanced Gradio interface
         with gr.Blocks(theme=gr.themes.Soft(
             primary_hue="blue",
@@ -574,32 +672,38 @@ def create_demo():
                         color: white;
                         font-weight: bold;
                     }
-                    /* Specific backgrounds for each emoji card */
-                    .feature-card:nth-child(1) {
-                        background: linear-gradient(135deg, #2193b0, #6dd5ed); /* Hospital emoji */
-                    }
-                    .feature-card:nth-child(2) {
-                        background: linear-gradient(135deg, #834d9b, #d04ed6); /* Magnifying glass emoji */
-                    }
-                    .feature-card:nth-child(3) {
-                        background: linear-gradient(135deg, #ff4b1f, #ff9068); /* Emergency emoji */
-                    }
-                    .feature-card:nth-child(4) {
-                        background: linear-gradient(135deg, #38ef7d, #11998e); /* Medical advice emoji */
-                    }
-                    /* Hover effect */
+                    .feature-card:nth-child(1) { background: linear-gradient(135deg, #2193b0, #6dd5ed); }
+                    .feature-card:nth-child(2) { background: linear-gradient(135deg, #834d9b, #d04ed6); }
+                    .feature-card:nth-child(3) { background: linear-gradient(135deg, #ff4b1f, #ff9068); }
+                    .feature-card:nth-child(4) { background: linear-gradient(135deg, #38ef7d, #11998e); }
                     .feature-card:hover {
                         transform: translateY(-5px);
                         box-shadow: 0 5px 15px rgba(0,0,0,0.2);
                     }
-                    }
-        
-                    /* Make emojis larger and more visible */
                     .feature-card span.emoji {
                         font-size: 2em;
                         display: block;
                         margin-bottom: 10px;
                     }
+                    .message-textbox textarea { resize: none; }
+                    #thumb-up, #thumb-down {
+                        min-width: 60px;
+                        padding: 8px;
+                        margin: 5px;
+                    }
+                    .chatbot-message {
+                        padding: 12px;
+                        margin: 8px 0;
+                        border-radius: 8px;
+                    }
+                    .user-message { background-color: #e3f2fd; }
+                    .assistant-message { background-color: #f5f5f5; }
+                    .feedback-section {
+                        margin-top: 20px;
+                        padding: 15px;
+                        border-radius: 8px;
+                        background-color: #f8f9fa;
+                    }
                 </style>
             """)
 
@@ -614,7 +718,6 @@ def create_demo():
             with gr.Row(elem_classes="header"):
                 gr.Markdown("""
                     # GP Medical Triage Assistant - Pearly
-                    
                     Welcome to your personal medical triage assistant. I'm here to help assess your symptoms and guide you to appropriate care.
                 """)
 
@@ -638,42 +741,6 @@ def create_demo():
                         <div>Medical Advice</div>
                     </div>
                 </div>
-                <style>
-                    /* Enable Enter key submission */
-                    .message-textbox textarea {
-                        resize: none;
-                    }
-                    
-                    /* Improve feedback buttons */
-                    #thumb-up, #thumb-down {
-                        min-width: 60px;
-                        padding: 8px;
-                        margin: 5px;
-                    }
-                    
-                    /* Improve chat messages */
-                    .chatbot-message {
-                        padding: 12px;
-                        margin: 8px 0;
-                        border-radius: 8px;
-                    }
-                    
-                    .user-message {
-                        background-color: #e3f2fd;
-                    }
-                    
-                    .assistant-message {
-                        background-color: #f5f5f5;
-                    }
-                    
-                    /* Feedback section */
-                    .feedback-section {
-                        margin-top: 20px;
-                        padding: 15px;
-                        border-radius: 8px;
-                        background-color: #f8f9fa;
-                    }
-                </style>
             """)
 
             # Chat Interface
@@ -696,8 +763,8 @@ def create_demo():
                             placeholder="Type your message here...",
                             lines=2,
                             scale=4,
-                            autofocus=True,  
-                            submit_on_enter=True 
+                            autofocus=True,
+                            submit_on_enter=True
                         )
                         submit = gr.Button("Send", variant="primary", scale=1)
 
@@ -708,37 +775,33 @@ def create_demo():
                     nhs_111_btn = gr.Button("📞 NHS 111 Info", variant="secondary")
                     booking_btn = gr.Button("📅 GP Booking", variant="secondary")
                     
-                    # Conversation Controls
+                    # Controls and Feedback
                     gr.Markdown("### Controls")
                     clear = gr.Button("🗑️ Clear Chat")
                     
-                    # Feedback Section
                     gr.Markdown("### Feedback")
                     with gr.Row():
                         feedback_positive = gr.Button("👍", elem_id="thumb-up")
                         feedback_negative = gr.Button("👎", elem_id="thumb-down")
                     
                     feedback_text = gr.Textbox(
-                        label="Additional comments (optional)",
-                        placeholder="Tell us more about your experience...",
+                        label="Additional comments",
+                        placeholder="Tell us more...",
                         lines=2,
-                        visible=True 
+                        visible=True
                     )
                     feedback_submit = gr.Button("Submit Feedback", visible=True)
-            # Examples Section
+
+            # Examples and Information
             with gr.Accordion("Example Messages", open=False):
-                gr.Examples(
-                    examples=[
-                        ["I've been having severe headaches for the past week"],
-                        ["I need to book a routine checkup"],
-                        ["I'm feeling very anxious lately and need help"],
-                        ["My child has had a fever for 2 days"],
-                        ["I need information about COVID-19 testing"]
-                    ],
-                    inputs=msg
-                )
+                gr.Examples([
+                    ["I've been having severe headaches for the past week"],
+                    ["I need to book a routine checkup"],
+                    ["I'm feeling very anxious lately and need help"],
+                    ["My child has had a fever for 2 days"],
+                    ["I need information about COVID-19 testing"]
+                ], inputs=msg)
 
-            # Information Accordions
             with gr.Accordion("NHS Services Guide", open=False):
                 gr.Markdown("""
                     ### Emergency Services (999)
@@ -758,74 +821,16 @@ def create_demo():
                 """)
 
             # Event Handlers
-            # Message submission handlers
-            msg.submit(
-                chat,
-                inputs=[msg, chatbot],
-                outputs=[chatbot]
-            ).then(
-                lambda: gr.update(value=""),
-                None,
-                [msg]
-            )
-            
-            submit.click(
-                chat,
-                inputs=[msg, chatbot],
-                outputs=[chatbot]
-            ).then(
-                lambda: gr.update(value=""),
-                None,
-                [msg]
-            )
-            
-            # Feedback handlers
-            feedback_positive.click(
-                lambda history: process_feedback(True, feedback_text.value, history),
-                inputs=[chatbot],
-                outputs=[feedback_text]
-            )
-            
-            feedback_negative.click(
-                lambda history: process_feedback(False, feedback_text.value, history),
-                inputs=[chatbot],
-                outputs=[feedback_text]
+            msg.submit(chat, [msg, chatbot], [chatbot]).then(
+                lambda: gr.update(value=""), None, [msg]
             )
             
-            feedback_submit.click(
-                lambda: gr.update(value=""),
-                outputs=[feedback_text]
+            submit.click(chat, [msg, chatbot], [chatbot]).then(
+                lambda: gr.update(value=""), None, [msg]
             )
-            
-            # Quick Action Button Handlers
-            def show_emergency_info():
-                return """🚨 Emergency Services (999)
-                - For life-threatening emergencies
-                - Severe chest pain
-                - Difficulty breathing
-                - Severe bleeding
-                - Loss of consciousness
-                """
-            
-            def show_nhs_111_info():
-                return """📞 NHS 111 Service
-                - Available 24/7
-                - Medical advice
-                - Local service information
-                - Urgent care guidance
-                """
-            
-            def show_booking_info():
-                return """📅 GP Booking Options
-                - Online booking
-                - Phone booking
-                - Routine appointments
-                - Urgent appointments
-                """
-            
-            emergency_btn.click(lambda: show_emergency_info(), outputs=[msg])
-            nhs_111_btn.click(lambda: show_nhs_111_info(), outputs=[msg])
-            booking_btn.click(lambda: show_booking_info(), outputs=[msg])
+
+            # Add queue for handling multiple users
+            demo.queue(concurrency_count=1, max_size=10)
 
         return demo
 
@@ -834,9 +839,23 @@ def create_demo():
         raise
 
 if __name__ == "__main__":
-    load_dotenv()  # Load environment variables
-    demo = create_demo()  # Launch the Gradio app
-    demo.launch(share=True)
-               
+    try:
+        # Initialize logging
+        logging.basicConfig(level=logging.INFO)
+        
+        # Load environment variables
+        load_dotenv()
+        
+        # Create and launch demo
+        demo = create_demo()
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True
+        )
+        
+    except Exception as e:
+        logger.error(f"Application startup failed: {e}")
+        raise