app.py

import gradio as gr
from transformers import pipeline

class HealthAssistant:
    def __init__(self):
        # Using a medical-specific model
        self.symptom_checker = pipeline(
            "text-classification", 
            model="palaksharma/bert-base-uncased-medical"
        )
        
        # Disease mapping for readable outputs
        self.label_mapping = {
            "LABEL_0": "Common Cold",
            "LABEL_1": "Influenza (Flu)",
            "LABEL_2": "Migraine",
            "LABEL_3": "Allergic Rhinitis",
            "LABEL_4": "Gastroenteritis"
        }
        
        # Enhanced knowledge base
        self.disease_info = {
            "Common Cold": {
                "symptoms": ["runny nose", "sneezing", "congestion", "sore throat"],
                "advice": "Rest, hydration, OTC cold medicines",
                "precautions": ["Hand hygiene", "Disinfect surfaces"]
            },
            "Influenza (Flu)": {
                "symptoms": ["fever", "cough", "body aches", "fatigue"],
                "advice": "Antivirals if prescribed, rest, fluids",
                "precautions": ["Flu vaccine", "Avoid crowds during outbreaks"]
            }
        }
    
    def predict_disease(self, symptoms):
        try:
            # Get model prediction
            result = self.symptom_checker(symptoms)
            
            # Map label to readable name
            label = result[0]['label']
            disease_name = self.label_mapping.get(label, "Unknown 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', 'Consult a doctor')}"
            
            return response
        
        except Exception as e:
            return f"Error analyzing symptoms: {str(e)}"

def create_interface():
    assistant = HealthAssistant()
    
    with gr.Blocks(title="Medical Symptom Checker", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🏥 AI Symptom Checker")
        
        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"],
                ["runny nose, sneezing, sore throat"],
                ["nausea, vomiting, diarrhea"]
            ],
            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()