app.py

import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Define the medical system prompt
MEDICAL_SYSTEM_PROMPT = """<bos><start_of_turn>system
You are Gemma, a medical AI assistant. You MUST ONLY answer health and medical-related questions.
USE BULLET POINTS TO STRUCTURE YOUR ANSWERS AND THEY MUST BE PUNCTUATED AND GRAMMATICALLY CORRECT.
Your responses should be concise and informative.
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."""

# Function to check if a query is medical-related
def is_medical_query(query):
    medical_keywords = [
        "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"
    ]
    
    query_lower = query.lower()
    return any(keyword in query_lower for keyword in medical_keywords)

# Response function for the chatbot
def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
    # Check if the query is medical-related
    if not is_medical_query(message):
        return "I'm specialized in medical topics only. I cannot answer this question. How can I assist with a health-related concern instead?"
    
    # Prepare the messages for the API
    messages = [{"role": "system", "content": MEDICAL_SYSTEM_PROMPT}]
    
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    
    # Generate the response
    response = ""
    
    try:
        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
            if token:
                response += token
                yield response
    except Exception as e:
        yield f"I apologize, but I encountered an error: {str(e)}. Please try again."

# Create the Gradio interface
demo = gr.ChatInterface(
    fn=respond,
    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}"
)

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