Update app.py

app.py

@@ -3,145 +3,105 @@ 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")
+        # Using a medical-specific model
+        self.symptom_checker = pipeline(
+            "text-classification", 
+            model="palaksharma/bert-base-uncased-medical"
+        )
         
-        # 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"
-                ]
+        # 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"]
             },
-            "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"
-                ]
+            "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):
-        symptoms = symptoms.lower()
-        matched = []
+        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
         
-        # 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))
+        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")
         
-        matched.sort(key=lambda x: x[1], reverse=True)
+        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
+                )
         
-        if matched:
-            return f"Possible conditions:\n" + "\n".join(
-                f"- {disease} ({count} matching symptoms)" 
-                for disease, count in matched[:3]  # Show top 3
-            )
+        # Examples
+        gr.Examples(
+            examples=[
+                ["headache, fever, body aches"],
+                ["runny nose, sneezing, sore throat"],
+                ["nausea, vomiting, diarrhea"]
+            ],
+            inputs=symptoms
+        )
         
-        # 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"]
-            )
+        analyze_btn.click(
+            fn=assistant.predict_disease,
+            inputs=symptoms,
+            outputs=output
+        )
         
-        # 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"]
+        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.
+        """)
     
-    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.
-    """)
+    return demo
 
-# For Hugging Face Spaces
-demo.launch()
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()