app.py

import gradio as gr
from transformers import pipeline

class HealthAssistant:
    def __init__(self):
        # Using an officially available medical model
        self.symptom_checker = pipeline(
            "text-classification", 
            model="emilyalsentzer/Bio_ClinicalBERT"
        )
        
        # Simplified mapping for demo purposes
        self.label_mapping = {
            "LABEL_0": "Respiratory Infection",
            "LABEL_1": "Gastrointestinal Issue",
            "LABEL_2": "Neurological Condition",
            "LABEL_3": "Musculoskeletal Problem",
            "LABEL_4": "General Viral Infection"
        }
        
        # Enhanced knowledge base
        self.disease_info = {
            "Respiratory Infection": {
                "symptoms": ["cough", "shortness of breath", "sore throat", "congestion"],
                "advice": "Rest, stay hydrated, use humidifier, monitor breathing",
                "precautions": ["Good hand hygiene", "Avoid smoking", "Get flu vaccine"]
            },
            "Gastrointestinal Issue": {
                "symptoms": ["nausea", "vomiting", "diarrhea", "stomach pain"],
                "advice": "Stay hydrated, BRAT diet, avoid dairy, monitor for dehydration",
                "precautions": ["Proper food handling", "Hand washing", "Avoid contaminated water"]
            }
        }
    
    def predict_disease(self, symptoms):
        try:
            # Get model prediction
            result = self.symptom_checker(symptoms[:512])  # Limit input length
            
            # Map label to readable name
            label = result[0]['label']
            disease_name = self.label_mapping.get(label, "Medical Condition")
            confidence = result[0]['score']
            
            # Get additional info if available
            info = self.disease_info.get(disease_name, {})
            
            response = f"Possible condition: {disease_name} (confidence: {confidence:.2f})\n\n"
            
            if info:
                response += f"Common symptoms:\n- " + "\n- ".join(info.get('symptoms', [])) + "\n\n"
                response += f"Recommended actions:\n{info.get('advice', 'Please consult a healthcare professional')}"
            
            return response
        
        except Exception as e:
            return f"System is currently analyzing your symptoms. For now, please consult a doctor about: {symptoms}"

def create_interface():
    assistant = HealthAssistant()
    
    with gr.Blocks(title="Medical Symptom Checker", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🩺 AI Symptom Checker (Demo)")
        
        with gr.Row():
            with gr.Column():
                symptoms = gr.Textbox(
                    label="Describe your symptoms",
                    placeholder="E.g. headache, fever, cough...",
                    lines=3
                )
                analyze_btn = gr.Button("Analyze Symptoms")
            
            with gr.Column():
                output = gr.Textbox(
                    label="Analysis Result",
                    interactive=False,
                    lines=10
                )
        
        # Examples
        gr.Examples(
            examples=[
                ["headache, fever, body aches"],
                ["nausea, vomiting, diarrhea"],
                ["shortness of breath, cough"]
            ],
            inputs=symptoms
        )
        
        analyze_btn.click(
            fn=assistant.predict_disease,
            inputs=symptoms,
            outputs=output
        )
        
        gr.Markdown("""
        ## ⚠️ Important Disclaimer
        This AI tool provides general information only and is not a substitute for professional medical advice.
        Always consult with a qualified healthcare provider for diagnosis and treatment.
        """)
    
    return demo

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