Added confidence scores

app.py

@@ -5,6 +5,7 @@ from typing import List, Tuple
 import gradio as gr
 import pandas as pd
 import spacy
+import torch
 from transformers import AutoModelForTokenClassification, AutoTokenizer
 
 try:
@@ -36,6 +37,7 @@ def predict(text, nlp, logger=None) -> Tuple[List[str], List[str]]:
 
     i_token = 0
     labels = []
+    scores = []
     for off, is_special_token, pred in zip(
         input_tokens["offset_mapping"][0],
         input_tokens["special_tokens_mask"][0],
@@ -47,17 +49,21 @@ def predict(text, nlp, logger=None) -> Tuple[List[str], List[str]]:
         if logger is not None:
             logger.info("{}, {}, {}".format(off, tokens[i_token], label))
         labels.append(label)
+        scores.append(
+            "{:.2f}".format(100 * float(torch.softmax(pred, dim=-1).detach().max()))
+        )
         i_token += 1
 
-    return tokens, labels
+    return tokens, labels, scores
 
 
 def text_analysis(text):
-    tokens, labels = predict(text, nlp, logger)
+    tokens, labels, scores = predict(text, nlp, logger)
     pos_count = pd.DataFrame(
         {
             "token": tokens,
             "etiqueta": labels,
+            "confiança": scores,
         }
     )
     pos_tokens = []