fix:show all prob test

app.py

@@ -24,6 +24,7 @@ emoji_to_emotion = {
 # Function to make predictions
 def predict_sentiment(text):
     start_time = timer()
+    
     inputs = model.tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     input_ids = inputs["input_ids"]
     attention_mask = inputs["attention_mask"]
@@ -32,13 +33,19 @@ def predict_sentiment(text):
         outputs = model(input_ids=input_ids, attention_mask=attention_mask)
 
     logits = outputs.logits
-    _, predicted_class = torch.max(logits, dim=1)
-    pred_time = round(timer() - start_time, 5)
+    probabilities = torch.nn.functional.softmax(logits, dim=1)
     
     # Map predicted class to emoji
-    result = emoji_to_emotion[predicted_class.item()]
+    predicted_class = torch.argmax(logits, dim=1).item()
+    result = emoji_to_emotion[predicted_class]
+    
+    # Create a dictionary of class probabilities
+    class_probabilities = {emoji_to_emotion[i]: float(probabilities[0, i]) for i in range(len(emoji_to_emotion))}
+    
+    # Calculate prediction time
+    pred_time = round(timer() - start_time, 5)
     
-    return result,pred_time
+    return class_probabilities, pred_time
 
 # Create title, description and article strings
 title = "Emoji-aware Sentiment Analysis using Roberta Model"
@@ -50,7 +57,7 @@ article = "Sentiment Analysis, also known as opinion mining, is a branch of Natu
 iface = gr.Interface(
     fn=predict_sentiment,
     inputs="text",
-    outputs=[gr.Label(num_top_classes=7, label="Predictions"), # what are the outputs?
+    outputs=[gr.Label(num_top_classes=7, label="Predictions"),
                              gr.Number(label="Prediction time (s)")],
     title=title,
     description=description,