Add class probs

app.py

@@ -20,17 +20,21 @@ def main():
             st.warning("Please enter some text to classify.")
 
 def classify_text(text):
-    results = text_classifier(text)
-    return results
+    results = text_classifier(text, return_all_scores=True)
+    scores_list = results[0]
+    total_score = sum(score['score'] for score in scores_list)
+    labeled_probabilities = {}
+    for score in scores_list:
+        label = score['label']
+        probability = (score['score'] / total_score) * 100
+        labeled_probabilities[label] = probability
+    return labeled_probabilities
 
 def display_results(results):
     st.subheader("Prediction:")
 
-    for result in results:
-        label = result['label']
-        score = result['score']
-        score_percent = score * 100
-        st.write(f"{label}: {score_percent:.2f}%")
+    for label, probability in results.items():
+        st.write(f"{label.lower()}: {probability:.2f}%")
 
 if __name__ == "__main__":
     main()