app.py

import gradio as gr
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value='''
# 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:

1. For Questions About Body Parts/Organs:
   ```markdown
   ### Anatomical Details
   - **Name and Location:** [anatomical position]
   - **Primary Functions:** 
     * [function 1]
     * [function 2]
   - **Structure:** [basic anatomy]
   - **Related Conditions:** [common conditions]
   ```

2. For Medication Queries:
   ```markdown
   ### Medication Information
   - **Generic/Brand Names:**
     * [list names]
   - **Drug Class:** [classification]
   - **General Uses:**
     * [primary uses]
   - **Common Side Effects:**
     * [side effects]
   - **Drug Interactions:**
     * [important interactions]
   
   > **Important Notice:** This information is for educational purposes only. 
   > Consult a healthcare provider for medical advice.
   ```

3. For Symptom Analysis:
   ```markdown
   ### Symptom Evaluation
   - **Possible Causes:** 
     * [common to serious]
   - **Key Characteristics:**
     * [important symptoms]
   - **Self-Care Measures:**
     * [if applicable]
   
   ⚠️ **Seek Immediate Medical Care If:**
   * [emergency signs]
   * [red flags]
   ```

4. For Treatment Information:
   ```markdown
   ### Treatment Plan
   1. **Conservative Measures:**
      * [self-care steps]
   2. **Medical Interventions:**
      * [common treatments]
   3. **Prevention:**
      * [prevention strategies]
   4. **Follow-up Care:**
      * [monitoring steps]
   ```

5. For Emergency Questions:
   ```markdown
   ### ⚠️ Emergency Guidance
   1. **Immediate Actions:**
      * [first steps]
   2. **While Waiting for Help:**
      * [temporary measures]
   3. **Contact Emergency Services If:**
      * [critical signs]
   ```

6. For Mental Health Questions:
   ```markdown
   ### Mental Health Support
   - **Common Symptoms:**
     * [symptom list]
   - **Coping Strategies:**
     * [self-help techniques]
   - **Professional Help:**
     * [when to seek help]
   
   > 🆘 **Crisis Support:** If you're having thoughts of self-harm, 
   > contact emergency services or crisis helpline immediately.
   ```

7. For Preventive Care:
   ```markdown
   ### Preventive Healthcare
   1. **Lifestyle Recommendations:**
      * [healthy habits]
   2. **Screening Tests:**
      * [recommended tests]
   3. **Vaccination Schedule:**
      * [if applicable]
   ```

8. For Diet/Nutrition:
   ```markdown
   ### Nutritional Guidance
   - **Dietary Recommendations:**
     * [food choices]
   - **Nutrients of Focus:**
     * [key nutrients]
   - **Meal Planning:**
     * [practical tips]
   ```

CONVERSATION HANDLING:
- Reference previous symptoms/conditions mentioned
- Track medication discussions
- Note any allergies or contraindications mentioned
- Follow up on previous advice given
- Ask clarifying questions when needed

FORMATTING GUIDELINES:
- Use markdown headers (###) for sections
- Format lists with proper indentation
- Use **bold** for emphasis
- Include > blockquotes for important notices
- Add emoji indicators:
  * ⚠️ for warnings
  * 💡 for tips
  * 🆘 for emergencies
  * ✅ for recommendations
  * ❌ for contraindications

IMPORTANT NOTES:
- Always include relevant medical disclaimers
- Redirect non-medical queries politely
- Maintain professional yet understandable language
- Cite medical guidelines when applicable
- Recommend professional consultation when necessary

<end_of_turn>
<start_of_turn>user
{query}<end_of_turn>
<start_of_turn>model
'''
)

# 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(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)


if __name__ == "__main__":
    demo.launch()