Update app.py

app.py

@@ -1,107 +1,165 @@
 import gradio as gr
 from transformers import pipeline
+import requests
 
-class HealthAssistant:
+class AIHealthAssistant:
     def __init__(self):
-        # Using an officially available medical model
+        # Initialize symptom checker model
         self.symptom_checker = pipeline(
-            "text-classification", 
-            model="emilyalsentzer/Bio_ClinicalBERT"
+            "text-classification",
+            model="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
         )
         
-        # 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"
-        }
+        # Initialize medical QA model
+        self.medical_qa = pipeline(
+            "question-answering",
+            model="deepset/roberta-base-squad2"
+        )
         
-        # 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"]
+        # Knowledge base (can be replaced with API calls)
+        self.disease_db = {
+            "influenza": {
+                "symptoms": ["fever", "cough", "sore throat", "runny nose", "body aches"],
+                "advice": "Rest, stay hydrated, take fever reducers like acetaminophen",
+                "precautions": ["Annual flu vaccine", "Frequent hand washing", "Avoid close contact"]
             },
-            "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"]
+            "migraine": {
+                "symptoms": ["severe headache", "nausea", "sensitivity to light", "aura"],
+                "advice": "Rest in dark room, take prescribed medication, apply cold compress",
+                "precautions": ["Identify triggers", "Maintain sleep schedule", "Stay hydrated"]
             }
         }
     
-    def predict_disease(self, symptoms):
+    def get_disease_from_symptoms(self, symptoms):
+        """Predict disease from symptoms using Hugging Face model"""
         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']
+            # For production, replace with a proper medical model
+            result = self.symptom_checker(symptoms)
             
-            # Get additional info if available
-            info = self.disease_info.get(disease_name, {})
+            # Map to diseases in our database
+            for disease in self.disease_db:
+                if disease in symptoms.lower():
+                    return disease
             
-            response = f"Possible condition: {disease_name} (confidence: {confidence:.2f})\n\n"
+            # Fallback to first disease (in real app, use proper mapping)
+            return list(self.disease_db.keys())[0]
+        
+        except Exception as e:
+            print(f"Model error: {e}")
+            return "unknown"
+    
+    def get_medical_info(self, disease, info_type):
+        """Get medical information from database or API"""
+        disease = disease.lower()
+        
+        # Check local database first
+        if disease in self.disease_db:
+            return self.disease_db[disease].get(info_type, "Information not available")
+        
+        # Fallback to API (example using hypothetical medical API)
+        try:
+            if info_type == "symptoms":
+                prompt = f"What are the symptoms of {disease}?"
+            elif info_type == "advice":
+                prompt = f"What is the treatment for {disease}?"
+            else:
+                prompt = f"What precautions should be taken for {disease}?"
             
-            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')}"
+            # In a real app, replace with actual API call:
+            # response = requests.get(f"https://medical-api.example.com/{disease}")
+            # return response.json().get(info_type)
             
-            return response
+            # For demo, using the QA model
+            context = f"{disease} is a medical condition. {self.get_medical_advice_from_api(disease)}"
+            result = self.medical_qa(question=prompt, context=context)
+            return result['answer']
         
         except Exception as e:
-            return f"System is currently analyzing your symptoms. For now, please consult a doctor about: {symptoms}"
+            print(f"API error: {e}")
+            return "Information not available"
+    
+    def get_medical_advice_from_api(self, disease):
+        """Simulate API response for demo purposes"""
+        api_responses = {
+            "diabetes": "Diabetes requires blood sugar monitoring, insulin therapy, and dietary changes.",
+            "hypertension": "Hypertension management includes medication, low-salt diet, and regular exercise."
+        }
+        return api_responses.get(disease.lower(), "Consult a healthcare professional for proper diagnosis and treatment.")
 
-def create_interface():
-    assistant = HealthAssistant()
+def create_demo():
+    assistant = AIHealthAssistant()
     
-    with gr.Blocks(title="Medical Symptom Checker", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🩺 AI Symptom Checker (Demo)")
+    def process_input(user_input, mode):
+        if mode == "Symptoms to Disease":
+            disease = assistant.get_disease_from_symptoms(user_input)
+            symptoms = assistant.get_medical_info(disease, "symptoms")
+            advice = assistant.get_medical_info(disease, "advice")
+            precautions = assistant.get_medical_info(disease, "precautions")
+            
+            output = (
+                f"🔍 Possible Condition: {disease.capitalize()}\n\n"
+                f"📋 Symptoms:\n- " + "\n- ".join(symptoms) + "\n\n"
+                f"💊 Recommended Actions:\n{advice}\n\n"
+                f"🛡️ Precautions:\n- " + "\n- ".join(precautions)
+            )
         
-        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")
+        elif mode == "Disease to Symptoms":
+            symptoms = assistant.get_medical_info(user_input, "symptoms")
+            advice = assistant.get_medical_info(user_input, "advice")
             
-            with gr.Column():
-                output = gr.Textbox(
-                    label="Analysis Result",
-                    interactive=False,
-                    lines=10
-                )
+            output = (
+                f"📋 Symptoms of {user_input.capitalize()}:\n- " + "\n- ".join(symptoms) + "\n\n"
+                f"💊 Recommended Actions:\n{advice}"
+            )
+        
+        return output
+    
+    with gr.Blocks(title="AI Health Assistant") as demo:
+        gr.Markdown("# 🏥 AI Health Assistant")
+        gr.Markdown("Enter symptoms or a disease name to get medical information")
+        
+        with gr.Row():
+            input_mode = gr.Radio(
+                choices=["Symptoms to Disease", "Disease to Symptoms"],
+                label="Input Mode"
+            )
+            user_input = gr.Textbox(
+                label="Input",
+                placeholder="Enter symptoms or disease name..."
+            )
+        
+        submit_btn = gr.Button("Get Medical Information")
+        output = gr.Textbox(label="Result", interactive=False, lines=10)
         
-        # Examples
+        # Example inputs
         gr.Examples(
             examples=[
-                ["headache, fever, body aches"],
-                ["nausea, vomiting, diarrhea"],
-                ["shortness of breath, cough"]
+                ["headache, nausea, sensitivity to light", "Symptoms to Disease"],
+                ["influenza", "Disease to Symptoms"],
+                ["fever, cough, sore throat", "Symptoms to Disease"]
             ],
-            inputs=symptoms
+            inputs=[user_input, input_mode],
+            outputs=output,
+            fn=process_input,
+            cache_examples=True
         )
         
-        analyze_btn.click(
-            fn=assistant.predict_disease,
-            inputs=symptoms,
+        submit_btn.click(
+            fn=process_input,
+            inputs=[user_input, input_mode],
             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.
+        This AI assistant provides general health information only and is not a substitute 
+        for professional medical advice, diagnosis, or treatment.
+        Always seek the advice of your physician or other qualified health provider 
+        with any questions you may have regarding a medical condition.
         """)
     
     return demo
 
 if __name__ == "__main__":
-    demo = create_interface()
+    demo = create_demo()
     demo.launch()