Upload 2 files

app.py

@@ -0,0 +1,213 @@
+from pathlib import Path
+from typing import List, Tuple
+
+import gradio as gr
+import numpy as np
+import torch
+from sudachipy import dictionary
+from sudachipy import tokenizer as sudachi_tokenizer
+from transformers import AutoModelForCausalLM, PreTrainedTokenizer, T5Tokenizer
+
+# trained model
+model_dir = Path(__file__).parents[0] / "model"
+device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+tokenizer = T5Tokenizer.from_pretrained(model_dir)
+tokenizer.do_lower_case = True
+trained_model = AutoModelForCausalLM.from_pretrained(model_dir)
+trained_model.to(device)
+
+# baseline model
+baseline_model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt2-medium")
+baseline_model.to(device)
+
+sudachi_tokenizer_obj = dictionary.Dictionary().create()
+mode = sudachi_tokenizer.Tokenizer.SplitMode.C
+
+
+def sudachi_tokenize(input_text: str) -> List[str]:
+    morphemes = sudachi_tokenizer_obj.tokenize(input_text, mode)
+    return [morpheme.surface() for morpheme in morphemes]
+
+
+def calc_offsets(tokens: List[str]) -> List[int]:
+    offsets = [0]
+    for token in tokens:
+        offsets.append(offsets[-1] + len(token))
+    return offsets
+
+
+def distribute_surprisals_to_characters(
+    tokens2surprisal: List[Tuple[str, float]]
+) -> List[Tuple[str, float]]:
+    tokens2surprisal_by_character: List[Tuple[str, float]] = []
+    for token, surprisal in tokens2surprisal:
+        token_len = len(token)
+        for character in token:
+            tokens2surprisal_by_character.append((character, surprisal / token_len))
+    return tokens2surprisal_by_character
+
+
+def calculate_surprisals_by_character(
+    input_text: str, model: AutoModelForCausalLM, tokenizer: PreTrainedTokenizer
+) -> Tuple[float, List[Tuple[str, float]]]:
+    input_tokens = [
+        token.replace("▁", "")
+        for token in tokenizer.tokenize(input_text)
+        if token != "▁"
+    ]
+    input_ids = tokenizer.encode(
+        "<s>" + input_text, add_special_tokens=False, return_tensors="pt"
+    ).to(device)
+
+    logits = model(input_ids)["logits"].squeeze(0)
+
+    surprisals = []
+    for i in range(logits.shape[0] - 1):
+        if input_ids[0][i + 1] == 9:
+            continue
+        logit = logits[i]
+        prob = torch.softmax(logit, dim=0)
+        neg_logprob = -torch.log(prob)
+        surprisals.append(neg_logprob[input_ids[0][i + 1]].item())
+    mean_surprisal = np.mean(surprisals)
+
+    tokens2surprisal: List[Tuple[str, float]] = []
+    for token, surprisal in zip(input_tokens, surprisals):
+        tokens2surprisal.append((token, surprisal))
+
+    char2surprisal = distribute_surprisals_to_characters(tokens2surprisal)
+
+    return mean_surprisal, char2surprisal
+
+
+def aggregate_surprisals_by_offset(
+    char2surprisal: List[Tuple[str, float]], offsets: List[int]
+) -> List[Tuple[str, float]]:
+    tokens2surprisal = []
+    for i in range(len(offsets) - 1):
+        start = offsets[i]
+        end = offsets[i + 1]
+        surprisal = sum([surprisal for _, surprisal in char2surprisal[start:end]])
+        token = "".join([char for char, _ in char2surprisal[start:end]])
+        tokens2surprisal.append((token, surprisal))
+
+    return tokens2surprisal
+
+
+def highlight_token(token: str, score: float):
+    if score > 0:
+        html_color = "#%02X%02X%02X" % (
+            255,
+            int(255 * (1 - score)),
+            int(255 * (1 - score)),
+        )
+    else:
+        html_color = "#%02X%02X%02X" % (
+            int(255 * (1 + score)),
+            int(255 * (1 + score)),
+            255,
+        )
+    return '<span style="background-color: {}; color: black">{}</span>'.format(
+        html_color, token
+    )
+
+
+def create_highlighted_text(label: str, tokens2scores: List[Tuple[str, float]]):
+    highlighted_text: str = "<h2><b>" + label + "</b></h2>"
+    for token, score in tokens2scores:
+        highlighted_text += highlight_token(token, score)
+    return highlighted_text
+
+
+def normalize_surprisals(
+    tokens2surprisal: List[Tuple[str, float]], log_scale: bool = False
+) -> List[Tuple[str, float]]:
+    if log_scale:
+        surprisals = [np.log(surprisal) for _, surprisal in tokens2surprisal]
+    else:
+        surprisals = [surprisal for _, surprisal in tokens2surprisal]
+    min_surprisal = np.min(surprisals)
+    max_surprisal = np.max(surprisals)
+    surprisals = [
+        (surprisal - min_surprisal) / (max_surprisal - min_surprisal)
+        for surprisal in surprisals
+    ]
+    assert min(surprisals) >= 0
+    assert max(surprisals) <= 1
+    return [
+        (token, surprisal)
+        for (token, _), surprisal in zip(tokens2surprisal, surprisals)
+    ]
+
+
+def calculate_surprisal_diff(
+    tokens2surprisal: List[Tuple[str, float]],
+    baseline_tokens2surprisal: List[Tuple[str, float]],
+    scale: float = 100.0,
+):
+    diff_tokens2surprisal = [
+        (token, (surprisal - baseline_surprisal) * 100)
+        for (token, surprisal), (_, baseline_surprisal) in zip(
+            tokens2surprisal, baseline_tokens2surprisal
+        )
+    ]
+    return diff_tokens2surprisal
+
+
+def main(input_text: str) -> Tuple[str, str, str]:
+    mean_surprisal, char2surprisal = calculate_surprisals_by_character(
+        input_text, trained_model, tokenizer
+    )
+    offsets = calc_offsets(sudachi_tokenize(input_text))
+    tokens2surprisal = aggregate_surprisals_by_offset(char2surprisal, offsets)
+    tokens2surprisal = normalize_surprisals(tokens2surprisal)
+    highlighted_text = create_highlighted_text("学習後モデル", tokens2surprisal)
+
+    (
+        baseline_mean_surprisal,
+        baseline_char2surprisal,
+    ) = calculate_surprisals_by_character(input_text, baseline_model, tokenizer)
+    baseline_tokens2surprisal = aggregate_surprisals_by_offset(
+        baseline_char2surprisal, offsets
+    )
+    baseline_tokens2surprisal = normalize_surprisals(baseline_tokens2surprisal)
+    baseline_highlighted_text = create_highlighted_text(
+        "学習前モデル", baseline_tokens2surprisal
+    )
+
+    diff_tokens2surprisal = calculate_surprisal_diff(
+        tokens2surprisal, baseline_tokens2surprisal, 100.0
+    )
+    diff_highlighted_text = create_highlighted_text("学習前後の差分", diff_tokens2surprisal)
+    return (
+        baseline_highlighted_text,
+        highlighted_text,
+        diff_highlighted_text,
+    )
+
+
+if __name__ == "__main__":
+    demo = gr.Interface(
+        fn=main,
+        title="読みにくい箇所を検出するAI（デモ）",
+        description="テキストを入力すると、読みにくさに応じてハイライトされて出力されます。",
+        show_label=True,
+        inputs=gr.Textbox(
+            lines=5,
+            label="テキスト",
+            placeholder="ここにテキストを入力してください。",
+        ),
+        outputs=[
+            gr.HTML(label="学習前モデル", show_label=True),
+            gr.HTML(label="学習後モデル", show_label=True),
+            gr.HTML(label="学習前後の差分", show_label=True),
+        ],
+        examples=[
+            "太郎が二郎を殴った。",
+            "太郎が二郎に殴った。",
+            "サイエンスインパクトラボは、国立研究開発法人科学技術振興機構（JST）の「科学と社会」推進部が行う共創プログラムです。「先端の研究開発を行う研究者」と「社会課題解決に取り組むプレイヤー」が約３ヶ月に渡って共創活動を行います。",
+            "近年、ニューラル言語モデルが自然言語の統語知識をどれほど有しているかを、容認性判断課題を通して検証する研究が行われてきている。しかし、このような言語モデルの統語的評価を行うためのデータセットは、主に英語を中心とした欧米の諸言語を対象に構築されてきた。本研究では、既存のデータセットの問題点を克服しつつ、このようなデータセットが構築されてこなかった日本語を対象とした初めてのデータセットである JCoLA (JapaneseCorpus of Linguistic Acceptability) を構築した上で、それを用いた言語モデルの統語的評価を行った。",
+        ],
+    )
+
+    demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+torch==1.12.1
+transformers==4.20.0
+sentencepiece==0.1.97
+sudachipy
+sudachidict_core