Update app.py

app.py

@@ -19,9 +19,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.cuda.amp import autocast
 from torch.utils.data import DataLoader
-import tensorflow as tf
-import keras
-import numpy as np
 
 # Hugging Face and Transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
@@ -39,16 +36,6 @@ from langchain_community.embeddings import HuggingFaceEmbeddings  # Updated impo
 from langchain_community.document_loaders import TextLoader  # Updated import
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-# Data Science and Visualization Libraries
-import pandas as pd
-import seaborn as sns
-import matplotlib.pyplot as plt
-from matplotlib.gridspec import GridSpec
-from sklearn.metrics import classification_report, confusion_matrix
-
-# Development and Testing
-import pytest
-from unittest.mock import Mock, patch
 
 # External Tools and APIs
 import wandb
@@ -108,54 +95,90 @@ class ModelManager:
             gc.collect()
 
 class PearlyBot:
-    def __init__(self, model_zip_path: str = "./checkpoint-500.zip", model_dir: str = "./checkpoint-500"):
-        self.model_dir = ModelManager.verify_and_extract_model(model_zip_path, model_dir)
-        self.setup_model(self.model_dir)
-        self.setup_rag()
-        self.conversation_history = []
-        self.last_interaction_time = time.time()
-        self.interaction_cooldown = 1.0  # seconds
-        
-    def setup_model(self, model_path: str):
-        """Initialize the model with proper error handling"""
+    def __init__(self):
         try:
-            logger.info("Starting model initialization...")
-            ModelManager.clear_gpu_memory()
+            # Use the correct model path from your space
+            self.repo_id = "Pearilsa/pearly_med_triage_chatbot_kagglex"
+            self.model_filename = "pearly_model.zip"
+            self.setup_model()
+            self.setup_rag()
+            self.conversation_history = []
+            self.last_interaction_time = time.time()
+            self.interaction_cooldown = 1.0
+        except Exception as e:
+            logger.error(f"Error initializing bot: {e}")
+            raise
+
+    def setup_model(self):
+        """Initialize model from Hugging Face space"""
+        try:
+            logger.info(f"Loading model from {self.repo_id}")
             
-            # Load tokenizer
-            try:
-                self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-                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
+            # Download and prepare model path
+            local_model_path = os.path.join(os.getcwd(), "models")
+            os.makedirs(local_model_path, exist_ok=True)
+            
+            # Load tokenizer and model from the space
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.repo_id,
+                token=os.getenv("HF_TOKEN"),  # Use your Hugging Face token
+                cache_dir=local_model_path
+            )
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            logger.info("Tokenizer loaded successfully")
+            
+            # Load model with 8-bit quantization
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.repo_id,
+                token=os.getenv("HF_TOKEN"),
+                device_map="auto",
+                load_in_8bit=True,
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True,
+                cache_dir=local_model_path
+            )
+            self.model.eval()
+            logger.info("Model loaded successfully")
             
-            # Load model
-            try:
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    model_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 in model setup: {str(e)}")
             raise
+            
 
     def setup_rag(self):
         try:
-            logger.info("Setting up RAG system...")
+            # Add configuration options
+            self.chunk_size = 300
+            self.chunk_overlap = 100
+            self.num_relevant_chunks = 3
+            
+            # Load knowledge base
+            knowledge_base = self._load_knowledge_base()
+            
+            # Setup embeddings with error handling
+            self.embeddings = self._initialize_embeddings()
+            
+            # Enhanced text splitting
+            texts = self._split_texts(knowledge_base)
+            
+            # Create vector store with metadata
+            self.vector_store = FAISS.from_texts(
+                texts,
+                self.embeddings,
+                metadatas=[{"source": f"chunk_{i}"} for i in range(len(texts))]
+            )
+            
+            # Add validation
+            self._validate_rag_setup()
+            
+        except Exception as e:
+            logger.error(f"RAG setup failed: {str(e)}")
+            raise
             # Load your knowledge base content
-            knowledge_base = {
-                "triage_scenarios.txt": """Medical Triage Scenarios and Responses:
+    def _load_knowledge_base(self):
+    # Add validation and error handling for knowledge base loading
+        return {
+            "triage_scenarios.txt": """Medical Triage Scenarios and Responses:
 
 EMERGENCY (999) SCENARIOS:
 1. Cardiovascular:
@@ -456,6 +479,48 @@ PROFESSIONAL CONDUCT:
             logger.error(f"Error setting up RAG: {str(e)}")
             raise
 
+    def _validate_rag_setup(self):
+        """Validate RAG system setup"""
+        try:
+            # Verify embeddings are working
+            test_text = "This is a test embedding"
+            test_embedding = self.embeddings.encode(test_text)
+            assert len(test_embedding) > 0
+            
+            # Verify vector store is operational
+            test_results = self.vector_store.similarity_search(test_text, k=1)
+            assert len(test_results) > 0
+            
+            logger.info("RAG system validation successful")
+            return True
+        except Exception as e:
+            logger.error(f"RAG system validation failed: {str(e)}")
+            raise
+
+    def _initialize_embeddings(self):
+        try:
+            return HuggingFaceEmbeddings(
+                model_name="sentence-transformers/all-MiniLM-L6-v2",
+                cache_folder="./embeddings_cache"  # Added caching
+            )
+        except Exception as e:
+            logger.error(f"Failed to initialize embeddings: {str(e)}")
+            raise
+
+    def _split_texts(self, knowledge_base):
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.chunk_size,
+            chunk_overlap=self.chunk_overlap,
+            length_function=len,
+            add_start_index=True
+        )
+        
+        all_texts = []
+        for content in knowledge_base.values():
+            texts = splitter.split_text(content)
+            all_texts.extend(texts)
+        return all_texts
+
     def get_relevant_context(self, query):
         try:
             docs = self.vector_store.similarity_search(query, k=3)
@@ -576,28 +641,8 @@ Guidelines:
         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:
         # Initialize bot
         bot = PearlyBot()
@@ -608,10 +653,7 @@ def create_demo():
                 if not message.strip():
                     return history
                 
-                # Generate response
                 response = bot.generate_response(message, history)
-                
-                # Update history with proper formatting
                 history.append({
                     "role": "user",
                     "content": message
@@ -621,7 +663,6 @@ def create_demo():
                     "content": response
                 })
                 return history
-                
             except Exception as e:
                 logger.error(f"Chat error: {e}")
                 return history + [{
@@ -642,11 +683,57 @@ def create_demo():
                     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="")
+
+        # Create Gradio interface
+        with gr.Blocks(theme=gr.themes.Soft(...)) as demo:
+            # 1. First, create all UI elements
+            # CSS styles
+            gr.HTML("""<style>...""")
+
+            # Emergency Banner
+            gr.HTML("""<div class="emergency-banner">...""")
+
+            # Header
+            with gr.Row(elem_classes="header"):
+                gr.Markdown("""# GP Medical Triage Assistant...""")
+
+            # Features Grid
+            gr.HTML("""<div class="features-grid">...""")
+
+            # Chat Interface
+            with gr.Row():
+                with gr.Column(scale=4):
+                    chatbot = gr.Chatbot(...)
+                    with gr.Row():
+                        msg = gr.Textbox(...)
+                        submit = gr.Button(...)
+
+                with gr.Column(scale=1):
+                    # Quick Actions
+                    emergency_btn = gr.Button("🚨 Emergency Info", variant="secondary")
+                    nhs_111_btn = gr.Button("📞 NHS 111 Info", variant="secondary")
+                    booking_btn = gr.Button("📅 GP Booking", variant="secondary")
+                    
+                    # Controls
+                    clear = gr.Button("🗑️ Clear Chat")
+                    
+                    # Feedback
+                    with gr.Row():
+                        feedback_positive = gr.Button("👍", elem_id="thumb-up")
+                        feedback_negative = gr.Button("👎", elem_id="thumb-down")
+                    feedback_text = gr.Textbox(...)
+                    feedback_submit = gr.Button(...)
+
+            # Examples and Guide
+            with gr.Accordion("Example Messages", open=False):
+                gr.Examples([...])
+
+            with gr.Accordion("NHS Services Guide", open=False):
+                gr.Markdown("""...""")
         
 
         # Create enhanced Gradio interface
@@ -711,33 +798,6 @@ def create_demo():
                     }
                 </style>
             """)
-            # Event Handlers - Moved inside the gr.Blocks context
-            msg.submit(chat, [msg, chatbot], [chatbot]).then(
-                lambda: gr.update(value=""), None, [msg]
-            )
-            
-            submit.click(chat, [msg, chatbot], [chatbot]).then(
-                lambda: gr.update(value=""), None, [msg]
-            )
-            
-            # Feedback handlers
-            feedback_positive.click(
-                lambda h: process_feedback(True, feedback_text.value, h),
-                inputs=[chatbot],
-                outputs=[feedback_text]
-            )
-            
-            feedback_negative.click(
-                lambda h: process_feedback(False, feedback_text.value, h),
-                inputs=[chatbot],
-                outputs=[feedback_text]
-            )
-            
-            # Clear chat
-            clear.click(lambda: None, None, chatbot)
-
-            # Add queue for handling multiple users
-            demo.queue(concurrency_count=1, max_size=10)
 
             # Emergency Banner
             gr.HTML("""
@@ -854,6 +914,64 @@ def create_demo():
 
         
 
+        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
+            """
+
+            # Chat handlers
+            msg.submit(chat, [msg, chatbot], [chatbot]).then(
+                lambda: gr.update(value=""), None, [msg]
+            )
+            
+            submit.click(chat, [msg, chatbot], [chatbot]).then(
+                lambda: gr.update(value=""), None, [msg]
+            )
+
+            # Quick action handlers
+            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])
+            
+            # Feedback handlers
+            feedback_positive.click(
+                lambda h: process_feedback(True, feedback_text.value, h),
+                inputs=[chatbot],
+                outputs=[feedback_text]
+            )
+            
+            feedback_negative.click(
+                lambda h: process_feedback(False, feedback_text.value, h),
+                inputs=[chatbot],
+                outputs=[feedback_text]
+            )
+            
+            # Clear chat
+            clear.click(lambda: None, None, chatbot)
+
+            # 3. Finally, add the queue
+            demo.queue(concurrency_count=1, max_size=10)
+
         return demo
 
     except Exception as e: