Support NAYOSE for span entities using BM25 of entity name sub-tokens

app.py

@@ -5,13 +5,18 @@ from pathlib import Path
 import gradio as gr
 import torch
 import unidic_lite
+from bm25s.hf import BM25HF, TokenizerHF
 from fugashi import GenericTagger
 from transformers import AutoModelForPreTraining, AutoTokenizer
 
 
+ALIAS_SEP = "|"
+
 repo_id = "studio-ousia/luxe"
 revision = "ja-v0.3"
 
+nayose_repo_id = "studio-ousia/luxe-nayose-bm25"
+
 ignore_category_patterns = [
     r"\d+年",
     r"楽曲 [ぁ-ん]",
@@ -77,6 +82,11 @@ def normalize_text(text: str) -> str:
     return unicodedata.normalize("NFKC", text)
 
 
+bm25_tokenizer = TokenizerHF(lower=True, splitter=tokenizer.tokenize, stopwords=None, stemmer=None)
+bm25_tokenizer.load_vocab_from_hub("studio-ousia/luxe-nayose-bm25")
+bm25_retriever = BM25HF.load_from_hub("studio-ousia/luxe-nayose-bm25")
+
+
 def get_texts_from_file(file_path):
     texts = []
     with open(file_path) as f:
@@ -146,7 +156,7 @@ def get_predicted_entity_spans(
 
 
 def get_topk_entities_from_texts(
-    texts: list[str], k: int = 5, entity_span_sensitivity: float = 1.0
+    texts: list[str], k: int = 5, entity_span_sensitivity: float = 1.0, nayose_coef: float = 0.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]] = []
@@ -172,7 +182,17 @@ def get_topk_entities_from_texts(
         topk_category_entities.append([id2category_entity[id_] for id_ in topk_category_entity_ids.tolist()])
 
         if model_outputs.entity_logits is not None:
-            _, topk_span_entity_ids = model_outputs.entity_logits[0, :, :500000].topk(k)
+            span_entity_logits = model_outputs.entity_logits[0, :, :500000]
+
+            if nayose_coef > 0.0:
+                nayose_queries = ["ja:" + text[start:end] for start, end in entity_spans]
+                nayose_query_tokens = bm25_tokenizer.tokenize(nayose_queries)
+                nayose_scores = torch.vstack(
+                    [torch.from_numpy(bm25_retriever.get_scores(tokens)) for tokens in nayose_query_tokens]
+                )
+                span_entity_logits += nayose_coef * nayose_scores
+
+            _, topk_span_entity_ids = span_entity_logits.topk(k)
             topk_span_entities.append(
                 [[id2normal_entity[id_] for id_ in ids] for ids in topk_span_entity_ids.tolist()]
             )
@@ -211,6 +231,7 @@ with gr.Blocks() as demo:
     texts = gr.State([])
     topk = gr.State(5)
     entity_span_sensitivity = gr.State(1.0)
+    nayose_coef = gr.State(0.0)
     batch_entity_spans = gr.State([])
     topk_normal_entities = gr.State([])
     topk_category_entities = gr.State([])
@@ -230,20 +251,29 @@ with gr.Blocks() as demo:
     entity_span_sensitivity_input.change(
         fn=lambda val: val, inputs=entity_span_sensitivity_input, outputs=entity_span_sensitivity
     )
+    nayose_coef_input = gr.Slider(
+        minimum=0.0, maximum=2.0, value=0.0, step=0.1, label="Nayose Coefficient", interactive=True
+    )
+    nayose_coef_input.change(fn=lambda val: val, inputs=nayose_coef_input, outputs=nayose_coef)
 
     texts.change(
         fn=get_topk_entities_from_texts,
-        inputs=[texts, topk, entity_span_sensitivity],
+        inputs=[texts, topk, entity_span_sensitivity, nayose_coef],
         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],
+        inputs=[texts, topk, entity_span_sensitivity, nayose_coef],
         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],
+        inputs=[texts, topk, entity_span_sensitivity, nayose_coef],
+        outputs=[batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities],
+    )
+    nayose_coef.change(
+        fn=get_topk_entities_from_texts,
+        inputs=[texts, topk, entity_span_sensitivity, nayose_coef],
         outputs=[batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities],
     )