huggingface Tokenizer の tokenize, encode, encode_plus などの違い

huggingface ライブラリを使っていると tokenize, encode, encode_plus などがよく出てきて混乱しがちなので改めてまとめておきます。

tokenize

言語モデルの vocabulary にしたがって入力文を分かち書きします。

$ pip install transformers[ja]

tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")
tokenizer.tokenize("私は元気です。")

['私', 'は', '元気', 'です', '。']

なお、convert_tokens_to_ids を使うと、語彙に対応する token id の配列に変換することができます。

tokenizer.convert_tokens_to_ids(['私', 'は', '元気', 'です', '。'])

[1325, 9, 12453, 2992, 8]

encode

先に述べた tokenize と convert_tokens_to_ids のステップを同時に行い、入力文から直接 token id の配列に変換します。

tokenizer.encode("私は元気です。")

[2, 1325, 9, 12453, 2992, 8, 3]

なお、現在の transformers ライブラリ (v4.11.3) ではこの encode の出力に関して、デフォルトの add_special_tokens オプションにより、配列の先頭と末尾にに特殊トークンを挿入します（これは言語モデルの事前学習の時点でそうされているので推奨操作です）。
BERT であれば、特殊トークンとは冒頭の [CLS] と文末の [SEP] のトークンのことで、これはそれぞれ id が 2, 3 に対応します。
少し古いバージョンですとこの挙動がデフォルトで入らないので、add_special_tokens=True という内容を引数に入れる必要があります。

# 古い transformers の場合
tokenizer.encode("私は元気です。", add_special_tokens=True)

encode_plus

先に述べた encode に加え、言語モデルの入力として必要な他の id を一緒に出力します。BERT であれば token type id と attention mask を一緒に出力します。

tokenizer.encode_plus("私は元気です。")

{'input_ids': [2, 1325, 9, 12453, 2992, 8, 3], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1]}

token type id は、特殊な fine-tuning を行わない限り、全部 0 として考えて問題ありません。
attention_mask は、トークンの実質的な長さを定めるときに重要です。言語モデルを使うときには、最大トークン数をある程度決めうちするのが普通なのですが、その最大トークン数全部が入力文で占められているわけではないので、使っていない部分に attention を貼らないようにするという意味です。

tokenizer.encode_plus("私は元気です。", padding='max_length', max_length=16)

{'input_ids': [2, 1325, 9, 12453, 2992, 8, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]}

上の例で、実質的なトークンがある部分のみ attention_mask の値が 1 になっているのが分かるでしょう。

また、encode_plus に return_tensors='pt' オプションをつけると、そのまま言語モデルに入力できる形式になるので便利です。

from transformers import AutoModelForMaskedLM, AutoTokenizer
model = AutoModelForMaskedLM.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")
tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")

inputs = tokenizer.encode_plus("私はとても[MASK]です。", return_tensors='pt')
outputs = model(**inputs)
tokenizer.convert_ids_to_tokens(outputs.logits[0][1:-1].argmax(axis=-1))

['私', 'は', 'とても', '幸せ', 'です', '。']

batch_encode_plus

複数の文を一度に処理したいときに便利です。

from transformers import AutoModelForMaskedLM, AutoTokenizer
model = AutoModelForMaskedLM.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")
tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")

sentences = [
  "私はとても[MASK]です。", 
  "昨日とても[MASK]なことがあった。", 
  "それはとても[MASK]な食べ物だね！"
]

inputs = tokenizer.batch_encode_plus(sentences, padding='max_length', max_length=16, return_tensors='pt')
outputs = model(**inputs)

for i in range(3):
    print(tokenizer.convert_ids_to_tokens(outputs.logits[i][1:-1].argmax(axis=-1)))

['私', 'は', 'とても', '幸せ', 'です', '。', '。', '」', '」', '」', 'そして', 'そう', 'そう', '。']
['昨', '##日', 'とても', '大変', 'な', 'こと', 'が', 'あっ', 'た', '。', '。', '残念', '残念', '」']
['それ', 'は', 'とても', '美味', 'な', '食べ物', 'だ', 'ね', '!', '。', '!', '大好き', '大好き', '!']