Update app.py

app.py

@@ -47,19 +47,14 @@ demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Textbox(value='''
-        You are Gemma, a comprehensive medical AI assistant specialized in all health-related topics. You handle various types of medical queries and maintain conversation context.
-
-SPECIALIZATIONS:
-- General medicine and medical conditions
-- Anatomy and physiology
-- Dermatology and skin conditions
-- Symptom analysis and interpretation
-- Medication information and interactions
-- Treatment approaches and procedures
-- Preventive care and wellness
-- Emergency guidance
-- Mental health support
-- Nutritional advice
+# Update MEDICAL_PROMPT to be more restrictive
+MEDICAL_PROMPT = PromptTemplate(
+    input_variables=["query"],
+    template="""<bos><start_of_turn>system
+You are Gemma, a medical AI assistant. You MUST ONLY answer health and medical-related questions.
+Your responses should be professional, accurate, and focused on medical topics only.
+For any non-medical questions, respond with a redirection to medical topics.
+For medication queries, provide general information and recommend consulting a healthcare professional.
 
 QUERY TYPES AND RESPONSE FORMATS:
 
@@ -196,7 +191,66 @@ IMPORTANT NOTES:
 <start_of_turn>user
 {query}<end_of_turn>
 <start_of_turn>model
-"""''', label="System message"),
+"""
+)
+
+# Update is_medical_query to be more comprehensive
+def is_medical_query(query):
+    medical_keywords_and_greetings = [
+        "health", "disease", "symptom", "doctor", "medicine", "medical", "treatment",
+        "hospital", "clinic", "diagnosis", "patient", "drug", "prescription", "therapy",
+        "cancer", "diabetes", "heart", "blood", "pain", "surgery", "vaccine", "infection",
+        "allergy", "diet", "nutrition", "vitamin", "exercise", "mental health", "depression",
+        "anxiety", "disorder", "syndrome", "chronic", "acute", "emergency", "pharmacy",
+        "dosage", "side effect", "contraindication", "body", "organ", "immune", "virus",
+        "bacterial", "fungal", "parasite", "genetic", "hereditary", "congenital", "prenatal",
+        "headaches", "ache", "stomach ache", "skin", "head", "arm", "leg", "chest", "back", "throat", "eye", "ear", "nose", "mouth"
+    ]
+    
+    # Remove greetings from the keyword list
+    medical_keywords = [word for word in medical_keywords_and_greetings if word not in ["hello", "hi", "greetings", "good morning", "good afternoon", "good evening", "hey"]]
+    
+    query_lower = query.lower()
+    return any(keyword in query_lower for keyword in medical_keywords)
+
+# Update chat_with_model to enforce medical-only responses
+def chat_with_model(message, history):
+    try:
+        context = "\n".join([f"User: {msg}\nAssistant: {res}" for msg, res in history])
+        full_query = f"{context}\nUser: {message}"
+        
+        if not is_medical_query(full_query):
+            return "I'm specialized in medical topics only. I cannot answer this question. How can I assist with a health-related concern instead?"
+        
+        response = medical_chain.run(query=full_query)
+        clean_response = response.split("<end_of_turn>")[0].strip()
+        
+        # Check if the response is medical-related
+        if not is_medical_query(clean_response):
+            return "I can only provide information on medical topics. Please ask a medical question."
+        
+        return clean_response
+        
+    except Exception as e:
+        return f"I apologize, but I encountered an error: {str(e)}. Please try again."
+
+# Update Gradio examples to be medical-specific
+iface = gr.ChatInterface(
+    fn=chat_with_model,
+    title="MedexDroid - Medical Assistant",
+    examples=[
+        "What are the symptoms of diabetes?",
+        "How can I improve my diet for heart health?",
+        "What is the treatment for a migraine?",
+        "What are the side effects of aspirin?",
+        "What are the causes of high blood pressure?"
+    ],
+    description="An AI Medical Assistant. Please ask health-related questions only.",
+    theme=gr.themes.Soft(),
+    css=".gradio-container {background-color: #f0f4f8}"
+)
+  
+''', label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(