return top 3 predictions

app.py

@@ -1,4 +1,5 @@
 import re
+import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
@@ -10,8 +11,9 @@ def tag_job_description(job_description):
     job_description = re.sub(r'[^a-zA-Z]', ' ', job_description)
     inputs = tokenizer(job_description, truncation=True, max_length=512, return_tensors="pt")
     outputs = model(**inputs)
-    predicted_tag = model.config.id2label[outputs.logits.argmax().item()]
-    return predicted_tag
+    topk_values, topk_indices = torch.topk(outputs.logits, 3)
+    predicted_tags = [model.config.id2label[i] for i in topk_indices[0]]
+    return predicted_tags
 
 
 demo = gr.Interface(fn=tag_job_description, inputs="text", outputs="text")