Add span-based entity linking

app.py

@@ -1,4 +1,8 @@
+from pathlib import Path
+
 import gradio as gr
+import unidic_lite
+from fugashi import GenericTagger
 from transformers import AutoModelForPreTraining, AutoTokenizer
 
 
@@ -26,6 +30,30 @@ normal_entity_embeddings = entity_embeddings[:num_normal_entities]
 category_entity_embeddings = entity_embeddings[num_normal_entities:]
 
 
+class MecabTokenizer:
+    def __init__(self):
+        unidic_dir = unidic_lite.DICDIR
+        mecabrc_file = Path(unidic_dir, "mecabrc")
+        mecab_option = f"-d {unidic_dir} -r {mecabrc_file}"
+        self.tagger = GenericTagger(mecab_option)
+
+    def __call__(self, text: str) -> list[tuple[str, str, tuple[int, int]]]:
+        outputs = []
+
+        end = 0
+        for node in self.tagger(text):
+            word = node.surface.strip()
+            pos = node.feature[0]
+            start = text.index(word, end)
+            end = start + len(word)
+            outputs.append((word, pos, (start, end)))
+
+        return outputs
+
+
+mecab_tokenizer = MecabTokenizer()
+
+
 def get_texts_from_file(file_path):
     texts = []
     with open(file_path) as f:
@@ -37,12 +65,33 @@ def get_texts_from_file(file_path):
     return texts
 
 
-def get_topk_entities_from_texts(texts: list[str], k: int = 5) -> tuple[list[list[str]], list[list[str]]]:
-    topk_normal_entities = []
-    topk_category_entities = []
+def get_noun_spans_from_text(text: str) -> list[tuple[int, int]]:
+    last_pos = None
+    noun_spans = []
+
+    for word, pos, (start, end) in mecab_tokenizer(text):
+        if pos == "名詞":
+            if len(noun_spans) > 0 and last_pos == "名詞":
+                noun_spans[-1] = (noun_spans[-1][0], end)
+            else:
+                noun_spans.append((start, end))
+
+        last_pos = pos
+
+    return noun_spans
+
+
+def get_topk_entities_from_texts(
+    texts: list[str], k: int = 5
+) -> tuple[list[list[str]], list[list[str]], list[list[list[str]]]]:
+    topk_normal_entities: list[list[str]] = []
+    topk_category_entities: list[list[str]] = []
+    topk_span_entities: list[list[list[str]]] = []
 
     for text in texts:
-        tokenized_examples = tokenizer(text, return_tensors="pt")
+        noun_spans = get_noun_spans_from_text(text)
+
+        tokenized_examples = tokenizer(text, entity_spans=noun_spans, return_tensors="pt")
         model_outputs = model(**tokenized_examples)
 
         _, topk_normal_entity_ids = model_outputs.topic_entity_logits[0].topk(k)
@@ -51,7 +100,10 @@ def get_topk_entities_from_texts(texts: list[str], k: int = 5) -> tuple[list[lis
         _, topk_category_entity_ids = model_outputs.topic_category_logits[0].topk(k)
         topk_category_entities.append([id2category_entity[id_] for id_ in topk_category_entity_ids.tolist()])
 
-    return topk_normal_entities, topk_category_entities
+        _, topk_span_entity_ids = model_outputs.entity_logits[0, :, :500000].topk(k)
+        topk_span_entities.append([[id2normal_entity[id_] for id_ in ids] for ids in topk_span_entity_ids.tolist()])
+
+    return topk_normal_entities, topk_category_entities, topk_span_entities
 
 
 def get_selected_entity(evt: gr.SelectData):
@@ -80,29 +132,52 @@ with gr.Blocks() as demo:
     texts = gr.State([])
     topk_normal_entities = gr.State([])
     topk_category_entities = gr.State([])
+    topk_span_entities = gr.State([])
     selected_entity = gr.State()
     similar_entities = gr.State([])
 
     text_input = gr.Textbox(label="Input Text")
-    texts_file = gr.File(label="Input texts")
+    texts_file = gr.File(label="Input Texts")
 
     text_input.change(fn=lambda text: [text], inputs=text_input, outputs=texts)
     texts_file.change(fn=get_texts_from_file, inputs=texts_file, outputs=texts)
-    texts.change(fn=get_topk_entities_from_texts, inputs=texts, outputs=[topk_normal_entities, topk_category_entities])
+    texts.change(
+        fn=get_topk_entities_from_texts,
+        inputs=texts,
+        outputs=[topk_normal_entities, topk_category_entities, topk_span_entities],
+    )
 
     gr.Markdown("---")
     gr.Markdown("## 出力エンティティ")
 
-    @gr.render(inputs=[texts, topk_normal_entities, topk_category_entities])
-    def render_topk_entities(texts, topk_normal_entities, topk_category_entities):
-        for text, normal_entities, category_entities in zip(texts, topk_normal_entities, topk_category_entities):
-            gr.Textbox(text, label="Text")
-            entities = gr.Dataset(
-                label="Entities",
-                components=["text"],
-                samples=[[entity] for entity in normal_entities + category_entities],
+    @gr.render(inputs=[texts, topk_normal_entities, topk_category_entities, topk_span_entities])
+    def render_topk_entities(texts, topk_normal_entities, topk_category_entities, topk_span_entities):
+        for text, normal_entities, category_entities, span_entities in zip(
+            texts, topk_normal_entities, topk_category_entities, topk_span_entities
+        ):
+            gr.HighlightedText(
+                value=[(word, pos if pos == "名詞" else None) for word, pos, _ in mecab_tokenizer(text)],
+                color_map={"名詞": "green"},
+                show_legend=True,
+                combine_adjacent=True,
+                adjacent_separator=" ",
+                label="Text",
             )
-            entities.select(fn=get_selected_entity, outputs=selected_entity)
+
+            # gr.Textbox(text, label="Text")
+            gr.Dataset(
+                label="Topic Entities", components=["text"], samples=[[entity] for entity in normal_entities]
+            ).select(fn=get_selected_entity, outputs=selected_entity)
+            gr.Dataset(
+                label="Topic Categories", components=["text"], samples=[[entity] for entity in category_entities]
+            ).select(fn=get_selected_entity, outputs=selected_entity)
+
+            noun_spans = get_noun_spans_from_text(text)
+            nouns = [text[start:end] for start, end in noun_spans]
+            for noun, entities in zip(nouns, span_entities):
+                gr.Dataset(
+                    label=f"Span Entities for {noun}", components=["text"], samples=[[entity] for entity in entities]
+                ).select(fn=get_selected_entity, outputs=selected_entity)
 
         gr.Markdown("---")
         gr.Markdown("## 選択されたエンティティの類似エンティティ")