the model classifier now returns probability distribution over all the labels

app.py

@@ -1,6 +1,7 @@
 from transformers import AutoModelForSequenceClassification,AutoTokenizer
-import gradio as gr
+from torch.nn.functional import softmax
 import torch
+import gradio as gr
 
 model_name="nebiyu29/hate_classifier"
 tokenizer=AutoTokenizer.from_pretrained(model_name)
@@ -15,8 +16,11 @@ def model_classifier(text):
             return f"the input text is {text}"
         else:
             encoded_input=tokenizer(text) #this is where the encoding happens
-            scores=model(encoded)[0]  #this is is the score for rach values
-            return scores
+            logits=model(**encoded)  #this is the logits of the labels
+            probs_label=softmax(logits,dim=-1)  #turning the probability distribution into normalize form
+            id2label=model.config.id2label
+            return_probs={id2label[i]:probs.item() for i,probs in enumerate(probs_label[0])}
+            return return_probs