Update app.py

app.py

@@ -11,7 +11,19 @@ models = {
 def classify_text(model_name, text):
     classifier = pipeline("text-classification", model=models[model_name], top_k=None)
     predictions = classifier(text)
-    return {pred["label"]: pred["score"] for pred in predictions[0]}
+
+    # Map the numerical labels to human-readable labels
+    label_mapping = {"0": "Male", "1": "Female"}
+
+    # Construct the output dictionary with human-readable labels
+    output_predictions = {}
+    for pred in predictions[0]:
+        # Ensure the label is treated as a string for dictionary lookup
+        numerical_label_str = str(pred["label"])
+        human_readable_label = label_mapping.get(numerical_label_str, numerical_label_str) # Use fallback if label not in mapping
+        output_predictions[human_readable_label] = pred["score"]
+
+    return output_predictions
 
 # Create the Gradio interface
 interface = gr.Interface(
@@ -20,7 +32,7 @@ interface = gr.Interface(
         gr.Dropdown(
             list(models.keys()),
             label="Select Model",
-            value="ModernBERT Base (Go-Emotions)"
+            value="ModernBERT Large (gender)"
         ),
         gr.Textbox(
             lines=2,
@@ -29,8 +41,8 @@ interface = gr.Interface(
         )
     ],
     outputs=gr.Label(num_top_classes=5),
-    title="🎭 ModernBERT Emotion Classifier",
-    description="Select a model and enter a sentence to see its associated emotions and confidence scores.",
+    title="ModernBERT gender Classifier",
+    description="Select a model and enter a sentence to see its associated gender and confidence scores.",
 )
 
 # Launch the app