app.py

import gradio as gr
from transformers import pipeline

class HealthAssistant:
    def __init__(self):
        # Initialize with smaller models for faster loading
        self.symptom_checker = pipeline("text-classification", model="distilbert-base-uncased")
        self.medical_qa = pipeline("text-generation", model="gpt2")
        
        # Knowledge base (in production, connect to a medical database)
        self.disease_db = {
            "flu": {
                "symptoms": ["fever", "cough", "sore throat", "muscle aches", "headache", "fatigue"],
                "advice": """1. Get plenty of rest
2. Stay hydrated
3. Use over-the-counter fever reducers like acetaminophen
4. See a doctor if symptoms worsen or last more than 10 days""",
                "precautions": [
                    "Get annual flu vaccine",
                    "Wash hands frequently with soap",
                    "Avoid close contact with sick individuals",
                    "Cover coughs and sneezes"
                ]
            },
            "common cold": {
                "symptoms": ["runny nose", "sneezing", "congestion", "sore throat", "cough", "mild headache"],
                "advice": """1. Rest and hydrate
2. Use saline nasal drops
3. Gargle with warm salt water for sore throat
4. Over-the-counter cold medicines may help""",
                "precautions": [
                    "Wash hands often with soap and water",
                    "Avoid touching your face with unwashed hands",
                    "Stay away from people who are sick",
                    "Disinfect frequently touched surfaces"
                ]
            },
            "migraine": {
                "symptoms": ["severe headache", "nausea", "sensitivity to light", "sensitivity to sound", "aura"],
                "advice": """1. Rest in a quiet, dark room
2. Apply cold compress to forehead
3. Take prescribed migraine medication
4. Try relaxation techniques""",
                "precautions": [
                    "Identify and avoid triggers (certain foods, stress, etc.)",
                    "Maintain regular sleep schedule",
                    "Stay hydrated",
                    "Consider preventive medications if frequent"
                ]
            }
        }
    
    def predict_disease(self, symptoms):
        symptoms = symptoms.lower()
        matched = []
        
        # Simple matching for demo (replace with actual model in production)
        for disease, data in self.disease_db.items():
            symptom_count = sum(1 for s in data["symptoms"] if s in symptoms)
            if symptom_count >= 2:
                matched.append((disease, symptom_count))
        
        matched.sort(key=lambda x: x[1], reverse=True)
        
        if matched:
            return f"Possible conditions:\n" + "\n".join(
                f"- {disease} ({count} matching symptoms)" 
                for disease, count in matched[:3]  # Show top 3
            )
        
        # Fallback to model if no matches
        result = self.symptom_checker(symptoms)
        return f"The model suggests this might be related to: {result[0]['label']} (confidence: {result[0]['score']:.2f})"
    
    def get_symptoms(self, disease):
        disease = disease.lower()
        if disease in self.disease_db:
            return "Common symptoms:\n" + "\n".join(
                f"- {symptom}" for symptom in self.disease_db[disease]["symptoms"]
            )
        
        # Fallback to model
        prompt = f"What are the common symptoms of {disease}?"
        result = self.medical_qa(prompt, max_length=100)
        return result[0]["generated_text"]
    
    def get_advice(self, condition):
        condition = condition.lower()
        if condition in self.disease_db:
            data = self.disease_db[condition]
            return (
                "Medical Advice:\n" + data["advice"] + "\n\n" +
                "Precautions:\n" + "\n".join(f"- {p}" for p in data["precautions"])
            )
        
        # Fallback to model
        prompt = f"What is the recommended medical advice and precautions for {condition}?"
        result = self.medical_qa(prompt, max_length=200)
        return result[0]["generated_text"]

# Initialize assistant
assistant = HealthAssistant()

# Create Gradio interface
with gr.Blocks(title="AI Health Assistant", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 AI Health Assistant")
    gr.Markdown("Describe your symptoms or ask about a disease to get information.")
    
    with gr.Tab("Diagnose from Symptoms"):
        symptom_input = gr.Textbox(label="Describe your symptoms (e.g., headache, fever, cough)")
        symptom_output = gr.Textbox(label="Possible Conditions", interactive=False)
        symptom_button = gr.Button("Analyze Symptoms")
    
    with gr.Tab("Disease Information"):
        disease_input = gr.Textbox(label="Enter a disease name")
        disease_symptoms = gr.Textbox(label="Common Symptoms", interactive=False)
        disease_advice = gr.Textbox(label="Medical Advice", interactive=False)
        disease_button = gr.Button("Get Disease Info")
    
    # Event handlers
    symptom_button.click(
        fn=assistant.predict_disease,
        inputs=symptom_input,
        outputs=symptom_output
    )
    
    disease_button.click(
        fn=assistant.get_symptoms,
        inputs=disease_input,
        outputs=disease_symptoms
    )
    
    disease_button.click(
        fn=assistant.get_advice,
        inputs=disease_input,
        outputs=disease_advice
    )

    # Disclaimer
    gr.Markdown("""
    ## ⚠️ Important Disclaimer
    This AI assistant provides general health information only and is not a substitute for professional medical advice. 
    Always consult with a qualified healthcare provider for diagnosis and treatment.
    """)

# For Hugging Face Spaces
demo.launch()