Update app.py

app.py

@@ -1,5 +1,6 @@
 # ✅ Optimized Triage Chatbot Code for Hugging Face Space (NVIDIA T4 GPU)
 # Covers: Memory optimizations, 4-bit quantization, lazy loading, FAISS caching, faster inference, safe Gradio UI
+# Includes: Proper Gradio history handling, response cleaning, safety checks
 
 import os
 import time
@@ -10,10 +11,8 @@ import psutil
 from datetime import datetime
 from huggingface_hub import login
 from dotenv import load_dotenv
-from datasets import load_dataset, load_from_disk, Dataset
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM,
-    TrainingArguments, Trainer,
     BitsAndBytesConfig
 )
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
@@ -33,10 +32,7 @@ class SecretsManager:
     def setup():
         load_dotenv()
         creds = {
-            'KAGGLE_USERNAME': os.getenv('KAGGLE_USERNAME'),
-            'KAGGLE_KEY': os.getenv('KAGGLE_KEY'),
             'HF_TOKEN': os.getenv('HF_TOKEN'),
-            'WANDB_KEY': os.getenv('WANDB_KEY')
         }
         if creds['HF_TOKEN']:
             login(token=creds['HF_TOKEN'])
@@ -44,9 +40,9 @@ class SecretsManager:
         return creds
 
 # ===========================
-# 🧠 CHATBOT CLASS
+# 🧠 MEDICAL CHATBOT CORE
 # ===========================
-class PearlyBot:
+class MedicalTriageBot:
     def __init__(self):
         self.tokenizer = None
         self.model = None
@@ -57,196 +53,232 @@ class PearlyBot:
         self.num_relevant_chunks = 3
         self.last_interaction_time = time.time()
         self.interaction_cooldown = 1.0
+        self.safety_phrases = [
+            "999", "111", "emergency", "GP", "NHS",
+            "consult a doctor", "seek medical attention"
+        ]
 
-    def setup_model_and_tokenizer(self, model_name="google/gemma-7b"):
+    def setup_model(self, model_name="google/gemma-7b-it"):
         if self.model is not None:
             return
-        logger.info("🚀 Loading model & tokenizer")
+        
+        logger.info("🚀 Initializing medical AI model")
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_use_double_quant=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=torch.float16
         )
+        
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.tokenizer.pad_token = self.tokenizer.eos_token
-        model = AutoModelForCausalLM.from_pretrained(
+        
+        base_model = AutoModelForCausalLM.from_pretrained(
             model_name,
             device_map="auto",
             quantization_config=bnb_config,
-            torch_dtype=torch.float16,
-            low_cpu_mem_usage=True
+            torch_dtype=torch.float16
         )
-        model = prepare_model_for_kbit_training(model)
-        lora_config = LoraConfig(
-            r=4,
-            lora_alpha=16,
+        
+        peft_config = LoraConfig(
+            r=8,
+            lora_alpha=32,
             target_modules=["q_proj", "v_proj"],
             lora_dropout=0.05,
-            bias="none",
             task_type="CAUSAL_LM"
         )
-        self.model = get_peft_model(model, lora_config)
-        self.model.to("cuda" if torch.cuda.is_available() else "cpu")
-        logger.info("✅ Model & tokenizer ready")
+        
+        self.model = get_peft_model(base_model, peft_config)
+        logger.info("✅ Medical AI model ready")
 
-    def setup_embeddings(self):
-        if self.embeddings is None:
-            logger.info("📌 Loading sentence-transformer embeddings")
-            self.embeddings = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-MiniLM-L6-v2",
-                cache_folder="./embeddings_cache"
-            )
-
-    def load_faiss_index(self):
-        logger.info("📁 Loading FAISS index")
-        if os.path.exists("index_store/index.faiss"):
-            self.vector_store = FAISS.load_local("index_store", self.embeddings)
-        else:
-            self.build_faiss_index()
-
-    def build_faiss_index(self):
-        logger.info("🔧 Building FAISS index from knowledge base")
-        knowledge_base = self._load_knowledge_base()
-        self.setup_embeddings()
-        texts = self._split_texts(knowledge_base)
-        self.vector_store = FAISS.from_texts(
-            texts,
+    def setup_rag_system(self):
+        logger.info("📚 Initializing medical knowledge base")
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name="sentence-transformers/all-MiniLM-L6-v2",
+            model_kwargs={"device": "cpu"}
+        )
+        
+        if not os.path.exists("medical_index/index.faiss"):
+            self.build_medical_index()
+            
+        self.vector_store = FAISS.load_local(
+            "medical_index",
             self.embeddings,
-            metadatas=[{"source": f"chunk_{i}"} for i in range(len(texts))]
+            allow_dangerous_deserialization=True
         )
-        self.vector_store.save_local("index_store")
 
-    def _load_knowledge_base(self):
-        kb_dir = "knowledge_base"
-        os.makedirs(kb_dir, exist_ok=True)
-        kb_files = {
-            "triage.txt": "Severe chest pain? Call 999. Persistent cough? Book GP.",
-            "emergency.txt": "Unconscious? 999. Breathing issues? 999.",
-            "cultural.txt": "Respect prayer times, language needs, traditional remedies.",
-            "gp_booking.txt": "I need to book a GP for routine care next week."
+    def build_medical_index(self):
+        medical_knowledge = {
+            "emergency_protocols.txt": """Emergency Protocols:
+            - Chest pain: Call 999 immediately
+            - Breathing difficulties: Urgent 999 call
+            - Severe bleeding: Apply pressure, call 999""",
+            
+            "triage_guidelines.txt": """Triage Guidelines:
+            - Persistent fever >48h: Contact 111
+            - Minor injuries: Visit urgent care
+            - Medication questions: Consult GP"""
         }
-        for file, content in kb_files.items():
-            with open(os.path.join(kb_dir, file), 'w') as f:
+        
+        os.makedirs("medical_knowledge", exist_ok=True)
+        for filename, content in medical_knowledge.items():
+            with open(f"medical_knowledge/{filename}", "w") as f:
                 f.write(content)
-        return kb_files
-
-    def _split_texts(self, kb):
-        splitter = RecursiveCharacterTextSplitter(
+        
+        text_splitter = RecursiveCharacterTextSplitter(
             chunk_size=self.chunk_size,
-            chunk_overlap=self.chunk_overlap,
-            length_function=len,
-            add_start_index=True
+            chunk_overlap=self.chunk_overlap
         )
-        texts = []
-        for text in kb.values():
-            chunks = splitter.split_text(text)
-            texts.extend(chunks)
-        return texts
+        
+        documents = []
+        for text in medical_knowledge.values():
+            documents.extend(text_splitter.split_text(text))
+        
+        vector_store = FAISS.from_texts(
+            documents,
+            self.embeddings,
+            metadatas=[{"source": f"doc_{i}"} for i in range(len(documents))]
+        )
+        vector_store.save_local("medical_index")
 
-    def _get_enhanced_context(self, query):
+    def get_medical_context(self, query):
         try:
-            results = self.vector_store.similarity_search_with_score(query, k=self.num_relevant_chunks)
-            context = [f"[Source: {doc.metadata.get('source', 'unknown')}]:\n{doc.page_content}"
-                       for doc, score in results if score < 0.8]
-            return "\n\n".join(context)
+            docs = self.vector_store.similarity_search(query, k=2)
+            return "\n".join([d.page_content for d in docs])
         except Exception as e:
             logger.error(f"Context error: {e}")
             return ""
 
     @torch.inference_mode()
-    def generate_response(self, message, history):
+    def generate_safe_response(self, message, history):
         try:
-            # Throttle
+            # Rate limiting
             if time.time() - self.last_interaction_time < self.interaction_cooldown:
                 time.sleep(self.interaction_cooldown)
-
-            self.setup_model_and_tokenizer()
-            self.setup_embeddings()
-            self.load_faiss_index()
-
-            context = self._get_enhanced_context(message)
-            conv_history = "\n".join([
-                f"User: {turn['content']}" if turn['role'] == 'user' else f"Assistant: {turn['content']}"
-                for turn in history[-3:]
-            ])
-
+            
+            # Convert Gradio history to conversational format
+            conversation = "\n".join(
+                [f"User: {user}\nAssistant: {bot}" for user, bot in history[-3:]]
+            )
+            
+            # Get medical context
+            context = self.get_medical_context(message)
+            
+            # Create safety-focused prompt
             prompt = f"""<start_of_turn>system
-Context:
+You are a medical triage assistant. Use this context:
 {context}
-Conversation:
-{conv_history}
+Current conversation:
+{conversation}
+Guidelines:
+1. Assess symptom severity
+2. Recommend appropriate care level
+3. Never diagnose or prescribe
+4. Always include safety netting
 <end_of_turn>
 <start_of_turn>user
 {message}
 <end_of_turn>
 <start_of_turn>assistant"""
-
-            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(self.model.device)
+            
+            # Generate response
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=1024
+            ).to(self.model.device)
+            
             outputs = self.model.generate(
                 **inputs,
-                max_new_tokens=128,
-                min_new_tokens=20,
-                do_sample=True,
+                max_new_tokens=256,
                 temperature=0.7,
                 top_p=0.9,
-                repetition_penalty=1.2,
-                no_repeat_ngram_size=3
+                repetition_penalty=1.2
             )
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            response = response.split("<start_of_turn>assistant")[-1].strip().split("<end_of_turn>")[0].strip()
+            
+            # Clean and validate response
+            full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            clean_response = full_response.split("<start_of_turn>assistant")[-1]
+            clean_response = clean_response.split("<end_of_turn>")[0].strip()
+            
+            # Ensure medical safety
+            if not any(phrase in clean_response.lower() for phrase in self.safety_phrases):
+                clean_response += "\n\nIf symptoms persist, please contact NHS 111."
+            
             self.last_interaction_time = time.time()
-            logger.info(f"💬 Memory use: {psutil.virtual_memory().percent}%")
-            return response
-
+            return clean_response[:500]  # Limit response length
+            
         except Exception as e:
             logger.error(f"Generation error: {e}")
-            return "I encountered an error. Please try again."
+            return "Please contact NHS 111 directly for urgent medical advice."
 
 # ===========================
-# 💬 GRADIO UI
+# 💬 SAFE GRADIO INTERFACE
 # ===========================
-def create_demo():
-    bot = PearlyBot()
-
-    def chat(message, history):
-        if not message.strip():
-            return history
-    
-        # Convert Gradio-style history to model-style format
-        structured_history = []
-        for user_msg, bot_msg in history:
-            structured_history.append({"role": "user", "content": user_msg})
-            structured_history.append({"role": "assistant", "content": bot_msg})
-    
-        # Generate bot response
-        response = bot.generate_response(message, structured_history)
-    
-        # Append new message pair
-        history.append([message, response])
-        return history
+def create_medical_interface():
+    bot = MedicalTriageBot()
+    bot.setup_model()
+    bot.setup_rag_system()
 
+    def handle_conversation(message, history):
+        try:
+            response = bot.generate_safe_response(message, history)
+            return history + [(message, response)]
+        except Exception as e:
+            logger.error(f"Conversation error: {e}")
+            return history + [(message, "System error - please refresh the page")]
 
-    with gr.Blocks() as demo:
-        chatbot = gr.Chatbot(
-    value=[["Hello!", "Hi, I’m Pearly, your GP triage assistant. I’m here to help you assess your symptoms and guide you to the right care. How are you feeling today?"]],
-    height=500,
-    show_label=False
-)
-        msg = gr.Textbox(label="Type your message")
-        send = gr.Button("Send")
-        clear = gr.Button("Clear Chat")
-
-        msg.submit(chat, [msg, chatbot], [chatbot]).then(lambda: gr.update(value=""), None, [msg])
-        send.click(chat, [msg, chatbot], [chatbot]).then(lambda: gr.update(value=""), None, [msg])
-        clear.click(lambda: [], None, chatbot)
+    with gr.Blocks(theme=gr.themes.Soft()) as interface:
+        gr.Markdown("# NHS Triage Assistant")
+        gr.HTML("""<div class="emergency-banner">🚨 In emergencies, always call 999 immediately</div>""")
+        
+        with gr.Row():
+            chatbot = gr.Chatbot(
+                value=[("", "Hello! I'm your NHS digital assistant. How can I help you today?")],
+                height=500,
+                label="Medical Triage Chat"
+            )
+            
+        with gr.Row():
+            message_input = gr.Textbox(
+                placeholder="Describe your symptoms...",
+                label="Your Message",
+                max_lines=3
+            )
+            submit_btn = gr.Button("Send", variant="primary")
+            
+        clear_btn = gr.Button("Clear History")
+        
+        # Event handlers
+        message_input.submit(
+            handle_conversation,
+            [message_input, chatbot],
+            [chatbot]
+        ).then(lambda: "", None, [message_input])
+        
+        submit_btn.click(
+            handle_conversation,
+            [message_input, chatbot],
+            [chatbot]
+        ).then(lambda: "", None, [message_input])
+        
+        clear_btn.click(
+            lambda: [("", "Hello! I'm your NHS digital assistant. How can I help you today?")],
+            None,
+            [chatbot]
+        )
 
-    return demo
+    return interface
 
 # ===========================
-# 🚀 MAIN
+# 🚀 LAUNCH APPLICATION
 # ===========================
 if __name__ == "__main__":
     SecretsManager.setup()
-    demo = create_demo()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    medical_app = create_medical_interface()
+    medical_app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )