Reset model and tokenizer when the demo is reloaded

app.py

@@ -33,9 +33,6 @@ ignore_category_patterns = [
     r"各年の",
 ]
 
-model = AutoModelForPreTraining.from_pretrained(repo_id, revision=revision, trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained(repo_id, revision=revision, trust_remote_code=True)
-
 
 class MecabTokenizer:
     def __init__(self):
@@ -65,11 +62,6 @@ 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: str | None):
     texts = []
     if file_path is not None:
@@ -104,7 +96,7 @@ 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]]:
+def get_token_spans(tokenizer, text: str) -> list[tuple[int, int]]:
     token_spans = []
     end = 0
     for token in tokenizer.tokenize(text):
@@ -147,12 +139,15 @@ def get_predicted_entity_spans(
 
 
 def get_topk_entities_from_texts(
+    models,
     texts: list[str],
     k: int = 5,
     entity_span_sensitivity: float = 1.0,
     nayose_coef: float = 1.0,
     entities_are_replaced: bool = False,
 ) -> tuple[list[list[tuple[int, int]]], list[list[str]], list[list[str]], list[list[list[str]]]]:
+    model, tokenizer, bm25_tokenizer, bm25_retriever = models
+
     batch_entity_spans: list[list[tuple[int, int]]] = []
     topk_normal_entities: list[list[str]] = []
     topk_category_entities: list[list[str]] = []
@@ -178,7 +173,7 @@ def get_topk_entities_from_texts(
     for text in texts:
         tokenized_examples = tokenizer(text, return_tensors="pt")
         model_outputs = model(**tokenized_examples)
-        token_spans = get_token_spans(text)
+        token_spans = get_token_spans(tokenizer, text)
         entity_spans = get_predicted_entity_spans(model_outputs.ner_logits[0], token_spans, entity_span_sensitivity)
         batch_entity_spans.append(entity_spans)
 
@@ -223,7 +218,9 @@ def get_selected_entity(evt: gr.SelectData):
     return evt.value[0]
 
 
-def get_similar_entities(query_entity: str, k: int = 10) -> list[str]:
+def get_similar_entities(models, query_entity: str, k: int = 10) -> list[str]:
+    model, tokenizer, _, _ = models
+
     query_entity_id = tokenizer.entity_vocab[query_entity]
 
     id2normal_entity = {
@@ -282,11 +279,14 @@ def get_new_entity_text_pairs_from_file(file_path: str | None) -> list[list[str]
 
 
 def replace_entities(
+    models,
     new_entity_text_pairs: list[tuple[str, str]],
     new_num_category_entities: int = 0,
     new_entity_counts: list[int] | None = None,
     new_padding_idx: int = 0,
 ) -> True:
+    model, tokenizer, bm25_tokenizer, bm25_retriever = models
+
     gr.Info("トークナイザのエンティティの語彙を置き換えています...", duration=5)
     new_entity_tokens = ENTITY_SPECIAL_TOKENS + [entity for entity, _ in new_entity_text_pairs]
 
@@ -367,6 +367,18 @@ def replace_entities(
 
 
 with gr.Blocks() as demo:
+    model = AutoModelForPreTraining.from_pretrained(repo_id, revision=revision, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(repo_id, revision=revision, trust_remote_code=True)
+    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")
+
+    # Hint: gr.State に callable を渡すと、それが state の初期値を設定するための関数とみなされて
+    # __call__ が引数なしで実行されてしまうため、gr.State の引数に model や tokenizer を単体で渡すとエラーになってしまう。
+    # ここでは、モデル一式のタプル（callable でない）を渡すことで、そのようなエラーを回避している。
+    # cf. https://www.gradio.app/docs/gradio/state#param-state-value
+    models = gr.State((model, tokenizer, bm25_tokenizer, bm25_retriever))
+
     texts = gr.State([])
 
     entities_are_replaced = gr.State(False)
@@ -407,7 +419,9 @@ with gr.Blocks() as demo:
     new_entity_text_pairs_file.change(
         fn=get_new_entity_text_pairs_from_file, inputs=new_entity_text_pairs_file, outputs=new_entity_text_pairs_input
     )
-    replace_entity_button.click(fn=replace_entities, inputs=new_entity_text_pairs_input, outputs=entities_are_replaced)
+    replace_entity_button.click(
+        fn=replace_entities, inputs=[models, new_entity_text_pairs_input], outputs=entities_are_replaced
+    )
 
     with gr.Accordion(label="ハイパーパラメータ", open=False):
         topk_input = gr.Number(5, label="エンティティ件数", interactive=True)
@@ -428,25 +442,30 @@ with gr.Blocks() as demo:
 
     texts.change(
         fn=get_topk_entities_from_texts,
-        inputs=[texts, topk, entity_span_sensitivity, nayose_coef, entities_are_replaced],
+        inputs=[models, texts, topk, entity_span_sensitivity, nayose_coef, entities_are_replaced],
         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, nayose_coef, entities_are_replaced],
+        inputs=[models, texts, topk, entity_span_sensitivity, nayose_coef, entities_are_replaced],
         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, nayose_coef, entities_are_replaced],
+        inputs=[models, texts, topk, entity_span_sensitivity, nayose_coef, entities_are_replaced],
         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, entities_are_replaced],
+        inputs=[models, texts, topk, entity_span_sensitivity, nayose_coef, entities_are_replaced],
+        outputs=[batch_entity_spans, topk_normal_entities, topk_category_entities, topk_span_entities],
+    )
+
+    entities_are_replaced.change(
+        fn=get_topk_entities_from_texts,
+        inputs=[models, texts, topk, entity_span_sensitivity, nayose_coef, entities_are_replaced],
         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("## 出力エンティティ")
@@ -499,7 +518,7 @@ with gr.Blocks() as demo:
     #     gr.Markdown("---")
     #     gr.Markdown("## 選択されたエンティティの類似エンティティ")
 
-    # selected_entity.change(fn=get_similar_entities, inputs=selected_entity, outputs=similar_entities)
+    # selected_entity.change(fn=get_similar_entities, inputs=[models, selected_entity], outputs=similar_entities)
 
     # @gr.render(inputs=[selected_entity, similar_entities])
     # def render_similar_entities(selected_entity, similar_entities):