Update app.py

app.py

@@ -1,8 +1,12 @@
 import os
 import gc
+import re
+import time
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
+from collections import defaultdict
+from datetime import datetime, timedelta
 
 # =============================
 # Configuration
@@ -12,22 +16,46 @@ MAX_NEW_TOKENS = 200
 TEMPERATURE = 0.5
 TOP_K = 50
 REPETITION_PENALTY = 1.1
+MAX_HISTORY_TURNS = 5  # Limit conversation history
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"🚀 Loading model from {MODEL_PATH} on {device}...")
 
+# =============================
+# Rate Limiting (Simple IP-based)
+# =============================
+rate_limit_store = defaultdict(list)
+MAX_REQUESTS_PER_MINUTE = 10
+
+def check_rate_limit(session_id):
+    """Simple rate limiting to prevent abuse"""
+    now = datetime.now()
+    rate_limit_store[session_id] = [
+        timestamp for timestamp in rate_limit_store[session_id]
+        if now - timestamp < timedelta(minutes=1)
+    ]
+    
+    if len(rate_limit_store[session_id]) >= MAX_REQUESTS_PER_MINUTE:
+        return False
+    
+    rate_limit_store[session_id].append(now)
+    return True
+
 # ==========================
 # Load Model & Tokenizer
 # =============================
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_PATH,
-    device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    low_cpu_mem_usage=True
-)
-
-print("✅ ChatDoctor model loaded successfully!\n")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        device_map="auto",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        low_cpu_mem_usage=True
+    )
+    print("✅ ChatDoctor model loaded successfully!\n")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    raise
 
 # =============================
 # Stop Criteria
@@ -57,43 +85,113 @@ MEDICAL_KEYWORDS = [
     "stomach", "head", "chest", "throat", "heart", "lung", "liver", "kidney", "brain",
     "doctor", "hospital", "medicine", "treatment", "therapy", "surgery", "disease",
     "illness", "blood", "test", "scan", "health", "diet", "nutrition", "stress", "sleep",
-    "weight", "vitamin", "fatigue", "anxiety", "depression"
+    "weight", "vitamin", "fatigue", "anxiety", "depression", "nausea", "dizziness",
+    "rash", "swelling", "injury", "bruise", "cold", "sneeze", "tired", "weak"
 ]
 
-CASUAL_ONLY_PATTERNS = [
-    "hey", "hi", "hello", "sup", "yo", "good morning", "good evening",
-    "how are you", "wassup", "hiya"
+# Emergency keywords that should trigger immediate medical attention warning
+EMERGENCY_KEYWORDS = [
+    "suicide", "kill myself", "end my life", "chest pain", "can't breathe",
+    "severe bleeding", "overdose", "poisoning", "unconscious", "seizure",
+    "stroke", "heart attack", "choking"
 ]
 
+CASUAL_PATTERNS = [
+    r"^(hey|hi|hello|sup|yo|wassup|hiya)\s*[\?\!\.]*$",
+    r"^good\s+(morning|evening|afternoon|night)\s*[\?\!\.]*$",
+    r"^how\s+are\s+you\s*[\?\!\.]*$",
+    r"^what'?s\s+up\s*[\?\!\.]*$",
+]
+
+
+def is_emergency_query(message):
+    """Detect if query contains emergency keywords"""
+    message_lower = message.lower()
+    return any(keyword in message_lower for keyword in EMERGENCY_KEYWORDS)
+
 
 def is_medical_query(message):
+    """Enhanced medical query detection"""
     message_lower = message.lower()
+    
+    # Check for medical keywords
     for keyword in MEDICAL_KEYWORDS:
         if keyword in message_lower:
             return True
-    question_words = ["what", "how", "why", "when", "where", "can", "should", "is", "are", "do", "does"]
-    has_question = any(q in message_lower.split()[:3] for q in question_words)
-    if has_question and len(message.split()) > 5:
+    
+    # Check for question patterns with sufficient length
+    question_words = ["what", "how", "why", "when", "where", "can", "should", "is", "are", "do", "does", "could", "would"]
+    words = message_lower.split()
+    has_question = any(q in words[:4] for q in question_words)
+    
+    if has_question and len(words) > 5:
         return True
+    
     return False
 
 
 def is_only_greeting(message):
-    message_lower = message.lower().strip().replace("!", "").replace("?", "").replace(".", "")
-    if len(message_lower.split()) <= 3:
-        for pattern in CASUAL_ONLY_PATTERNS:
-            if message_lower == pattern or message_lower.startswith(pattern):
-                return True
+    """Improved greeting detection using regex"""
+    message_clean = message.lower().strip()
+    
+    # Remove punctuation for matching
+    message_clean = re.sub(r'[!?.]+$', '', message_clean)
+    
+    # Check if it matches any casual pattern
+    for pattern in CASUAL_PATTERNS:
+        if re.match(pattern, message_clean):
+            return True
+    
+    return False
+
+
+# =============================
+# Safety Filter
+# =============================
+DANGEROUS_PATTERNS = [
+    r"take\s+\d+\s+(pills|tablets|capsules)",
+    r"inject\s+(yourself|myself)",
+    r"(don't|do not)\s+go\s+to\s+(hospital|doctor|emergency)",
+    r"ignore\s+(doctor|medical|professional)",
+]
+
+def contains_dangerous_advice(response):
+    """Check if response contains potentially dangerous medical advice"""
+    response_lower = response.lower()
+    
+    for pattern in DANGEROUS_PATTERNS:
+        if re.search(pattern, response_lower):
+            return True
+    
     return False
 
 
 # =============================
 # Get Response
 # =============================
-def get_response(user_input, history_context):
+def get_response(user_input, history_context, session_id="default"):
+    """Generate response with enhanced safety and quality checks"""
+    
+    # Rate limiting check
+    if not check_rate_limit(session_id):
+        return "⏰ You've made too many requests. Please wait a minute before trying again."
+    
+    # Emergency detection
+    if is_emergency_query(user_input):
+        return (
+            "🚨 **EMERGENCY DETECTED** 🚨\n\n"
+            "If you are experiencing a medical emergency, please:\n"
+            "• Call emergency services immediately (911 in US, 999 in UK, 112 in EU)\n"
+            "• Go to the nearest emergency room\n"
+            "• Contact your local emergency hotline\n\n"
+            "This AI cannot provide emergency medical care. Please seek immediate professional help."
+        )
+    
+    # Greeting detection
     if is_only_greeting(user_input):
         return "👋 Hello! I'm ChatDoctor — your AI medical assistant. Please tell me about any health symptoms or medical concerns you'd like to discuss."
 
+    # Non-medical query handling
     if not is_medical_query(user_input):
         return (
             "Hello! I'm ChatDoctor, an AI medical assistant specialized in health and wellness.\n\n"
@@ -104,17 +202,21 @@ def get_response(user_input, history_context):
             "Please describe your health concern in detail to get started."
         )
 
+    # Build prompt with limited history
     human_prefix = "Patient:"
     doctor_prefix = "ChatDoctor:"
     system_instruction = (
         "You are ChatDoctor, a professional medical AI assistant. "
-        "You provide accurate, concise, and empathetic responses to health-related questions only.\n\n"
-        "If the question is non-medical, politely redirect back to medical topics.\n"
+        "You provide accurate, concise, and empathetic responses to health-related questions only.\n"
+        "Always recommend consulting a healthcare professional for serious conditions.\n"
+        "Never provide dosage instructions or tell patients to avoid seeking professional help.\n\n"
     )
 
-    # Build history
+    # Limit history to prevent token overflow
+    limited_history = history_context[-MAX_HISTORY_TURNS:] if len(history_context) > MAX_HISTORY_TURNS else history_context
+    
     history_text = [system_instruction]
-    for human, assistant in history_context:
+    for human, assistant in limited_history:
         if human:
             history_text.append(f"{human_prefix} {human}")
         if assistant:
@@ -122,45 +224,68 @@ def get_response(user_input, history_context):
     history_text.append(f"{human_prefix} {user_input}")
 
     prompt = "\n".join(history_text) + f"\n{doctor_prefix} "
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-
-    stop_words = ["Patient:", "\nPatient:", "Patient :", "\n\nPatient"]
-    stop_ids = [tokenizer.encode(word, add_special_tokens=False) for word in stop_words]
-    stopping_criteria = StoppingCriteriaList([StopOnTokens(stop_ids)])
-
-    with torch.no_grad():
-        output_ids = model.generate(
-            input_ids,
-            max_new_tokens=MAX_NEW_TOKENS,
-            do_sample=True,
-            temperature=TEMPERATURE,
-            top_k=TOP_K,
-            repetition_penalty=REPETITION_PENALTY,
-            stopping_criteria=stopping_criteria,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id
-        )
-
-    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)[len(prompt):].strip()
-
-    for stop_word in ["Patient:", "Patient :", "\nPatient", "Patient"]:
-        if stop_word in response:
-            response = response.split(stop_word)[0].strip()
-            break
-
-    response = response.strip()
-    if any(x in response.lower() for x in ["chatbot", "api key", "error", "cloud"]):
-        response = (
-            "I apologize for the confusion — I'm ChatDoctor, trained to assist with medical and health-related topics only. "
-            "Please tell me about your symptoms or health concerns."
-        )
-
-    del input_ids, output_ids
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
-    return response
+    
+    try:
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+        
+        # Stop words for cleaner output
+        stop_words = ["Patient:", "\nPatient:", "Patient :", "\n\nPatient"]
+        stop_ids = [tokenizer.encode(word, add_special_tokens=False) for word in stop_words]
+        stopping_criteria = StoppingCriteriaList([StopOnTokens(stop_ids)])
+
+        with torch.no_grad():
+            output_ids = model.generate(
+                input_ids,
+                max_new_tokens=MAX_NEW_TOKENS,
+                do_sample=True,
+                temperature=TEMPERATURE,
+                top_k=TOP_K,
+                repetition_penalty=REPETITION_PENALTY,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id
+            )
+
+        response = tokenizer.decode(output_ids[0], skip_special_tokens=True)[len(prompt):].strip()
+
+        # Clean up response
+        for stop_word in ["Patient:", "Patient :", "\nPatient", "Patient"]:
+            if stop_word in response:
+                response = response.split(stop_word)[0].strip()
+                break
+
+        response = response.strip()
+        
+        # Safety filter
+        if contains_dangerous_advice(response):
+            response = (
+                "I apologize, but I cannot provide that specific medical advice. "
+                "Please consult with a qualified healthcare professional who can properly evaluate your situation."
+            )
+        
+        # Filter out inappropriate content
+        if any(x in response.lower() for x in ["chatbot", "api key", "error", "cloud", "sorry, i don't have"]):
+            response = (
+                "I apologize for the confusion. I'm ChatDoctor, trained to assist with medical and health-related topics. "
+                "Please tell me more about your symptoms or health concerns so I can help you better."
+            )
+        
+        # Add disclaimer for serious conditions
+        serious_conditions = ["cancer", "tumor", "heart disease", "stroke", "diabetes complications"]
+        if any(condition in response.lower() for condition in serious_conditions):
+            response += "\n\n⚠️ **Important:** Please consult a healthcare professional for proper diagnosis and treatment."
+
+        # Clean up memory
+        del input_ids, output_ids
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return response
+        
+    except Exception as e:
+        print(f"Error generating response: {e}")
+        return "I apologize, but I encountered an error processing your request. Please try rephrasing your question or try again later."
 
 
 # =============================
@@ -171,46 +296,82 @@ custom_css = """
     text-align: center;
     background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
     color: white;
-    padding: 20px;
-    border-radius: 10px;
+    padding: 25px;
+    border-radius: 12px;
     margin-bottom: 20px;
+    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
 }
-#header h1 { margin: 0; font-size: 2.3em; }
-#header p { margin: 5px 0 0; font-size: 1em; opacity: 0.9; }
+#header h1 { margin: 0; font-size: 2.5em; font-weight: 700; }
+#header p { margin: 5px 0 0; font-size: 1.1em; opacity: 0.95; }
 .disclaimer {
     background-color: #fff3cd;
-    border: 1px solid #ffc107;
+    border-left: 4px solid #ffc107;
     border-radius: 8px;
-    padding: 15px;
+    padding: 18px;
     margin: 20px 0;
     color: #856404;
 }
+.disclaimer h3 { margin-top: 0; color: #d39e00; }
+.emergency-warning {
+    background-color: #f8d7da;
+    border-left: 4px solid #dc3545;
+    border-radius: 8px;
+    padding: 15px;
+    margin: 15px 0;
+    color: #721c24;
+}
+footer {
+    margin-top: 30px;
+    padding: 15px;
+    text-align: center;
+    color: #6c757d;
+    font-size: 0.9em;
+}
 """
 
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    session_state = gr.State(value=str(time.time()))  # Unique session ID
+    
     gr.HTML("""
         <div id="header">
             <h1>🩺 ChatDoctor AI Assistant</h1>
             <p>Your AI-powered medical consultation partner</p>
         </div>
     """)
+    
     gr.HTML("""
         <div class="disclaimer">
             <h3>⚠️ Medical Disclaimer</h3>
-            <p>This AI assistant is for informational purposes only. 
-            It is NOT a substitute for professional medical advice, diagnosis, or treatment.</p>
+            <p><strong>This AI assistant is for informational purposes only.</strong> 
+            It is NOT a substitute for professional medical advice, diagnosis, or treatment. 
+            Always seek the advice of your physician or qualified health provider with any questions 
+            you may have regarding a medical condition.</p>
+        </div>
+    """)
+    
+    gr.HTML("""
+        <div class="emergency-warning">
+            <h4>🚨 In Case of Emergency</h4>
+            <p>If you are experiencing a medical emergency, call emergency services immediately 
+            (911 in US, 999 in UK, 112 in EU) or go to the nearest emergency room.</p>
         </div>
     """)
 
     chatbot = gr.Chatbot(
-        height=480,
-        placeholder="<div style='text-align:center;padding:40px;'><h3>👋 Welcome to ChatDoctor!</h3><p>Describe your symptoms or ask a health-related question to begin.</p></div>",
+        height=500,
+        placeholder="<div style='text-align:center;padding:50px;'><h3>👋 Welcome to ChatDoctor!</h3><p style='color:#6c757d;'>Describe your symptoms or ask a health-related question to begin.</p><p style='color:#dc3545;margin-top:15px;'><strong>Remember:</strong> This is not a replacement for professional medical care.</p></div>",
         show_label=False,
         avatar_images=(None, "🤖"),
     )
 
     with gr.Row():
-        msg = gr.Textbox(placeholder="Type your medical concern here...", show_label=False, scale=9, container=False)
+        msg = gr.Textbox(
+            placeholder="Type your medical concern here... (e.g., 'I have a headache for 3 days')",
+            show_label=False,
+            scale=9,
+            container=False,
+            lines=1
+        )
         send_btn = gr.Button("Send 📤", scale=1, variant="primary")
 
     with gr.Row():
@@ -218,44 +379,62 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         retry_btn = gr.Button("🔄 Retry", scale=1)
 
     with gr.Accordion("⚙️ Advanced Settings", open=False):
-        temp_slider = gr.Slider(0.1, 1.0, TEMPERATURE, 0.1, label="Temperature")
+        temp_slider = gr.Slider(0.1, 1.0, TEMPERATURE, 0.1, label="Temperature (Lower = More Focused)")
         max_tok_slider = gr.Slider(50, 500, MAX_NEW_TOKENS, 50, label="Max Tokens")
-        top_k_slider = gr.Slider(1, 100, TOP_K, 1, label="Top-K")
+        top_k_slider = gr.Slider(1, 100, TOP_K, 1, label="Top-K Sampling")
 
     def user_message(user_msg, history):
+        if not user_msg.strip():
+            return "", history
         return "", history + [[user_msg, None]]
 
-    def bot_response(history, temp, max_tok, topk):
+    def bot_response(history, temp, max_tok, topk, session_id):
+        if not history or history[-1][1] is not None:
+            return history
+            
         global TEMPERATURE, MAX_NEW_TOKENS, TOP_K
         TEMPERATURE, MAX_NEW_TOKENS, TOP_K = temp, int(max_tok), int(topk)
+        
         user_msg = history[-1][0]
-        bot_msg = get_response(user_msg, history[:-1])
+        bot_msg = get_response(user_msg, history[:-1], session_id)
         history[-1][1] = bot_msg
         return history
 
-    def retry_last(history, temp, max_tok, topk):
+    def retry_last(history, temp, max_tok, topk, session_id):
         if not history:
             return history
         user_msg = history[-1][0]
-        bot_msg = get_response(user_msg, history[:-1])
+        bot_msg = get_response(user_msg, history[:-1], session_id)
         history[-1][1] = bot_msg
         return history
 
     msg.submit(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, [chatbot, temp_slider, max_tok_slider, top_k_slider], chatbot
+        bot_response, [chatbot, temp_slider, max_tok_slider, top_k_slider, session_state], chatbot
     )
     send_btn.click(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, [chatbot, temp_slider, max_tok_slider, top_k_slider], chatbot
+        bot_response, [chatbot, temp_slider, max_tok_slider, top_k_slider, session_state], chatbot
     )
     clear_btn.click(lambda: None, None, chatbot, queue=False)
-    retry_btn.click(retry_last, [chatbot, temp_slider, max_tok_slider, top_k_slider], chatbot)
-
-    gr.HTML(f"<footer><center><p>🧠 Powered by LLaMA-based ChatDoctor | Device: {device.upper()}</p></center></footer>")
+    retry_btn.click(retry_last, [chatbot, temp_slider, max_tok_slider, top_k_slider, session_state], chatbot)
+
+    gr.HTML(f"""
+        <footer>
+            <p><strong>🧠 Powered by LLaMA-based ChatDoctor</strong></p>
+            <p>Device: {device.upper()} | Rate Limit: {MAX_REQUESTS_PER_MINUTE} requests/minute</p>
+            <p style='font-size:0.85em;margin-top:10px;'>
+                This AI provides general health information only. Always consult healthcare professionals for medical advice.
+            </p>
+        </footer>
+    """)
 
 # =============================
 # Launch App
 # =============================
 if __name__ == "__main__":
-    print("\n💡 Launching ChatDoctor Gradio Interface...")
+    print("\n💡 Launching Enhanced ChatDoctor Gradio Interface...")
+    print(f"📊 Configuration:")
+    print(f"   - Max History Turns: {MAX_HISTORY_TURNS}")
+    print(f"   - Rate Limit: {MAX_REQUESTS_PER_MINUTE} requests/minute")
+    print(f"   - Device: {device.upper()}")
     demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)