Update app.py

app.py

@@ -9,16 +9,12 @@ from sentence_transformers import SentenceTransformer
 from peft import get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training
 import faiss
 import numpy as np
-from tqdm import tqdm
 from datasets import load_dataset
-from dataclasses import dataclass, field
 from datetime import datetime
 import json
 from huggingface_hub import login
 from dotenv import load_dotenv
 
-
-
 # Load environment variables
 load_dotenv()
 
@@ -30,62 +26,32 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # Retrieve secrets securely from environment variables
-kaggle_username = os.getenv("KAGGLE_USERNAME")
-kaggle_key = os.getenv("KAGGLE_KEY")
 hf_token = os.getenv("HF_TOKEN")
-wandb_key = os.getenv("WANDB_API_KEY")
-
-# Log in to Hugging Face
 if hf_token:
     login(token=hf_token)
-else:
-    logger.warning("Hugging Face token not found in environment variables.")
-
-@dataclass
-class AdaptiveBotConfig:
-    """Configuration for adaptive medical triage bot"""
-    MODEL_NAME: str = "google/gemma-7b"
-    EMBEDDING_MODEL: str = "sentence-transformers/all-MiniLM-L6-v2"
-    
-    # LoRA parameters
-    LORA_R: int = 8
-    LORA_ALPHA: int = 16
-    LORA_DROPOUT: float = 0.1
-    LORA_TARGET_MODULES: List[str] = field(default_factory=lambda: ["q_proj", "v_proj"])
-    
-    # Training parameters
-    MAX_LENGTH: int = 512
-    BATCH_SIZE: int = 1
-    LEARNING_RATE: float = 1e-4
-    
-    # Adaptive learning parameters
-    MIN_FEEDBACK_FOR_UPDATE: int = 5
-    FEEDBACK_HISTORY_SIZE: int = 100
-    LEARNING_RATE_DECAY: float = 0.95
 
 class AdaptiveMedicalBot:
     def __init__(self):
-        self.config = AdaptiveBotConfig()
+        self.config = self.AdaptiveBotConfig()
         self.setup_models()
         self.load_datasets()
         self.setup_adaptive_learning()
-        self.document_relevance = {}
-        
+        self.conversation_history = []  # Maintain conversation history
+
+    class AdaptiveBotConfig:
+        MODEL_NAME = "google/gemma-7b"
+        EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+        LORA_R = 8
+        LORA_ALPHA = 16
+        LORA_DROPOUT = 0.1
+        LORA_TARGET_MODULES = ["q_proj", "v_proj"]
+        MAX_LENGTH = 512
+        BATCH_SIZE = 1
+        LEARNING_RATE = 1e-4
+
     def setup_adaptive_learning(self):
         """Initialize adaptive learning components"""
         self.feedback_history = []
-        self.conversation_patterns = {}
-        self.learning_buffer = []
-        
-        # Load existing learning data if available
-        try:
-            if os.path.exists('learning_data.json'):
-                with open('learning_data.json', 'r') as f:
-                    data = json.load(f)
-                    self.conversation_patterns = data.get('patterns', {})
-                    self.feedback_history = data.get('feedback', [])
-        except Exception as e:
-            logger.warning(f"Could not load learning data: {e}")
 
     def setup_models(self):
         """Initialize models with LoRA and quantization"""
@@ -95,21 +61,15 @@ class AdaptiveMedicalBot:
                 bnb_4bit_quant_type="nf4",
                 bnb_4bit_compute_dtype=torch.float16
             )
-            
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.config.MODEL_NAME,
-                trust_remote_code=True
-            )
-            
+
+            self.tokenizer = AutoTokenizer.from_pretrained(self.config.MODEL_NAME)
             base_model = AutoModelForCausalLM.from_pretrained(
                 self.config.MODEL_NAME,
                 quantization_config=bnb_config,
-                device_map="auto",
-                trust_remote_code=True
+                device_map="auto"
             )
-            
+
             base_model = prepare_model_for_kbit_training(base_model)
-            
             lora_config = LoraConfig(
                 r=self.config.LORA_R,
                 lora_alpha=self.config.LORA_ALPHA,
@@ -118,720 +78,185 @@ class AdaptiveMedicalBot:
                 bias="none",
                 task_type=TaskType.CAUSAL_LM
             )
-            
+
             self.model = get_peft_model(base_model, lora_config)
             self.embedding_model = SentenceTransformer(self.config.EMBEDDING_MODEL)
-            
         except Exception as e:
             logger.error(f"Error setting up models: {e}")
             raise
 
     def load_datasets(self):
-        """Load and process datasets for RAG"""
+        """Load and prepare datasets for RAG"""
         try:
             datasets = {
                 "medqa": load_dataset("medalpaca/medical_meadow_medqa", split="train[:500]"),
                 "diagnosis": load_dataset("wasiqnauman/medical-diagnosis-synthetic", split="train[:500]"),
                 "persona": load_dataset("AlekseyKorshuk/persona-chat", split="train[:500]")
             }
-            
+
             self.documents = []
-            
             for dataset_name, dataset in datasets.items():
                 for item in dataset:
                     if dataset_name == "persona":
                         if isinstance(item.get('personality'), list):
-                            self.documents.append({
-                                'text': " ".join(item['personality']),
-                                'type': 'persona'
-                            })
+                            self.documents.append({'text': " ".join(item['personality']), 'type': 'persona'})
                     else:
                         if 'input' in item and 'output' in item:
-                            self.documents.append({
-                                'text': f"{item['input']}\n{item['output']}",
-                                'type': dataset_name
-                            })
-            
+                            self.documents.append({'text': f"{item['input']}\n{item['output']}", 'type': dataset_name})
+
             self._create_index()
-            
         except Exception as e:
             logger.error(f"Error loading datasets: {e}")
             raise
 
     def _create_index(self):
         """Create FAISS index for RAG"""
-        sample_embedding = self.embedding_model.encode("sample text")
-        self.index = faiss.IndexFlatIP(sample_embedding.shape[0])
-        
-        batch_size = 32
-        for i in range(0, len(self.documents), batch_size):
-            batch = self.documents[i:i + batch_size]
-            texts = [doc['text'] for doc in batch]
-            embeddings = self.embedding_model.encode(texts)
-            self.index.add(np.array(embeddings))
-
-    def analyze_conversation_context(self, message: str, history: List[tuple]) -> Dict[str, Any]:
-        """Analyze conversation context to determine appropriate follow-up questions"""
         try:
-            # Extract key information
-            mentioned_symptoms = set()
-            time_indicators = set()
-            severity_indicators = set()
-            
-            # Analyze current message and history
-            for msg in [message] + [h[0] for h in (history or [])]:
-                msg_lower = msg.lower()
-                
-                # Update conversation patterns
-                pattern_key = self._extract_pattern_key(msg_lower)
-                if pattern_key in self.conversation_patterns:
-                    self.conversation_patterns[pattern_key]['frequency'] += 1
-                else:
-                    self.conversation_patterns[pattern_key] = {
-                        'frequency': 1,
-                        'successful_responses': []
-                    }
-            
-            return {
-                'needs_follow_up': True,  # Always encourage follow-up questions
-                'conversation_depth': len(history) if history else 0,
-                'pattern_key': pattern_key
-            }
-            
-        except Exception as e:
-            logger.error(f"Error analyzing conversation: {e}")
-            return {'needs_follow_up': True}
+            sample_embedding = self.embedding_model.encode("sample text")
+            self.index = faiss.IndexFlatIP(sample_embedding.shape[0])
 
-    def _extract_pattern_key(self, message: str) -> str:
-        """Extract conversation pattern key for learning"""
-        # Simplified pattern extraction - can be enhanced based on learning
-        words = message.lower().split()
-        return " ".join(sorted(set(words))[:5])
+            embeddings = [self.embedding_model.encode(doc['text']) for doc in self.documents]
+            self.index.add(np.array(embeddings))
+        except Exception as e:
+            logger.error(f"Error creating FAISS index: {e}")
+            raise
 
-    def generate_follow_up_questions(self, context: Dict[str, Any], message: str) -> List[str]:
-        """Generate contextual follow-up questions"""
+    def generate_follow_up_questions(self, message: str, context: Dict[str, Any]) -> List[str]:
+        """Generate follow-up questions based on context"""
         try:
-            # Use the model to generate follow-up questions
-            prompt = f"""Given the patient message: "{message}"
-            Generate relevant follow-up questions to better understand their situation. 
-            Focus on: timing, severity, associated symptoms, impact on daily life.
-            Do not make diagnoses or suggest treatments.
-            
+            prompt = f"""Patient message: "{message}"
+            Generate relevant follow-up questions focusing on timing, severity, associated symptoms, and impact on daily life.
             Questions:"""
-            
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                max_length=self.config.MAX_LENGTH,
-                truncation=True
-            ).to(self.model.device)
-            
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=100,
-                temperature=0.7,
-                do_sample=True
-            )
-            
+
+            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=self.config.MAX_LENGTH).to(self.model.device)
+            outputs = self.model.generate(inputs['input_ids'], max_new_tokens=50, temperature=0.7, do_sample=True)
             questions = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return questions.split('\n')
-            
+            return questions.split("\n")
         except Exception as e:
             logger.error(f"Error generating follow-up questions: {e}")
             return ["Could you tell me more about when this started?"]
 
-    def update_from_feedback(self, message: str, response: str, feedback: int):
-        """Process feedback for adaptive learning"""
-        try:
-            self.feedback_history.append({
-                'message': message,
-                'response': response,
-                'feedback': feedback,
-                'timestamp': datetime.now().isoformat()
-            })
-            
-            # Update conversation patterns
-            pattern_key = self._extract_pattern_key(message)
-            if pattern_key in self.conversation_patterns:
-                if feedback > 0:
-                    self.conversation_patterns[pattern_key]['successful_responses'].append(response)
-            
-            # Save learning data periodically
-            if len(self.feedback_history) % 10 == 0:
-                self._save_learning_data()
-            
-            # Update model if enough feedback
-            if len(self.feedback_history) >= self.config.MIN_FEEDBACK_FOR_UPDATE:
-                self._update_model_from_feedback()
-                
-        except Exception as e:
-            logger.error(f"Error processing feedback: {e}")
+    def assess_symptom_severity(self, message: str) -> str:
+        """Assess severity based on keywords in the message"""
+        if "severe" in message.lower() or "emergency" in message.lower():
+            return "emergency"
+        elif "persistent" in message.lower() or "moderate" in message.lower():
+            return "urgent"
+        return "routine"
 
-    def store_interaction(self, interaction: Dict[str, Any]):
-        """Store interaction for adaptive learning"""
+    def generate_response(self, message: str) -> Dict[str, Any]:
+        """Generate a response based on the message"""
         try:
-            self.learning_buffer.append(interaction)
-            
-            # Update learning patterns
-            if len(self.learning_buffer) >= 10:
-                self._update_learning_model()
-                self._save_learning_data()
-                self.learning_buffer = []
-                
-        except Exception as e:
-            logger.error(f"Error storing interaction: {e}")
+            severity = self.assess_symptom_severity(message)
+            response = ""
 
-    def _update_learning_model(self):
-        """Update model based on accumulated learning"""
-        try:
-            # Process learning buffer
-            successful_interactions = [
-                interaction for interaction in self.learning_buffer
-                if interaction.get('feedback', 0) > 0
-            ]
-            
-            if successful_interactions:
-                # Update conversation patterns
-                for interaction in successful_interactions:
-                    pattern_key = self._extract_pattern_key(interaction['message'])
-                    if pattern_key in self.conversation_patterns:
-                        self.conversation_patterns[pattern_key]['successful_responses'].append(
-                            interaction['response']
-                        )
-                
-                # Update document relevance
-                for interaction in successful_interactions:
-                    for doc in interaction.get('relevant_docs', []):
-                        doc_key = doc['text'][:100]
-                        if doc_key in self.document_relevance:
-                            self.document_relevance[doc_key]['success_count'] += 1
-                            
-                logger.info("Updated learning model with new patterns")
-                
-        except Exception as e:
-            logger.error(f"Error updating learning model: {e}")
-    
-    def generate_context_questions(self, message: str, history: List[tuple], context: Dict[str, Any]) -> List[str]:
-        """Generate context-aware follow-up questions"""
-        try:
-            # Create dynamic question generation prompt
-            prompt = f"""Based on the conversation context, generate appropriate follow-up questions.
-            Consider:
-            - Understanding the main concern
-            - Timeline and progression
-            - Impact on daily life
-            - Related symptoms or factors
-            - Previous treatments or consultations
-            
-            Current message: {message}
-            Context: {context}
-            
-            Generate questions:"""
-            
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                max_length=self.config.MAX_LENGTH,
-                truncation=True
-            ).to(self.model.device)
-            
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=150,
-                temperature=0.7,
-                do_sample=True
-            )
-            
-            questions = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return [q.strip() for q in questions.split("\n") if "?" in q]
-            
-        except Exception as e:
-            logger.error(f"Error generating context questions: {e}")
-            return ["Could you tell me more about your concerns?"]
+            # Retrieve relevant documents from FAISS
+            query_embedding = self.embedding_model.encode([message])
+            _, indices = self.index.search(query_embedding, k=5)
+            relevant_docs = [self.documents[idx]['text'] for idx in indices[0]]
 
-    def _save_learning_data(self):
-        """Save learning data to disk"""
-        try:
-            data = {
-                'patterns': self.conversation_patterns,
-                'feedback': self.feedback_history[-100:]  # Keep last 100 entries
-            }
-            with open('learning_data.json', 'w') as f:
-                json.dump(data, f)
-        except Exception as e:
-            logger.error(f"Error saving learning data: {e}")
+            prompt = f"""As a compassionate medical assistant, analyze the patient message: "{message}".
+            Consider relevant knowledge and the following documents:\n{relevant_docs}.
+            Respond with empathy, follow-up questions, and care guidance."""
 
-    def _update_model_from_feedback(self):
-        """Update model based on feedback"""
-        try:
-            positive_feedback = [f for f in self.feedback_history if f['feedback'] > 0]
-            if len(positive_feedback) >= self.config.MIN_FEEDBACK_FOR_UPDATE:
-                # Prepare training data from successful interactions
-                training_data = []
-                for feedback in positive_feedback:
-                    training_data.append({
-                        'input_ids': self.tokenizer(
-                            feedback['message'],
-                            return_tensors='pt'
-                        ).input_ids,
-                        'labels': self.tokenizer(
-                            feedback['response'],
-                            return_tensors='pt'
-                        ).input_ids
-                    })
-                
-                # Update model (simplified for example)
-                logger.info("Updating model from feedback")
-                self.feedback_history = []  # Clear history after update
-                
-        except Exception as e:
-            logger.error(f"Error updating model from feedback: {e}")
+            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=self.config.MAX_LENGTH).to(self.model.device)
+            outputs = self.model.generate(inputs['input_ids'], max_new_tokens=100, temperature=0.7, do_sample=True)
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    def analyze_symptom_context(self, message: str, history: List[tuple]) -> Dict[str, Any]:
-        """Enhanced symptom analysis with context learning"""
-        try:
-            current_symptoms = set()
-            temporal_info = {}
-            related_conditions = set()
-            conversation_depth = len(history) if history else 0
-            
-            # Analyze full conversation context
-            all_messages = [message] + [h[0] for h in (history or [])]
-            all_responses = [h[1] for h in (history or [])]
+            follow_ups = self.generate_follow_up_questions(message, {})
+            response += f"\n{follow_ups[0]}"
             
-            # Extract conversation context
-            context = {
-                'symptoms_mentioned': current_symptoms,
-                'temporal_info': temporal_info,
-                'conversation_depth': conversation_depth,
-                'needs_clarification': True,
-                'specialist_referral_needed': False,
-                'previous_questions': set()
-            }
-            
-            if history:
-                # Learn from previous interactions
-                for prev_msg, prev_resp in history:
-                    if "?" in prev_resp:
-                        context['previous_questions'].add(prev_resp.split("?")[0] + "?")
-            
-            return context
-            
-        except Exception as e:
-            logger.error(f"Error in symptom analysis: {e}")
-            return {'needs_clarification': True}
+            # Append response to conversation history
+            self.conversation_history.append((message, response))
 
-    def generate_targeted_questions(self, symptoms: Dict[str, Any], history: List[tuple]) -> List[str]:
-        """Generate context-aware follow-up questions"""
-        try:
-            # Use the model to generate relevant questions based on context
-            context_prompt = f"""Based on the patient's symptoms and conversation history, generate 3 specific follow-up questions.
-            Focus on:
-            1. Symptom details (duration, severity, patterns)
-            2. Impact on daily life
-            3. Related symptoms or conditions
-            4. Previous treatments or consultations
-            
-            Do not ask about:
-            - Questions already asked
-            - Diagnostic conclusions
-            - Treatment recommendations
-            
-            Current context: {symptoms}
-            Previous questions asked: {symptoms.get('previous_questions', set())}
-            
-            Generate questions:"""
-            
-            inputs = self.tokenizer(
-                context_prompt,
-                return_tensors="pt",
-                max_length=self.config.MAX_LENGTH,
-                truncation=True
-            ).to(self.model.device)
-            
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=150,
-                temperature=0.7,
-                do_sample=True
-            )
-            
-            questions = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return [q.strip() for q in questions.split("\n") if "?" in q]
-            
-        except Exception as e:
-            logger.error(f"Error generating questions: {e}")
-            return ["Could you tell me more about your symptoms?"]
+            # Add care level guidance
+            if severity == "emergency":
+                response += "\nThis seems urgent. Please call 999 immediately."
+            elif severity == "urgent":
+                response += "\nConsider calling NHS 111 for urgent assistance."
 
-    def analyze_medical_context(self, message: str, history: List[tuple]) -> Dict[str, Any]:
-        """Comprehensive medical context analysis"""
-        try:
-            # Initialize context tracking
-            context = {
-                'conversation_depth': len(history) if history else 0,
-                'needs_follow_up': True,
-                'previous_interactions': [],
-                'care_pathway': 'initial_triage',
-                'consultation_type': 'general',
-            }
-            
-            # Analyze current conversation flow
-            all_messages = [message] + [h[0] for h in (history or [])]
-            
-            # Build contextual understanding
-            for msg in all_messages:
-                msg_lower = msg.lower()
-                
-                # Track conversation patterns
-                pattern_key = self._extract_pattern_key(msg_lower)
-                if pattern_key in self.conversation_patterns:
-                    self.conversation_patterns[pattern_key]['frequency'] += 1
-                    context['previous_patterns'] = self.conversation_patterns[pattern_key]
-                
-                # Update learning patterns
-                self._update_learning_patterns(msg_lower, context)
-            
-            return context
-            
+            return {'response': response}
         except Exception as e:
-            logger.error(f"Error in medical context analysis: {e}")
-            return {'needs_follow_up': True}
-    
-    def generate_adaptive_response(self, message: str, history: List[tuple] = None) -> Dict[str, Any]:
-        """Generate comprehensive triage response"""
+            logger.error(f"Error generating response: {e}")
+            return {
+                'response': "I'm experiencing technical issues. If this is an emergency, please call 999 immediately.",
+            }
+
+    def handle_feedback(self, message: str, response: str, feedback: int):
+        """Update model based on feedback"""
         try:
-            # Analyze medical context
-            context = self.analyze_medical_context(message, history)
-            
-            # Retrieve relevant knowledge
-            query_embedding = self.embedding_model.encode([message])
-            _, indices = self.index.search(query_embedding, k=5)
-            relevant_docs = [self.documents[idx] for idx in indices[0]]
-            
-            # Build conversation history
-            conv_history = "\n".join([f"Patient: {h[0]}\nPearly: {h[1]}" for h in (history or [])])
-            
-            # Create dynamic prompt based on context
-            prompt = f"""As Pearly, a compassionate GP medical triage assistant, help assess the patient's needs and provide appropriate guidance.
-    
-    Previous Conversation:
-    {conv_history}
-    
-    Current Message: {message}
-    
-    Medical Knowledge Context:
-    {[doc['text'] for doc in relevant_docs]}
-    
-    Guidelines:
-    - Show empathy and understanding
-    - Ask relevant follow-up questions
-    - Guide to appropriate care level (GP, 111, emergency services)
-    - Consider all aspects of patient care
-    - Do not diagnose or recommend treatments
-    - Focus on understanding concerns and proper healthcare guidance
-    
-    Response:"""
-    
-            # Generate base response
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                max_length=self.config.MAX_LENGTH,
-                truncation=True
-            ).to(self.model.device)
-            
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=300,
-                temperature=0.7,
-                do_sample=True
-            )
-            
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            # Generate contextual follow-up questions
-            if context['needs_follow_up']:
-                follow_ups = self.generate_context_questions(message, history, context)
-                if follow_ups:
-                    response = f"{response}\n\n{follow_ups[0]}"
-            
-            # Store interaction for learning
-            self.store_interaction({
+            self.feedback_history.append({
                 'message': message,
                 'response': response,
-                'context': context,
-                'relevant_docs': relevant_docs,
+                'feedback': feedback,
                 'timestamp': datetime.now().isoformat()
             })
-            
-            return {
-                'response': response,
-                'context': context
-            }
-            
-        except Exception as e:
-            logger.error(f"Error generating response: {e}")
-            return {
-                'response': "I apologize, but I'm having technical difficulties. If this is an emergency, please call 999 immediately. For urgent concerns, call 111.",
-                'context': {}
-            }
 
-    
-    def _add_specialist_guidance(self, response: str, context: Dict[str, Any]) -> str:
-        """Add specialist referral guidance to response"""
-        try:
-            specialist_prompt = f"""Based on the symptoms and context, suggest appropriate specialist care pathways. 
-            Context: {context}
-            Current response: {response}
-            
-            Add appropriate specialist referral guidance:"""
-            
-            inputs = self.tokenizer(
-                specialist_prompt,
-                return_tensors="pt",
-                max_length=self.config.MAX_LENGTH,
-                truncation=True
-            ).to(self.model.device)
-            
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=150,
-                temperature=0.7,
-                do_sample=True
-            )
-            
-            specialist_guidance = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            return f"{response}\n\n{specialist_guidance}"
-            
-        except Exception as e:
-            logger.error(f"Error adding specialist guidance: {e}")
-            return response
-    
-    def update_learning_from_interaction(self, interaction: Dict[str, Any]):
-        """Update adaptive learning system from interaction"""
-        try:
-            # Extract key information
-            message = interaction['message']
-            response = interaction['response']
-            context = interaction['context']
-            relevant_docs = interaction.get('relevant_docs', [])
-            
-            # Update conversation patterns
-            pattern_key = self._extract_pattern_key(message)
-            if pattern_key in self.conversation_patterns:
-                self.conversation_patterns[pattern_key]['frequency'] += 1
-                if context.get('successful_response'):
-                    self.conversation_patterns[pattern_key]['successful_responses'].append(response)
-            
-            # Update document relevance scores
-            for doc in relevant_docs:
-                doc_key = doc['text'][:100]  # Use first 100 chars as key
-                if doc_key in self.document_relevance:
-                    self.document_relevance[doc_key]['usage_count'] += 1
-                    if context.get('successful_response'):
-                        self.document_relevance[doc_key]['success_count'] += 1
-            
-            # Save learning data periodically
-            if len(self.learning_buffer) >= 10:
-                self._save_learning_data()
-                self.learning_buffer = []
-                
+            if len(self.feedback_history) >= 10:
+                # Implement learning updates from feedback
+                self.feedback_history = []  # Reset history after learning update
         except Exception as e:
-            logger.error(f"Error updating learning system: {e}")
+            logger.error(f"Error processing feedback: {e}")
 
 def create_demo():
-    """Create enhanced Gradio interface with advanced features and styling"""
+    """Set up Gradio interface for the chatbot"""
     try:
         bot = AdaptiveMedicalBot()
-        
-        def chat(message: str, history: List[Dict[str, str]], urgency_filter: str = "all"):
+
+        def chat(message: str, history: List[Dict[str, str]]):
             try:
-                # Convert history to the format expected by the bot
                 bot_history = [(h["user"], h["bot"]) for h in history] if history else []
-                
-                # Generate response
-                response_data = bot.generate_adaptive_response(message, bot_history)
+                response_data = bot.generate_response(message)
                 response = response_data['response']
-                
-                # Format response for Gradio chat
+
                 history.append({"role": "user", "content": message})
                 history.append({"role": "assistant", "content": response})
-                
                 return history
-
             except Exception as e:
                 logger.error(f"Chat error: {e}")
                 return history + [
                     {"role": "user", "content": message},
-                    {"role": "assistant", "content": "I apologize, but I'm experiencing technical difficulties. For emergencies, please call 999."}
+                    {"role": "assistant", "content": "I'm experiencing technical difficulties. For emergencies, call 999."}
                 ]
 
         def process_feedback(feedback: str, history: List[Dict[str, str]], comment: str = ""):
-            """Process feedback with optional comment"""
             try:
                 if history and len(history) >= 2:
                     last_user_msg = history[-2]["content"]
                     last_bot_msg = history[-1]["content"]
-                    bot.update_from_feedback(
-                        last_user_msg,
-                        last_bot_msg,
-                        1 if feedback == "👍" else -1,
-                        comment=comment
-                    )
+                    bot.handle_feedback(last_user_msg, last_bot_msg, 1 if feedback == "👍" else -1)
             except Exception as e:
                 logger.error(f"Error processing feedback: {e}")
 
-        # Create enhanced Gradio interface
-        with gr.Blocks(theme=gr.themes.Soft(
-            primary_hue="blue",
-            secondary_hue="indigo",
-            neutral_hue="slate",
-            font=gr.themes.GoogleFont("Inter")
-        )) as demo:
-            # Custom CSS for enhanced styling
-            gr.HTML("""
-                <style>
-                    .container { max-width: 900px; margin: auto; }
-                    .header { text-align: center; padding: 20px; }
-                    .emergency-banner {
-                        background-color: #ff4444;
-                        color: white;
-                        padding: 10px;
-                        text-align: center;
-                        font-weight: bold;
-                        margin-bottom: 20px;
-                    }
-                    .features-grid {
-                        display: grid;
-                        grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
-                        gap: 20px;
-                        padding: 20px;
-                    }
-                    .feature-card {
-                        background: #f8f9fa;
-                        padding: 15px;
-                        border-radius: 10px;
-                        text-align: center;
-                    }
-                </style>
-            """)
+        with gr.Blocks() as demo:
+            chatbot = gr.Chatbot(value=[{"role": "assistant", "content": "Hello! I'm Pearly, your GP Triage medical assistant. How can I help you today?"}],
+                height=500,
+                elem_id="chatbot",
+                type="messages",
+                show_label=False
+            )
+            
+            msg = gr.Textbox(
+                label="Your message",
+                placeholder="Type your message here...",
+                lines=2
+            )
+            submit = gr.Button("Send", variant="primary")
+            
+            feedback = gr.Radio(
+                choices=["👍", "👎"],
+                label="Was this response helpful?",
+                visible=True
+            )
+            feedback_text = gr.Textbox(
+                label="Additional comments (optional)",
+                placeholder="Tell us more about your experience...",
+                lines=2
+            )
             
-            # Emergency Banner
-            gr.HTML("""
-                <div class="emergency-banner">
-                    🚨 For medical emergencies, always call 999 immediately 🚨
-                </div>
-            """)
-
-            # Header Section
-            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.
-                """)
-
-            # Main Features Grid
-            gr.HTML("""
-                <div class="features-grid">
-                    <div class="feature-card">
-                        🏥 GP Appointments
-                    </div>
-                    <div class="feature-card">
-                        🔍 Symptom Assessment
-                    </div>
-                    <div class="feature-card">
-                        ⚡ Urgent Care Guide
-                    </div>
-                    <div class="feature-card">
-                        💊 Medical Advice
-                    </div>
-                </div>
-            """)
-
-            # Chat Interface
-            with gr.Row():
-                with gr.Column(scale=4):
-                    chatbot = gr.Chatbot(
-                        value=[{
-                            "role": "assistant",
-                            "content": "Hello! I'm Pearly, your GP medical assistant. How can I help you today?"
-                        }],
-                        height=500,
-                        elem_id="chatbot",
-                        type="messages",
-                        show_label=False
-                    )
-                    
-                    with gr.Row():
-                        msg = gr.Textbox(
-                            label="Your message",
-                            placeholder="Type your message here...",
-                            lines=2,
-                            scale=4
-                        )
-                        submit = gr.Button("Send", variant="primary", scale=1)
-
-                with gr.Column(scale=1):
-                    # Quick Actions Panel
-                    gr.Markdown("### 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")
-                    
-                    # Conversation Controls
-                    gr.Markdown("### Controls")
-                    clear = gr.Button("🗑️ Clear Chat")
-                    
-                    # Feedback Section
-                    gr.Markdown("### Feedback")
-                    feedback = gr.Radio(
-                        choices=["👍", "👎"],
-                        label="Was this response helpful?",
-                        visible=True
-                    )
-                    feedback_text = gr.Textbox(
-                        label="Additional comments (optional)",
-                        placeholder="Tell us more about your experience...",
-                        lines=2
-                    )
-
-            # Examples Section
-            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
-                )
-
-            # Information Accordions
-            with gr.Accordion("NHS Services Guide", open=False):
-                gr.Markdown("""
-                    ### Emergency Services (999)
-                    - Life-threatening emergencies
-                    - Severe injuries
-                    - Suspected heart attack or stroke
-                    
-                    ### NHS 111
-                    - Urgent but non-emergency situations
-                    - Medical advice needed
-                    - Unsure where to go
-                    
-                    ### GP Services
-                    - Routine check-ups
-                    - Non-urgent medical issues
-                    - Prescription renewals
-                """)
-
             # Event Handlers
             submit.click(
-                chat,
+                fn=chat,
                 inputs=[msg, chatbot],
                 outputs=[chatbot]
             ).then(
@@ -841,7 +266,7 @@ def create_demo():
             )
             
             msg.submit(
-                chat,
+                fn=chat,
                 inputs=[msg, chatbot],
                 outputs=[chatbot]
             ).then(
@@ -850,63 +275,56 @@ def create_demo():
                 msg
             )
             
-            clear.click(
-                lambda: [[], ""],
-                None,
-                [chatbot, msg]
-            )
-            
             feedback.change(
-                process_feedback,
+                fn=process_feedback,
                 inputs=[feedback, chatbot, feedback_text],
                 outputs=[]
             )
-            
-            # 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])
+
+            # Clear Chat Handler
+            clear = gr.Button("🗑️ Clear Chat")
+            clear.click(lambda: [[], ""], None, [chatbot, msg])
+
+            # Additional Information Sections
+            gr.HTML("""
+                <div style="padding: 20px;">
+                    <h2>Quick Actions</h2>
+                    <button onclick="document.getElementById('chatbot').value += 'Emergency Info: For emergencies, call 999.'">Emergency Info</button>
+                    <button onclick="document.getElementById('chatbot').value += 'NHS 111 Info: For urgent but non-emergency situations, call 111.'">NHS 111 Info</button>
+                    <button onclick="document.getElementById('chatbot').value += 'GP Booking Info: For routine appointments with your GP.'">GP Booking</button>
+                </div>
+            """)
+
+            gr.Markdown("### Example Messages")
+            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.Markdown("""
+                ### NHS Services Guide
+                **999 - Emergency Services**: For life-threatening emergencies, severe injuries, heart attack, stroke.
+                
+                **NHS 111**: Available 24/7 for urgent but non-life-threatening situations, medical advice, and guidance.
+                
+                **GP Services**: Routine check-ups, non-urgent medical issues, and prescription renewals.
+            """)
 
         return demo
-        
+
     except Exception as e:
         logger.error(f"Error creating demo: {e}")
         raise
 
 if __name__ == "__main__":
-    # Initialize environment
-    load_dotenv()
-    
-    # Set up HuggingFace login if token exists
-    hf_token = os.getenv("HF_TOKEN")
-    if hf_token:
-        login(token=hf_token)
-    
-    # Launch demo
+    # Launch Gradio Interface
     demo = create_demo()
-    demo.launch()
+    demo.launch()
+
+