Use ja-v0.3 model, introduce entity span sensitivity

app.py

@@ -2,13 +2,14 @@ import re
 from pathlib import Path
 
 import gradio as gr
+import torch
 import unidic_lite
 from fugashi import GenericTagger
 from transformers import AutoModelForPreTraining, AutoTokenizer
 
 
 repo_id = "studio-ousia/luxe"
-revision = "ja-v0.2"
+revision = "ja-v0.3"
 
 ignore_category_patterns = [
     r"\d+年",
@@ -98,17 +99,63 @@ def get_noun_spans_from_text(text: str) -> list[tuple[int, int]]:
     return noun_spans
 
 
+def get_token_spans(text: str) -> list[tuple[int, int]]:
+    token_spans = []
+    end = 0
+    for token in tokenizer.tokenize(text):
+        token = token.removeprefix("##")
+        start = text.index(token, end)
+        end = start + len(token)
+        token_spans.append((start, end))
+
+    return [(0, 0)] + token_spans + [(end, end)]  # count for "[CLS]" and "[SEP]"
+
+
+def get_predicted_entity_spans(
+    ner_logits: torch.Tensor, token_spans: list[tuple[int, int]], entity_span_sensitivity: float = 1.0
+) -> list[tuple[int, int]]:
+    length = ner_logits.size(-1)
+    assert ner_logits.size() == (length, length)  # not batched
+
+    ner_probs = torch.sigmoid(ner_logits).triu()
+    probs_sorted, sort_idxs = ner_probs.flatten().sort(descending=True)
+
+    predicted_entity_spans = []
+    for p, i in zip(probs_sorted, sort_idxs.tolist()):
+        if p < 10.0 ** (-1.0 * entity_span_sensitivity):
+            break
+
+        start_idx = i // length
+        end_idx = i % length
+
+        start = token_spans[start_idx][0]
+        end = token_spans[end_idx][1]
+
+        for ex_start, ex_end in predicted_entity_spans:
+            if not (start < end <= ex_start or ex_end <= start < end):
+                break
+        else:
+            predicted_entity_spans.append((start, end))
+
+    return sorted(predicted_entity_spans)
+
+
 def get_topk_entities_from_texts(
-    texts: list[str], k: int = 5
-) -> tuple[list[list[str]], list[list[str]], list[list[list[str]]]]:
+    texts: list[str], k: int = 5, entity_span_sensitivity: float = 1.0
+) -> tuple[list[list[tuple[int, int]]], list[list[str]], list[list[str]], list[list[list[str]]]]:
+    batch_entity_spans: list[list[tuple[int, int]]] = []
     topk_normal_entities: list[list[str]] = []
     topk_category_entities: list[list[str]] = []
     topk_span_entities: list[list[list[str]]] = []
 
     for text in texts:
-        noun_spans = get_noun_spans_from_text(text)
+        tokenized_examples = tokenizer(text, return_tensors="pt")
+        model_outputs = model(**tokenized_examples)
+        token_spans = get_token_spans(text)
+        entity_spans = get_predicted_entity_spans(model_outputs.ner_logits[0], token_spans, entity_span_sensitivity)
+        batch_entity_spans.append(entity_spans)
 
-        tokenized_examples = tokenizer(text, entity_spans=noun_spans, return_tensors="pt")
+        tokenized_examples = tokenizer(text, entity_spans=entity_spans or None, return_tensors="pt")
         model_outputs = model(**tokenized_examples)
 
         model_outputs.topic_category_logits[:, ignore_category_entity_ids] = float("-inf")
@@ -119,10 +166,15 @@ def get_topk_entities_from_texts(
         _, 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()])
 
-        _, 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()])
+        if model_outputs.entity_logits is not None:
+            _, 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()]
+            )
+        else:
+            topk_span_entities.append([])
 
-    return topk_normal_entities, topk_category_entities, topk_span_entities
+    return batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities
 
 
 def get_selected_entity(evt: gr.SelectData):
@@ -152,6 +204,9 @@ with gr.Blocks() as demo:
     gr.Markdown("## テキスト（直接入力またはファイルアップロード）")
 
     texts = gr.State([])
+    topk = gr.State(5)
+    entity_span_sensitivity = gr.State(1.0)
+    batch_entity_spans = gr.State([])
     topk_normal_entities = gr.State([])
     topk_category_entities = gr.State([])
     topk_span_entities = gr.State([])
@@ -159,31 +214,60 @@ with gr.Blocks() as demo:
     similar_entities = gr.State([])
 
     text_input = gr.Textbox(label="Input Text")
-    texts_file = gr.File(label="Input Texts")
-
     text_input.change(fn=lambda text: [text], inputs=text_input, outputs=texts)
+    texts_file = gr.File(label="Input Texts")
     texts_file.change(fn=get_texts_from_file, inputs=texts_file, outputs=texts)
+    topk_input = gr.Number(5, label="Top K", interactive=True)
+    topk_input.change(fn=lambda val: val, inputs=topk_input, outputs=topk)
+    entity_span_sensitivity_input = gr.Slider(
+        minimum=0.1, maximum=5.0, value=1.0, step=0.1, label="Entity Span Sensitivity", interactive=True
+    )
+    entity_span_sensitivity_input.change(
+        fn=lambda val: val, inputs=entity_span_sensitivity_input, outputs=entity_span_sensitivity
+    )
+
     texts.change(
         fn=get_topk_entities_from_texts,
-        inputs=texts,
-        outputs=[topk_normal_entities, topk_category_entities, topk_span_entities],
+        inputs=[texts, topk, entity_span_sensitivity],
+        outputs=[batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities],
+    )
+    topk.change(
+        fn=get_topk_entities_from_texts,
+        inputs=[texts, topk, entity_span_sensitivity],
+        outputs=[batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities],
     )
+    entity_span_sensitivity.change(
+        fn=get_topk_entities_from_texts,
+        inputs=[texts, topk, entity_span_sensitivity],
+        outputs=[batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities],
+    )
+
+    topk_input.change(inputs=topk_input, outputs=topk)
 
     gr.Markdown("---")
     gr.Markdown("## 出力エンティティ")
 
-    @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.render(inputs=[texts, batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities])
+    def render_topk_entities(
+        texts, batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities
+    ):
+        for text, entity_spans, normal_entities, category_entities, span_entities in zip(
+            texts, batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities
         ):
+            highlighted_text_value = []
+            cur = 0
+            for start, end in entity_spans:
+                if cur < start:
+                    highlighted_text_value.append((text[cur:start], None))
+
+                highlighted_text_value.append((text[start:end], "Entity"))
+                cur = end
+
+            if cur < len(text):
+                highlighted_text_value.append((text[cur:], None))
+
             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",
+                value=highlighted_text_value, color_map={"Entity": "green"}, combine_adjacent=False, label="Text"
             )
 
             # gr.Textbox(text, label="Text")
@@ -194,11 +278,12 @@ with gr.Blocks() as demo:
                 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):
+            span_texts = [text[start:end] for start, end in entity_spans]
+            for span_text, entities in zip(span_texts, span_entities):
                 gr.Dataset(
-                    label=f"Span Entities for {noun}", components=["text"], samples=[[entity] for entity in entities]
+                    label=f"Span Entities for {span_text}",
+                    components=["text"],
+                    samples=[[entity] for entity in entities],
                 ).select(fn=get_selected_entity, outputs=selected_entity)
 
         gr.Markdown("---")