prpobability update

app.py

@@ -48,15 +48,22 @@ def classify(text):
         logits_task1, logits_task2 = model(input_ids, attention_mask)
     
     # Softmax to get probabilities
-    probs_task1 = F.softmax(logits_task1, dim=1).cpu().numpy()
-    probs_task2 = F.softmax(logits_task2, dim=1).cpu().numpy()
+    probs_task1 = F.softmax(logits_task1, dim=1).cpu().numpy()[0]  # Extract the first batch item
+    probs_task2 = F.softmax(logits_task2, dim=1).cpu().numpy()[0]  # Extract the first batch item
 
     # Predict label with highest probability
     pred_task1 = label_mapping_task1[probs_task1.argmax()]
     pred_task2 = label_mapping_task2[probs_task2.argmax()]
 
-    # Return the two results separately
-    return pred_task1, pred_task2
+    # Format probabilities as percentages
+    probs_task1_str = ", ".join([f"{label}: {prob*100:.2f}%" for label, prob in zip(label_mapping_task1, probs_task1)])
+    probs_task2_str = ", ".join([f"{label}: {prob*100:.2f}%" for label, prob in zip(label_mapping_task2, probs_task2)])
+
+    # Combine label predictions with their probabilities
+    result_task1 = f"{pred_task1} ({probs_task1_str})"
+    result_task2 = f"{pred_task2} ({probs_task2_str})"
+
+    return result_task1, result_task2
 
 # Gradio Interface
 iface = gr.Interface(fn=classify,