Update app.py

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 from transformers import pipeline
 
 # Step 2: Load our AI Model
-# This time, we are using a text-classification pipeline with our chosen health model.
 print("Loading the Health Analysis model...")
 health_classifier = pipeline(
     "text-classification",
@@ -11,23 +10,39 @@ health_classifier = pipeline(
 )
 print("Model loaded successfully!")
 
+# NEW: Step 2.5 - Create the Label Dictionary
+# The model outputs generic names like 'label_0'. This dictionary maps them 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 function for the app
-# This function takes the user's symptom text and returns a formatted prediction.
 def analyze_symptoms(symptoms):
-    # Use the pipeline to get the prediction. It returns a list of dictionaries.
     predictions = health_classifier(symptoms)
-    # For this model, the top prediction is the first item in the list.
     top_prediction = predictions[0]
     
-    # Extract the predicted disease and the confidence score.
-    disease = top_prediction['label']
+    # MODIFIED: Look up the generic label in our new dictionary
+    generic_label = top_prediction['label']
+    disease = label_map.get(generic_label, "Unknown Label") # .get() safely handles unknown labels
     confidence_score = top_prediction['score']
     
-    # Return the results in a clear, readable string.
     return f"Predicted Condition: {disease}\nConfidence: {confidence_score:.2f}"
 
-# Step 4: Define the content for our app
-# We'll create a title, a detailed description, and a very important disclaimer.
+# Step 4: Define the content for our app (No changes here)
 app_title = "Symptom to Disease Classifier"
 app_description = """
 Enter a list of symptoms, and this AI will predict a possible related medical condition.
@@ -40,7 +55,7 @@ This is NOT a medical diagnostic tool. The predictions are generated by an AI mo
 This tool should not be used for any real-world medical decision-making.
 """
 
-# Step 5: Create and Launch the Gradio Web App Interface
+# Step 5: Create and Launch the Gradio Web App Interface (No changes here)
 app = gr.Interface(
     fn=analyze_symptoms,
     inputs=gr.Textbox(
@@ -50,7 +65,7 @@ app = gr.Interface(
     outputs="text",
     title=app_title,
     description=app_description,
-    article=disclaimer,  # The 'article' parameter is great for adding extra text like our disclaimer.
+    article=disclaimer,
     examples=[
         ["I have a persistent cough, high fever, and body aches."],
         ["Feeling very thirsty and needing to urinate frequently."],