ModernBERTで日本語文章の穴埋め

世は大LLM時代ですが、BERT系のモデルも、特定の用途には適しており改めて学習する価値があると思います。BERT系のモデルの主力の一つであるModernBERTについて、日本語対応しているモデルを動作させてみます。

理解の補助のために、すべての行に対してLLMにコメントアウトさせています。

from transformers import AutoTokenizer, AutoModelForMaskedLM  # トークナイザーとマスク付き言語モデルを使う
import torch  # PyTorchテンソル操作と推論に利用

model_name = "llm-jp/llm-jp-modernbert-base"  # 利用するBERTモデル名
tokenizer = AutoTokenizer.from_pretrained(model_name)  # 事前学習済みのトークナイザーをロード
model = AutoModelForMaskedLM.from_pretrained(model_name)  # マスク付きLMモデル本体をロード

mask = tokenizer.mask_token  # モデルが使う[MASK]トークン文字列
print("mask_token:", mask)  # マスクトークンを確認出力

text = f"日本の首都は{mask}です。"  # マスク穴埋め対象の入力文
inputs = tokenizer(text, return_tensors="pt")  # トークン化しテンソル化
print(inputs)
outputs = model(**inputs)  # モデルで前向き推論して各トークンのロジットを得る
print(tokenizer.mask_token_id)
mask_positions = (inputs["input_ids"][0] == tokenizer.mask_token_id).nonzero(
    as_tuple=False).flatten()  # 入力中の[MASK]位置をインデックスとして取得
print("mask_positions:", mask_positions.tolist())  # マスク位置を確認出力

k = 10  # 上位何件の候補を出すか
logits = outputs.logits[0, mask_positions[0]]  # マスク位置に対応する語彙ごとのロジット
topk = torch.topk(logits, k)  # ロジット上位k件を抽出

for score, tid in zip(topk.values, topk.indices):  # 候補ごとにスコアとトークンIDを取り出す
    print(tokenizer.decode([tid.item()]), float(score))  # デコードして単語とスコアを表示