also display second most probable label in table

app.py

@@ -103,9 +103,11 @@ def predict(text, model_id, tokenizer_id):
     return probs
 
 
-def get_most_probable_label(probs):
-    label = id2label[probs.argmax()]
-    probability = f"{round(100 * probs.max(), 2)}%"
+def get_most_probable_label(probs, idx=1):
+    sorted_indices = probs.argsort()[::-1]
+    selected_idx = sorted_indices[idx - 1]
+    label = id2label[selected_idx]
+    probability = f"{round(100 * probs[selected_idx], 2)}%"
     return label, probability
 
 
@@ -206,8 +208,9 @@ def predict_wrapper(text, language):
     results_heatmap = []
     for sentence in tqdm(sentences):
         probs = predict(sentence, model_id, tokenizer_id)
-        label, probability = get_most_probable_label(probs)
-        results.append([sentence, label, probability])
+        label1, probability1 = get_most_probable_label(probs, 1)
+        label2, probability2 = get_most_probable_label(probs, 2)
+        results.append([sentence, label1, probability1, label2, probability2])
         results_heatmap.append({"sentence": sentence, "emotions": probs})
 
     # let's see...
@@ -256,8 +259,8 @@ with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column(scale=7):
             result_table = gr.Dataframe(
-                headers=["Sentence", "Prediction", "Confidence"],
-                column_widths=["65%", "25%", "10%"],
+                headers=["Sentence", "Prediction (1)", "Confidence (1)", "Prediction (2)", "Confidence (2)"],
+                column_widths=["40%", "20%", "10%", "20%", "10%"],
                 wrap=True,  # important
             )
         with gr.Column(scale=3):