Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import gc
 import torch
 import gradio as gr
-from transformers import LlamaTokenizer, LlamaForCausalLM, StoppingCriteria, StoppingCriteriaList
+from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
 
 # =============================
 # Configuration
@@ -13,26 +13,24 @@ TEMPERATURE = 0.5
 TOP_K = 50
 REPETITION_PENALTY = 1.1
 
-# Detect device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Loading model from {MODEL_PATH} on {device}...")
+print(f"🚀 Loading model from {MODEL_PATH} on {device}...")
 
 # =============================
-# Load Tokenizer and Model
+# Load Model & Tokenizer
 # =============================
-tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)
-model = LlamaForCausalLM.from_pretrained(
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH,
     device_map="auto",
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     low_cpu_mem_usage=True
 )
 
-generator = model.generate
 print("✅ ChatDoctor model loaded successfully!\n")
 
 # =============================
-# Stopping Criteria
+# Stop Criteria
 # =============================
 class StopOnTokens(StoppingCriteria):
     def __init__(self, stop_ids):
@@ -49,132 +47,89 @@ class StopOnTokens(StoppingCriteria):
                         return True
         return False
 
+
 # =============================
 # Medical Keywords and Validation
 # =============================
 MEDICAL_KEYWORDS = [
-    # Symptoms
-    "pain", "ache", "symptom", "hurt", "sore", "discomfort", "suffering",
-    # Common conditions
-    "fever", "cough", "cold", "flu", "infection", "allergy", "diabetes", "pressure",
-    "asthma", "migraine", "nausea", "vomit", "diarrhea", "constipation",
-    # Body parts
-    "heart", "stomach", "head", "back", "chest", "throat", "lung", "kidney",
-    "liver", "brain", "skin", "eye", "ear", "nose", "tooth", "teeth", "joint",
-    "muscle", "bone", "neck", "shoulder", "knee", "ankle", "foot", "hand",
-    # Medical terms
-    "doctor", "hospital", "clinic", "emergency", "ambulance", "medication",
-    "medicine", "prescription", "diagnosis", "treatment", "therapy", "cure",
-    "sick", "ill", "disease", "condition", "disorder", "syndrome",
-    # Injuries
-    "injury", "wound", "cut", "bruise", "fracture", "sprain", "burn", "bleed",
-    # Vitals and tests
-    "blood", "pressure", "temperature", "pulse", "breathing", "test", "scan",
-    # Mental health
-    "stress", "anxiety", "depression", "mental", "sleep", "insomnia", "tired",
-    "fatigue", "exhausted", "mood", "panic", "worry",
-    # Lifestyle/wellness
-    "diet", "nutrition", "exercise", "weight", "vitamin", "supplement", "healthy",
-    "wellness", "fitness", "eating", "appetite", "lifestyle", "food", "fruit",
-    "vegetable", "meal", "breakfast", "lunch", "dinner", "snack", "drink",
-    "water", "hydration", "protein", "carb", "fat", "calorie", "sugar",
-    "cholesterol", "gym", "workout", "run", "walk", "yoga", "sport",
-    # Serious conditions
-    "cancer", "tumor", "surgery", "stroke", "attack", "seizure", "diabetic",
-    # Questions about health
-    "health", "medical", "feel", "feeling", "comfortable", "uncomfortable",
-    "recommendation", "recommend", "advice", "suggest", "should i", "better",
-    "improve", "prevent", "avoid", "good for", "bad for"
+    "pain", "ache", "symptom", "hurt", "sore", "discomfort", "fever", "cough", "flu",
+    "infection", "allergy", "diabetes", "pressure", "asthma", "migraine", "vomit",
+    "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"
 ]
 
 CASUAL_ONLY_PATTERNS = [
-    "hey", "hi", "hello", "sup", "what's up", "whats up", "yo",
-    "good morning", "good evening", "good afternoon", "good night",
-    "how are you", "how r u", "wassup", "hiya", "greetings"
+    "hey", "hi", "hello", "sup", "yo", "good morning", "good evening",
+    "how are you", "wassup", "hiya"
 ]
 
+
 def is_medical_query(message):
-    """Check if the message contains medical-related content"""
     message_lower = message.lower()
-    
-    # Check for medical keywords
     for keyword in MEDICAL_KEYWORDS:
         if keyword in message_lower:
             return True
-    
-    # Check for question words combined with longer messages (might be medical)
     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 it has a question word and is longer than 5 words, might be medical
     if has_question and len(message.split()) > 5:
         return True
-    
     return False
 
+
 def is_only_greeting(message):
-    """Check if message is ONLY a casual greeting with no medical content"""
-    message_lower = message.lower().strip()
-    
-    # Remove punctuation for checking
-    message_clean = message_lower.replace("!", "").replace("?", "").replace(".", "").strip()
-    
-    # Check if it's a short greeting (3 words or less)
-    if len(message_clean.split()) <= 3:
+    message_lower = message.lower().strip().replace("!", "").replace("?", "").replace(".", "")
+    if len(message_lower.split()) <= 3:
         for pattern in CASUAL_ONLY_PATTERNS:
-            if message_clean == pattern or message_clean.startswith(pattern):
+            if message_lower == pattern or message_lower.startswith(pattern):
                 return True
-    
     return False
 
+
 # =============================
-# Get Response Function
+# Get Response
 # =============================
 def get_response(user_input, history_context):
-    """Generate response from ChatDoctor model"""
-    
-    # STRICT FILTERING: Only allow medical queries to reach the model
-    if not is_medical_query(user_input):
-        return "Hello! I'm ChatDoctor, an AI medical assistant specialized in health and medical topics. I can help you with:\n\n• Symptoms and health concerns\n• Medical conditions and treatments\n• General health advice\n• Wellness and prevention\n\nPlease describe any health-related symptoms or medical questions you have, and I'll do my best to assist you."
-    
-    human_invitation = "Patient: "
-    doctor_invitation = "ChatDoctor: "
-
-    # Enhanced system instruction
-    system_instruction = """You are ChatDoctor, a professional medical AI assistant. You ONLY discuss health, medical symptoms, treatments, and wellness topics. 
-
-If a patient greets you or asks non-medical questions, you must respond professionally: "I'm ChatDoctor, here to help with your health concerns. What medical symptoms or health questions can I assist you with today?"
+    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."
 
-Now continue the medical consultation:
+    if not is_medical_query(user_input):
+        return (
+            "Hello! I'm ChatDoctor, an AI medical assistant specialized in health and wellness.\n\n"
+            "I can help you with:\n"
+            "• Symptoms and medical conditions\n"
+            "• Treatment and prevention advice\n"
+            "• Fitness, diet, and mental health tips\n\n"
+            "Please describe your health concern in detail to get started."
+        )
 
-"""
+    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"
+    )
 
-    # Build conversation from history
+    # Build history
     history_text = [system_instruction]
     for human, assistant in history_context:
         if human:
-            history_text.append(human_invitation + human)
+            history_text.append(f"{human_prefix} {human}")
         if assistant:
-            history_text.append(doctor_invitation + assistant)
-    
-    # Add current user input with medical context reinforcement
-    if not is_medical_query(user_input):
-        user_input = f"{user_input} [Medical consultation context]"
-    
-    history_text.append(human_invitation + user_input)
+            history_text.append(f"{doctor_prefix} {assistant}")
+    history_text.append(f"{human_prefix} {user_input}")
 
-    # Build conversation prompt
-    prompt = "\n".join(history_text) + "\n" + doctor_invitation
+    prompt = "\n".join(history_text) + f"\n{doctor_prefix} "
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
-    # Define stop words and their token IDs
     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)])
 
-    # Generate model response
     with torch.no_grad():
-        output_ids = generator(
+        output_ids = model.generate(
             input_ids,
             max_new_tokens=MAX_NEW_TOKENS,
             do_sample=True,
@@ -186,49 +141,30 @@ Now continue the medical consultation:
             eos_token_id=tokenizer.eos_token_id
         )
 
-    # Decode and clean response
-    full_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    response = full_output[len(prompt):].strip()
-    
-    # Remove any "Patient:" that might have slipped through
-    for stop_word in ["Patient:", "Patient :", "\nPatient:", "\nPatient", "Patient"]:
+    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()
-    
-    # Post-processing: Check if response seems off-topic
-    response_lower = response.lower()
-    chatbot_service_keywords = ["chatbot", "service", "error code", "cloud-based", "platform"]
-    
-    if any(keyword in response_lower for keyword in chatbot_service_keywords):
-        # Model went off-topic, force redirect
-        response = "I apologize for any confusion. I'm ChatDoctor, and I'm specifically designed to help with medical and health-related questions. Could you please tell me about any health symptoms or medical concerns you're experiencing?"
-
-    # Free memory
+    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()
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
     return response
 
-# =============================
-# Gradio Chat Function
-# =============================
-def chat_function(message, history):
-    """Gradio chat interface function"""
-    if not message.strip():
-        return ""
-    
-    try:
-        response = get_response(message, history)
-        return response
-    except Exception as e:
-        return f"Error: {str(e)}"
 
 # =============================
-# Custom CSS
+# Gradio Interface
 # =============================
 custom_css = """
 #header {
@@ -239,18 +175,8 @@ custom_css = """
     border-radius: 10px;
     margin-bottom: 20px;
 }
-
-#header h1 {
-    margin: 0;
-    font-size: 2.5em;
-}
-
-#header p {
-    margin: 10px 0 0 0;
-    font-size: 1.1em;
-    opacity: 0.9;
-}
-
+#header h1 { margin: 0; font-size: 2.3em; }
+#header p { margin: 5px 0 0; font-size: 1em; opacity: 0.9; }
 .disclaimer {
     background-color: #fff3cd;
     border: 1px solid #ffc107;
@@ -259,153 +185,77 @@ custom_css = """
     margin: 20px 0;
     color: #856404;
 }
-
-.disclaimer h3 {
-    margin-top: 0;
-    color: #856404;
-}
-
-footer {
-    text-align: center;
-    margin-top: 30px;
-    color: #666;
-    font-size: 0.9em;
-}
 """
 
-# =============================
-# Gradio Interface
-# =============================
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    # Header
     gr.HTML("""
         <div id="header">
             <h1>🩺 ChatDoctor AI Assistant</h1>
-            <p>Your AI-powered medical conversation partner</p>
+            <p>Your AI-powered medical consultation partner</p>
         </div>
     """)
-    
-    # Disclaimer
     gr.HTML("""
         <div class="disclaimer">
             <h3>⚠️ Medical Disclaimer</h3>
-            <p><strong>Important:</strong> This AI assistant is for informational and educational purposes only. 
-            It is NOT a substitute for professional medical advice, diagnosis, or treatment. 
-            Always seek the advice of your physician or other qualified health provider with any questions 
-            you may have regarding a medical condition. Never disregard professional medical advice or 
-            delay in seeking it because of something you have read here.</p>
+            <p>This AI assistant is for informational purposes only. 
+            It is NOT a substitute for professional medical advice, diagnosis, or treatment.</p>
         </div>
     """)
-    
-    # Chatbot Interface
+
     chatbot = gr.Chatbot(
-        height=500,
-        placeholder="<div style='text-align: center; padding: 40px;'><h3>👋 Welcome to ChatDoctor!</h3><p>I'm here to discuss your health concerns. Please describe your symptoms or health question.</p></div>",
+        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>",
         show_label=False,
         avatar_images=(None, "🤖"),
     )
-    
+
     with gr.Row():
-        msg = gr.Textbox(
-            placeholder="Describe your health symptoms or medical concern here...",
-            show_label=False,
-            scale=9,
-            container=False
-        )
-        submit_btn = gr.Button("Send 📤", scale=1, variant="primary")
-    
+        msg = gr.Textbox(placeholder="Type your medical concern here...", show_label=False, scale=9, container=False)
+        send_btn = gr.Button("Send 📤", scale=1, variant="primary")
+
     with gr.Row():
         clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
         retry_btn = gr.Button("🔄 Retry", scale=1)
-    
-    # Examples
-    gr.Examples(
-        examples=[
-            "I have a persistent headache for 3 days. What should I do?",
-            "What are the symptoms of diabetes?",
-            "How can I improve my sleep quality?",
-            "I have a fever and sore throat. Should I be concerned?",
-            "What are some natural ways to reduce stress?",
-        ],
-        inputs=msg,
-        label="💡 Example Medical Questions"
-    )
-    
-    # Settings (collapsed by default)
+
     with gr.Accordion("⚙️ Advanced Settings", open=False):
-        temperature_slider = gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=TEMPERATURE,
-            step=0.1,
-            label="Temperature (Creativity)",
-            info="Higher values make responses more creative but less focused"
-        )
-        max_tokens_slider = gr.Slider(
-            minimum=50,
-            maximum=500,
-            value=MAX_NEW_TOKENS,
-            step=50,
-            label="Max Response Length",
-            info="Maximum number of tokens in response"
-        )
-        top_k_slider = gr.Slider(
-            minimum=1,
-            maximum=100,
-            value=TOP_K,
-            step=1,
-            label="Top K",
-            info="Limits vocabulary selection"
-        )
-    
-    # Footer
-    gr.HTML("""
-        <footer>
-            <p>Powered by ChatDoctor Model | Built with Gradio</p>
-            <p>Device: """ + device.upper() + """ | Model: LLaMA-based Medical AI</p>
-        </footer>
-    """)
-    
-    # Event handlers
+        temp_slider = gr.Slider(0.1, 1.0, TEMPERATURE, 0.1, label="Temperature")
+        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")
+
     def user_message(user_msg, history):
         return "", history + [[user_msg, None]]
-    
-    def bot_response(history, temp, max_tok, top_k_val):
+
+    def bot_response(history, temp, max_tok, topk):
         global TEMPERATURE, MAX_NEW_TOKENS, TOP_K
-        TEMPERATURE = temp
-        MAX_NEW_TOKENS = int(max_tok)
-        TOP_K = int(top_k_val)
-        
+        TEMPERATURE, MAX_NEW_TOKENS, TOP_K = temp, int(max_tok), int(topk)
         user_msg = history[-1][0]
-        bot_msg = chat_function(user_msg, history[:-1])
+        bot_msg = get_response(user_msg, history[:-1])
         history[-1][1] = bot_msg
         return history
-    
-    # Connect events
+
+    def retry_last(history, temp, max_tok, topk):
+        if not history:
+            return history
+        user_msg = history[-1][0]
+        bot_msg = get_response(user_msg, history[:-1])
+        history[-1][1] = bot_msg
+        return history
+
     msg.submit(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, [chatbot, temperature_slider, max_tokens_slider, top_k_slider], chatbot
+        bot_response, [chatbot, temp_slider, max_tok_slider, top_k_slider], chatbot
     )
-    
-    submit_btn.click(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, [chatbot, temperature_slider, max_tokens_slider, top_k_slider], 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
     )
-    
     clear_btn.click(lambda: None, None, chatbot, queue=False)
-    
-    def retry_last():
-        return None
-    
-    retry_btn.click(retry_last, 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>")
 
 # =============================
-# Launch Interface
+# Launch App
 # =============================
 if __name__ == "__main__":
-    print("\n🚀 Launching ChatDoctor Gradio Interface...")
+    print("\n💡 Launching ChatDoctor Gradio Interface...")
     demo.queue()
-    demo.launch(
-        server_name="0.0.0.0",  # Accessible from network
-        server_port=7860,
-        share=False,  # Set to True to create public link
-        show_error=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)