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]'
Tokenizerクラス
基本的な処理はPreTrainedTokenizerBaseクラスにて実装されているのでこちらを参照