Update app.py

app.py

@@ -3,6 +3,7 @@ pip.main(['install', 'torch'])
 pip.main(['install', 'transformers'])
 
 import torch
+import torch.nn as nn
 import gradio as gr
 import transformers
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -32,10 +33,18 @@ def inference(prompt):
     with torch.no_grad():
         logits = model(**inputs).logits
 
-    predicted_class_id = logits.argmax().item()
+    # for binary classification
+    sigmoid = nn.Sigmoid()
+    bi_prob = sigmoid(logits)
+
+    predicted_class_id = bi_prob.argmax().item()
     class_id = model.config.id2label[predicted_class_id]
 
-    return class_id
+    return {
+        "class_id": class_id,
+        "clean_prob": bi_prob[0][0].item(),
+        "unclean_prob": bi_prob[0][1].item()
+        }
 
 demo = gr.Interface(
     fn=inference,