no comments need more coffee

app.py

@@ -12,7 +12,7 @@ def predict(text):
 
     probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
     predicted_class = torch.argmax(probs, dim=-1).item()
-    return label_dict[str(predicted_class)], probs[0][predicted_class].item()
+    return id2label[predicted_class], probs[0][predicted_class].item()
 
 if __name__ == '__main__':
     model_path = "Dunateo/roberta-cwe-classifier-kelemia-v0.2"
@@ -24,11 +24,14 @@ if __name__ == '__main__':
     # get the dict file
     label_dict_file = hf_hub_download(repo_id=model_path, filename="label_dict.json")
 
-    global label_dict
+
     with open(label_dict_file, "r") as f:
         content = f.read()
     label_dict = json.loads(content)
 
+    global id2label
+    id2label = {v: k for k, v in label_dict.items()}
+
     # gradio specific to create an IHM
     iface = gr.Interface(
         fn=predict,