Update app.py

app.py

@@ -12,15 +12,20 @@ def process_input(text_input, labels_or_premise, mode):
         results = {label: score for label, score in zip(prediction['labels'], prediction['scores'])}
         return results, ''
     else:  # NLI mode
-        prediction = nli_classifier([{"text": text_input, "text_pair": labels_or_premise}])
-        results = {pred['label']: pred['score'] for pred in prediction}
+        prediction = nli_classifier([{"text": text_input, "text_pair": labels_or_premise}])[0]
+        # Force showing all three labels
+        results = {
+            "entailment": prediction.get("score", 0) if prediction.get("label") == "entailment" else 0,
+            "contradiction": prediction.get("score", 0) if prediction.get("label") == "contradiction" else 0,
+            "neutral": prediction.get("score", 0) if prediction.get("label") == "neutral" else 0
+        }
         return results, ''
 
 def update_interface(mode):
     if mode == "Zero-Shot Classification":
         return gr.update(label="🏷️ Categories", placeholder="Enter comma-separated categories...")
     else:
-        return gr.update(label="Hypothesis", placeholder="Enter a hypothesis to compare with the premise...")
+        return gr.update(label="🔎 Hypothesis", placeholder="Enter a hypothesis to compare with the premise...")
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 ModernBERT Text Analysis")