Show scores for top-k entities

app.py

@@ -167,22 +167,34 @@ def get_topk_entities_from_texts(
         model_outputs = model(**tokenized_examples)
         token_spans = get_token_spans(tokenizer, text)
         entity_spans = get_predicted_entity_spans(model_outputs.ner_logits[0], token_spans, entity_span_sensitivity)
-        entity_spans = entity_spans[:tokenizer.max_entity_length]
+        entity_spans = entity_spans[: tokenizer.max_entity_length]
         batch_entity_spans.append(entity_spans)
 
         tokenized_examples = tokenizer(text, entity_spans=entity_spans or None, truncation=True, return_tensors="pt")
         model_outputs = model(**tokenized_examples)
 
         if model_outputs.topic_entity_logits is not None:
-            _, topk_normal_entity_ids = model_outputs.topic_entity_logits[0].topk(entity_k)
-            topk_normal_entities.append([id2normal_entity[id_] for id_ in topk_normal_entity_ids.tolist()])
+            topk_normal_entity_scores, topk_normal_entity_ids = model_outputs.topic_entity_logits[0].topk(entity_k)
+            topk_normal_entities.append(
+                [
+                    f"{id2normal_entity[id_]} ({score:.3f})"
+                    for score, id_ in zip(topk_normal_entity_scores, topk_normal_entity_ids.tolist())
+                ]
+            )
         else:
             topk_normal_entities.append([])
 
         if model_outputs.topic_category_logits is not None:
             model_outputs.topic_category_logits[:, ignore_category_entity_ids] = float("-inf")
-            _, topk_category_entity_ids = model_outputs.topic_category_logits[0].topk(category_k)
-            topk_category_entities.append([id2category_entity[id_] for id_ in topk_category_entity_ids.tolist()])
+            topk_category_entity_scores, topk_category_entity_ids = model_outputs.topic_category_logits[0].topk(
+                category_k
+            )
+            topk_category_entities.append(
+                [
+                    f"{id2category_entity[id_]} ({score:.3f})"
+                    for score, id_ in zip(topk_category_entity_scores, topk_category_entity_ids.tolist())
+                ]
+            )
         else:
             topk_category_entities.append([])
 
@@ -197,9 +209,12 @@ def get_topk_entities_from_texts(
                 )
                 span_entity_logits += nayose_coef * nayose_scores
 
-            _, topk_span_entity_ids = span_entity_logits.topk(entity_k)
+            topk_span_entity_scores, topk_span_entity_ids = span_entity_logits.topk(entity_k)
             topk_span_entities.append(
-                [[id2normal_entity[id_] for id_ in ids] for ids in topk_span_entity_ids.tolist()]
+                [
+                    [f"{id2normal_entity[id_]} ({score:.3f})" for score, id_ in zip(scores, ids)]
+                    for scores, ids in zip(topk_span_entity_scores, topk_span_entity_ids.tolist())
+                ]
             )
         else:
             topk_span_entities.append([])