Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from transformers import pipeline
 
-# Step 1: Load our AI Model
+# Step 1: Load the AI Model
 print("Loading the Health Analysis model...")
 health_classifier = pipeline(
     "text-classification",
@@ -9,33 +9,49 @@ health_classifier = pipeline(
 )
 print("Model loaded successfully!")
 
-# Step 2: Define the main function for the app
+# Step 2: Create the Label-to-Name Dictionary
+# This is the translator that maps the model's output to real disease names.
+label_map = {
+    'LABEL_0': 'Fungal infection', 'LABEL_1': 'Allergy', 'LABEL_2': 'GERD', 
+    'LABEL_3': 'Chronic cholestasis', 'LABEL_4': 'Drug Reaction', 'LABEL_5': 'Peptic ulcer disease', 
+    'LABEL_6': 'AIDS', 'LABEL_7': 'Diabetes', 'LABEL_8': 'Gastroenteritis', 
+    'LABEL_9': 'Bronchial Asthma', 'LABEL_10': 'Hypertension', 'LABEL_11': 'Migraine', 
+    'LABEL_12': 'Cervical spondylosis', 'LABEL_13': 'Paralysis (brain hemorrhage)', 
+    'LABEL_14': 'Jaundice', 'LABEL_15': 'Malaria', 'LABEL_16': 'Chicken pox', 
+    'LABEL_17': 'Dengue', 'LABEL_18': 'Typhoid', 'LABEL_19': 'Hepatitis A', 
+    'LABEL_20': 'Hepatitis B', 'LABEL_21': 'Hepatitis C', 'LABEL_22': 'Hepatitis D', 
+    'LABEL_23': 'Hepatitis E', 'LABEL_24': 'Alcoholic hepatitis', 'LABEL_25': 'Tuberculosis', 
+    'LABEL_26': 'Common Cold', 'LABEL_27': 'Pneumonia', 'LABEL_28': 'Dimorphic hemmorhoids(piles)', 
+    'LABEL_29': 'Heart attack', 'LABEL_30': 'Varicose veins', 'LABEL_31': 'Hypothyroidism', 
+    'LABEL_32': 'Hyperthyroidism', 'LABEL_33': 'Hypoglycemia', 'LABEL_34': 'Osteoarthristis', 
+    'LABEL_35': 'Arthritis', 'LABEL_36': '(vertigo) Paroymsal Positional Vertigo', 
+    'LABEL_37': 'Acne', 'LABEL_38': 'Urinary tract infection', 'LABEL_39': 'Psoriasis', 
+    'LABEL_40': 'Impetigo'
+}
+
+# Step 3: Define the main analysis function
 def analyze_symptoms(symptoms):
-    # The pipeline returns a list containing one dictionary, e.g., [{'label': 'Allergy', 'score': 0.98}]
-    # We select the first item from the list to get the dictionary.
+    # Get the prediction dictionary from the model, e.g., {'label': 'LABEL_2', 'score': 0.96}
     prediction_dict = health_classifier(symptoms)[0]
     
-    # Extract the disease name and score from the dictionary.
-    disease = prediction_dict['label']
+    # Get the generic label (e.g., 'LABEL_2')
+    generic_label = prediction_dict['label']
+    
+    # Look up the real disease name in our dictionary
+    disease = label_map.get(generic_label, "Unknown Condition")
+    
+    # Get the confidence score
     confidence_score = prediction_dict['score']
     
-    # Format the results into a clean, readable string for the user.
+    # Format the final, clean string for the user
     return f"Predicted Condition: {disease}\nConfidence Score: {confidence_score:.2f}"
 
-# Step 3: Define the content for our app
+# Step 4: Define the app's text content
 app_title = "Symptom to Disease Classifier"
-app_description = """
-Enter a list of symptoms, and this AI will predict a possible related medical condition.
-This tool is for educational purposes only and is based on the 'shanover/disease_classifier_base' model from Hugging Face.
-"""
-disclaimer = """
-**⚠️ IMPORTANT DISCLAIMER ⚠️**
-This is NOT a medical diagnostic tool. The predictions are generated by an AI model and may be inaccurate.
-**Always consult with a qualified healthcare professional for any medical advice or diagnosis.**
-This tool should not be used for any real-world medical decision-making.
-"""
+app_description = "Enter symptoms to predict a possible related medical condition. For educational use only."
+disclaimer = "⚠️ **IMPORTANT DISCLAIMER:** This is NOT a medical diagnostic tool. Always consult a qualified healthcare professional for medical advice."
 
-# Step 4: Create and Launch the Gradio Web App Interface
+# Step 5: Create and Launch the Gradio Interface
 app = gr.Interface(
     fn=analyze_symptoms,
     inputs=gr.Textbox(
@@ -53,5 +69,4 @@ app = gr.Interface(
     ]
 )
 
-# Go live!
 app.launch()