Update app.py

app.py

@@ -7,10 +7,10 @@ MODEL_NAME = "GhadeerALbadani/mmbert-Multilingual_detection_of_hate_speech"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
-labels = {
-    0: "Not Hate Speech",
-    1: "Hate Speech"
-}
+labels = [
+    "Not Hate Speech",
+    "Hate Speech"
+]
 
 def predict(text):
 
@@ -24,9 +24,12 @@ def predict(text):
     with torch.no_grad():
         outputs = model(**inputs)
 
-    prediction = torch.argmax(outputs.logits, dim=1).item()
+    probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
 
-    return labels[prediction]
+    return {
+        labels[0]: float(probs[0]),
+        labels[1]: float(probs[1])
+    }
 
 demo = gr.Interface(
     fn=predict,