app.py

"""
Phase 5: Gradio Chatbot Interface — Professional Edition (Hugging Face Version).

Launches an interactive web UI where a patient can type symptoms
and receive an empathetic, medically-informed response.
"""
import os
import gradio as gr
import re
import json
import torch
import gc
from huggingface_hub import hf_hub_download

import sys
sys.path.insert(0, os.path.dirname(__file__))
from config import EMOTION_LABELS, RAG_TOP_K
from model import EmotionAwareMedicalChatbot
from rag_engine import RAGEngine

# ============================================================
# 1. DOWNLOAD MODEL & INDEX FROM HUB
# ============================================================
REPO_ID = "shreenikethjoshi/MedEmotion-BioGPT-Large"

print("Downloading SOTA artifacts from Hub...")
checkpoint_path = hf_hub_download(repo_id=REPO_ID, filename="best_model_inference.pt")
index_path = hf_hub_download(repo_id=REPO_ID, filename="faiss_index.bin")
passages_path = hf_hub_download(repo_id=REPO_ID, filename="faiss_index_passages.npy")


class ChatbotInference:
    """Production-grade inference wrapper for the Emotion-Aware Medical Chatbot."""

    def __init__(self, ckpt_path: str, idx_path: str, pass_path: str):
        print("Loading model...")
        self.model = EmotionAwareMedicalChatbot()
        ckpt = torch.load(ckpt_path, map_location="cpu")
        self.model.load_state_dict(ckpt["model_state_dict"])
        self.model.to("cpu").eval()
        
        del ckpt
        gc.collect()
        print("  ✅ Model loaded")

        print("Loading RAG engine...")
        self.rag = RAGEngine(index_path=idx_path, passages_path=pass_path)
        try:
            self.rag.load_index()
            self.use_rag = True
            print("  ✅ RAG index loaded")
        except Exception as e:
            self.use_rag = False
            print(f"  ⚠️  No RAG index found — running without retrieval ({str(e)})")

    # ==========================================================
    # Intent Router — handles non-medical inputs without GPU
    # ==========================================================
    INTENTS = {
        "greeting": {
            "patterns": [
                "hi", "hello", "hey", "good morning", "good evening",
                "good afternoon", "howdy", "greetings", "sup", "yo",
                "hii", "hiii", "hiiii", "helo", "namaste", "hola",
            ],
            "response": (
                "Hello! 👋 I'm your **Emotion-Aware Medical Assistant**.\n\n"
                "I can help you with:\n"
                "• 🩺 Symptom analysis & medical advice\n"
                "• 🎭 Emotional support & empathetic responses\n"
                "• 💊 Medication information\n\n"
                "**How can I help you today?** Please describe your symptoms "
                "or ask a health-related question."
            ),
        },
        "thanks": {
            "patterns": [
                "thank", "thanks", "thank you", "thx", "appreciate",
                "thanks a lot", "thank you so much", "great help",
            ],
            "response": (
                "You're welcome! 😊 I'm glad I could help.\n\n"
                "**Important reminders:**\n"
                "• Always consult a qualified doctor for proper diagnosis\n"
                "• Follow prescribed medications as directed\n"
                "• Don't hesitate to seek emergency care if symptoms worsen\n\n"
                "Take care of your health! 🏥💚"
            ),
        },
        "farewell": {
            "patterns": [
                "bye", "goodbye", "see you", "take care", "good night",
                "see ya", "gotta go", "leaving", "exit", "quit",
            ],
            "response": (
                "Goodbye! Take care of yourself 💚\n\n"
                "**Remember:**\n"
                "• Stay hydrated and get adequate rest\n"
                "• Follow up with your doctor if any symptoms persist\n"
                "• I'm available anytime you need medical guidance\n\n"
                "Stay healthy! 🌟"
            ),
        },
        "about": {
            "patterns": [
                "who are you", "what are you", "what can you do",
                "help", "about", "features", "capabilities",
                "how do you work", "tell me about yourself",
            ],
            "response": (
                "🏥 **Emotion-Aware Medical Chatbot**\n\n"
                "I'm an AI-powered healthcare assistant built with a SOTA architecture:\n\n"
                "**🧠 How I work:**\n"
                "1. **Longformer Encoder** — Reads your full message (up to 4096 tokens)\n"
                "2. **Syntax GCN** — Understands clinical grammar & relationships\n"
                "3. **Emotion Detector** — Identifies your emotional state (fear, sadness, etc.)\n"
                "4. **RAG Engine** — Retrieves relevant info from 180k+ real doctor conversations\n"
                "5. **BioGPT Decoder** — Generates an empathetic, medically-informed response\n\n"
                "⚠️ **Disclaimer:** I'm a research prototype. Always consult a real doctor for medical decisions."
            ),
        },
        "how_are_you": {
            "patterns": [
                "how are you", "how r u", "how do you do",
                "whats up", "what's up", "wassup",
            ],
            "response": (
                "I'm doing great, thank you for asking! 😊\n\n"
                "I'm here and ready to help with any health concerns you may have. "
                "Please feel free to describe your symptoms or ask a medical question."
            ),
        },
        "ok_ack": {
            "patterns": [
                "ok", "okay", "fine", "alright", "got it",
                "understood", "i see", "hmm", "ohh", "oh",
            ],
            "response": (
                "Got it! 👍 Is there anything else you'd like to know about your health? "
                "Feel free to ask me any medical question."
            ),
        },
    }

    EMERGENCY_KEYWORDS = [
        "suicide", "kill myself", "want to die", "end my life",
        "self harm", "self-harm", "cutting myself",
        "heart attack", "chest pain and can't breathe",
        "can't breathe", "choking", "unconscious",
        "overdose", "poisoning", "severe bleeding",
    ]

    EMERGENCY_RESPONSE = (
        "🚨 **EMERGENCY — PLEASE SEEK IMMEDIATE HELP** 🚨\n\n"
        "If you or someone is in immediate danger:\n\n"
        "📞 **Emergency Services:** Call **112** (India) or **911** (US)\n"
        "📞 **Suicide Prevention Helpline (India):** 9820466726 (iCall)\n"
        "📞 **Mental Health Helpline:** 1800-599-0019 (KIRAN)\n"
        "📞 **National Suicide Prevention (US):** 988\n\n"
        "⚠️ **I am an AI and cannot provide emergency medical care.** "
        "Please reach out to a medical professional or emergency services immediately.\n\n"
        "**You are not alone. Help is available. 💚**"
    )

    MEDICAL_KEYWORDS = [
        "pain", "hurt", "ache", "fever", "cough", "cold", "sick",
        "blood", "symptom", "disease", "medicine", "doctor", "treatment",
        "diagnosed", "surgery", "infection", "allergy", "headache",
        "stomach", "vomit", "nausea", "diarrhea", "rash", "swelling",
        "dizzy", "fatigue", "weakness", "breathing", "chest", "heart",
        "diabetes", "pressure", "cancer", "thyroid", "kidney", "liver",
        "pregnant", "pregnancy", "period", "menstrual", "anxiety",
        "depression", "insomnia", "sleep", "weight", "diet", "vitamin",
        "tablet", "capsule", "dose", "mg", "prescription", "test",
        "report", "scan", "xray", "mri", "ct scan", "biopsy",
        "suffering", "problem", "condition", "disorder", "chronic",
        "acute", "remedy", "cure", "medication", "drug", "antibiotic",
    ]

    def _detect_intent(self, message):
        """Route non-medical messages to predefined responses."""
        msg = message.strip().lower()
        msg_clean = re.sub(r'[!.,?;:\'"]+$', '', msg)

        # 1. Emergency check (highest priority)
        for keyword in self.EMERGENCY_KEYWORDS:
            if keyword in msg:
                return self.EMERGENCY_RESPONSE

        # 2. Check defined intents
        for intent_data in self.INTENTS.values():
            for pattern in intent_data["patterns"]:
                if (msg_clean == pattern or
                    msg_clean.startswith(pattern + " ") or
                    msg_clean.endswith(" " + pattern)):
                    return intent_data["response"]

        # 3. Empty or gibberish detection
        if len(msg.strip()) < 2:
            return self.INTENTS["greeting"]["response"]

        # Only alphabetic characters check (no real words)
        alpha_only = re.sub(r'[^a-zA-Z\s]', '', msg)
        if len(alpha_only.split()) == 0:
            return ("I couldn't understand your message. 🤔\n\n"
                    "Could you please describe your symptoms or health concern in more detail? "
                    "For example: *\"I have a headache and fever for 2 days\"*")

        # 4. Short non-medical message → friendly prompt
        if (len(msg.split()) <= 2 and
            not any(w in msg for w in self.MEDICAL_KEYWORDS)):
            return self.INTENTS["greeting"]["response"]

        return None  # It's a medical query → use the AI model

    # ==========================================================
    # Main response generation
    # ==========================================================
    @torch.no_grad()
    def respond(self, patient_message: str, max_new_tokens: int = 150):
        """
        Generate a response using the FULL trained pipeline.
        Returns: (response_text, emotion_detected, emotion_scores, rag_passages)
        """
        # 0. Intent routing (instant, no GPU)
        intent_response = self._detect_intent(patient_message)
        if intent_response:
            return intent_response, "neutral", {}, []

        # 1. Emotion detection
        emotion_probs = self.model.get_emotion_embedding([patient_message])
        probs = emotion_probs[0].cpu().tolist()
        emotion_scores = {
            EMOTION_LABELS[i]: round(probs[i], 4)
            for i in range(len(EMOTION_LABELS))
        }
        emotion_detected = max(emotion_scores, key=emotion_scores.get)

        # 2. RAG retrieval
        rag_passages = []
        if getattr(self, 'use_rag', False):
            results = self.rag.retrieve(patient_message, top_k=RAG_TOP_K)
            rag_passages = [r["passage"] for r in results]

        # 3. Generate via full pipeline (Longformer → GCN → Fusion → Prefix → BioGPT)
        dep_edges = json.dumps([])
        generated_texts, _ = self.model.generate_with_context(
            patient_texts=[patient_message],
            dep_edges_json=[dep_edges],
            max_new_tokens=280,
            temperature=0.45,
            top_p=0.9,
            do_sample=True,
        )

        response_text = generated_texts[0] if generated_texts else ""

        # 4. Post-process: clean hallucinations and data leaks
        response_text = self._clean_response(response_text)

        # 5. Safety check: if response is empty after cleaning
        if not response_text or len(response_text.strip()) < 10:
            response_text = (
                "Based on your symptoms, I recommend consulting a healthcare professional "
                "for a proper diagnosis. In the meantime, stay hydrated and get adequate rest."
            )

        return response_text, emotion_detected, emotion_scores, rag_passages

    @staticmethod
    def _clean_response(text):
        if not text: return ""
        
        # 1. Brutally remove any URLs or links
        text = re.sub(r'https?://\S+|www\.\S+', '', text, flags=re.IGNORECASE)
        text = re.sub(r'urldefense\S+', '', text, flags=re.IGNORECASE)
        
        # 2. Remove annoying forum phrases like "ask a direct question on this forum"
        text = re.sub(r'(Let me know if you have|Feel free to ask|You can ask a direct question|Please meet an oral).*$', '', text, flags=re.IGNORECASE | re.DOTALL)
        
        # 3. Safely remove doctor signatures WITHOUT breaking the word "doctor"
        text = re.sub(r'(Regards|Sincerely|Best wishes|Warm regards|Thanks|Wishing),?\s*(Dr\.?|Doctor)?\s*[A-Z][a-z]+.*$', '', text, flags=re.IGNORECASE | re.DOTALL)
        
        # Split into sentences to remove leaked patient history
        sentences = re.split(r'(?<=[.!?])\s+', text.strip())
        
        if len(sentences) > 1:
            cleaned = []
            for sent in sentences:
                if re.search(r'\b\d{1,2}\s*(year|yr)\s*old\b', sent, re.IGNORECASE): continue
                if re.search(r'\b(visit(ing)?\s+her|visit(ing)?\s+him|his\s+wife|her\s+husband)\b', sent, re.IGNORECASE): continue
                if re.search(r'^(I am|I have been|My name|Dear doctor)', sent, re.IGNORECASE): continue
                if 'healthcaremagic' in sent.lower() or 'urldefense' in sent.lower(): continue
                cleaned.append(sent)
            text = ' '.join(cleaned)
            
        if len(text) > 600:
            last_period = text[:600].rfind('.')
            if last_period > 100: text = text[:last_period + 1]

        # 4. Prevent mid-sentence cut-offs
        if text and text[-1] not in '.!?':
            last_punct = max(text.rfind('.'), text.rfind('!'), text.rfind('?'))
            if last_punct > 0: text = text[:last_punct + 1]
            else: text += "."
                
        # 5. Clean up any weird leftover punctuation like "_ _;!!"
        text = re.sub(r'[_;\-\|]{2,}.*$', '', text)
                
        return text.strip()


# Load the instance
chatbot = ChatbotInference(checkpoint_path, index_path, passages_path)

# ==============================================================
# Gradio UI — Professional Medical Chatbot Interface
# ==============================================================
def launch_gradio(chatbot: ChatbotInference):
    """Launch a professional Gradio web interface."""

    def chat_fn(message, history):
        """Process user message and return formatted response."""
        if not message or not message.strip():
            return "Please type a message to get started! 😊"

        response, emotion, scores, passages = chatbot.respond(message)

        if not scores:
            return response

        emotion_bar = " | ".join([
            f"{k}: {v:.0%}"
            for k, v in sorted(scores.items(), key=lambda x: -x[1])[:3]
        ])

        output = f"{response}\n\n"
        output += f"---\n"
        output += f"🎭 **Detected Emotion:** {emotion.upper()}\n"
        output += f"📊 {emotion_bar}\n"

        if passages:
            output += f"\n📚 **Retrieved Medical Context:**\n"
            for i, p in enumerate(passages[:3], 1):
                output += f"  {i}. {p}\n\n"

        return output

    # Build the interface
    demo = gr.ChatInterface(
        fn=chat_fn,
        title="🏥 Emotion-Aware Medical Chatbot",
        description=(
            "An AI-powered healthcare assistant that **detects your emotional state** "
            "and provides **empathetic, medically-informed responses** powered by "
            "Longformer, GCN, and BioGPT.\n\n"
            "⚠️ **Disclaimer:** This is a research prototype. Always consult a qualified doctor "
            "for medical diagnosis and treatment."
        ),
        examples=[
            "I have been having severe chest pain for the last 3 days and I'm really scared",
            "My child has a high fever of 103°F and won't stop crying. What should I do?",
            "I've been feeling very depressed lately and can't sleep at night",
            "What are the side effects of metformin for diabetes?",
            "I have persistent headache with blurred vision since morning",
            "I'm 6 months pregnant and experiencing severe back pain",
        ],
    )

    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)

if __name__ == "__main__":
    launch_gradio(chatbot)