prediction probability feature

app.py

@@ -72,9 +72,16 @@ id2label = {0: "REJECTED", 1: "ACCEPTED"}
 # when submit button clicked, run the model and get result
 if st.button("Submit"):
     with torch.no_grad():
-        logits = model(**inputs).logits
+        outputs = model(**inputs)
+        probability = torch.nn.functional.softmax(outputs.logits, dim=1)
 
-    predicted_class_id = logits.argmax().item()
+    predicted_class_id = probability.argmax().item()
     pred_label = id2label[predicted_class_id]
     st.title("Predicted Patentability")
+    if probability[0][0] > probability[0][1]:
+        st.write("Rejection Score:")
+        st.write(probability[0][0])
+    else:
+        st.write("Acceptance Score:")
+        st.write(probability[0][1])
     st.write(pred_label)