Update app.py

app.py

@@ -1,165 +1,50 @@
 import gradio as gr
 from transformers import pipeline
-import requests
 
-class AIHealthAssistant:
-    def __init__(self):
-        # Initialize symptom checker model
-        self.symptom_checker = pipeline(
-            "text-classification",
-            model="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
-        )
-        
-        # Initialize medical QA model
-        self.medical_qa = pipeline(
-            "question-answering",
-            model="deepset/roberta-base-squad2"
-        )
-        
-        # 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"]
-            },
-            "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 get_disease_from_symptoms(self, symptoms):
-        """Predict disease from symptoms using Hugging Face model"""
-        try:
-            # For production, replace with a proper medical model
-            result = self.symptom_checker(symptoms)
-            
-            # Map to diseases in our database
-            for disease in self.disease_db:
-                if disease in symptoms.lower():
-                    return disease
-            
-            # 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"
+# Initialize medical LLM (using smaller model for Spaces compatibility)
+med_llm = pipeline("text-generation", model="distilgpt2")
+
+def diagnose(symptoms):
+    prompt = f"""Act as an AI doctor. Analyze these symptoms:
+{symptoms}
+
+Possible diagnoses from most to least likely:
+1."""
     
-    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}?"
-            
-            # 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)
-            
-            # 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:
-            print(f"API error: {e}")
-            return "Information not available"
+    diagnosis = med_llm(
+        prompt,
+        max_length=400,
+        num_return_sequences=1,
+        temperature=0.7
+    )[0]['generated_text']
     
-    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.")
+    # Extract just the diagnosis part
+    return diagnosis.split("1.", 1)[-1].strip()
 
-def create_demo():
-    assistant = AIHealthAssistant()
+with gr.Blocks() as demo:
+    gr.Markdown("# 🌍 Global Symptom Checker")
+    gr.Markdown("Enter symptoms to get possible diagnoses")
+    
+    with gr.Row():
+        symptoms = gr.Textbox(label="Describe your symptoms", lines=3)
+        submit = gr.Button("Diagnose")
+    
+    output = gr.Textbox(label="Possible Conditions", interactive=False)
     
-    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)
-            )
-        
-        elif mode == "Disease to Symptoms":
-            symptoms = assistant.get_medical_info(user_input, "symptoms")
-            advice = assistant.get_medical_info(user_input, "advice")
-            
-            output = (
-                f"📋 Symptoms of {user_input.capitalize()}:\n- " + "\n- ".join(symptoms) + "\n\n"
-                f"💊 Recommended Actions:\n{advice}"
-            )
-        
-        return output
+    gr.Examples(
+        examples=[
+            ["fever, headache, muscle pain"],
+            ["cough, shortness of breath, fatigue"],
+            ["rash, joint pain, red eyes"]
+        ],
+        inputs=symptoms
+    )
     
-    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)
-        
-        # Example inputs
-        gr.Examples(
-            examples=[
-                ["headache, nausea, sensitivity to light", "Symptoms to Disease"],
-                ["influenza", "Disease to Symptoms"],
-                ["fever, cough, sore throat", "Symptoms to Disease"]
-            ],
-            inputs=[user_input, input_mode],
-            outputs=output,
-            fn=process_input,
-            cache_examples=True
-        )
-        
-        submit_btn.click(
-            fn=process_input,
-            inputs=[user_input, input_mode],
-            outputs=output
-        )
-        
-        gr.Markdown("""
-        ## ⚠️ Important Disclaimer
-        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.
-        """)
+    submit.click(diagnose, inputs=symptoms, outputs=output)
     
-    return demo
+    gr.Markdown("""
+    ⚠️ **Disclaimer**: This AI provides general information only. 
+    Always consult a real doctor for medical advice.
+    """)
 
-if __name__ == "__main__":
-    demo = create_demo()
-    demo.launch()
+demo.launch()