Open3

huggingface transformers メモ

MLMタスク(日本語)

実行環境

  • MacOS
  • Python 3.9.9
  • transformers==4.12.5
  • torch==1.10.0

スクリプト

from transformers import BertJapaneseTokenizer, ElectraForPreTraining

tokenizer = BertJapaneseTokenizer.from_pretrained(
    'Cinnamon/electra-small-japanese-generator',
    mecab_kwargs={"mecab_option": "-d /usr/local/lib/mecab/dic/mecab-ipadic-neologd"})

text = f"私は{tokenizer.mask_token}へ行きます。"
encoded_input = tokenizer.encode(text, return_tensors="pt")
# tokenizer.decode(encoded_input.squeeze(0))
# '[CLS] 私 は [MASK] へ 行き ます 。 [SEP]'
mask_token_index = torch.where(encoded_input == tokenizer.mask_token_id)[1]

# 推論
token_logits = generator(encoded_input)[0]
mask_token_logits  = token_logits[0, mask_token_index,:]

# 上位5トークンを出力
top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()
for token in top_5_tokens:
    print(text.replace(tokenizer.mask_token, f"[{tokenizer.decode([token])}]"))

MLM問題の作り方

from transformers import BertJapaneseTokenizer, BertForMaskedLM
import torch

tokenizer = BertJapaneseTokenizer.from_pretrained(
      "cl-tohoku/bert-base-japanese-whole-word-masking")
model = BertForMaskedLM.from_pretrained(
      "cl-tohoku/bert-base-japanese-whole-word-masking")

inputs = tokenizer("問題を試しに作ってみましょう!", return_tensors="pt")

# [0,1)の一様分布を生成
rand = torch.rand(inputs.input_ids.shape) < 0.15
# [CLS], [SEP]以外でmaskする場所を決める
mask_arr = rand * (inputs.input_ids != tokenizer.cls_token_id) * (inputs.input_ids != tokenizer.sep_token_id)
# Trueのindexをリストアップ
selection = torch.flatten(mask_arr.squeeze(0).nonzero()).tolist()
# 入力をmaskする
inputs.input_ids[0, selection] = tokenizer.mask_token_id
tokenizer.decode(inputs.input_ids.squeeze(0))
# '[CLS] 問題 [MASK] [MASK] に 作っ て み ましょ う! [SEP]'
ログインするとコメントできます