日本語BERTに新しい単語（ユーザ辞書）を追加してみる

huggingface/transformersで扱える日本語BERT（cl-tohoku/bert-base-japanese-whole-word-masking）に新しい単語を追加する方法を調べていたので、ここにメモしておこうと思います。

同じことで悩んでいる人に届けば幸いです。

参考

といっても、以下の内容をまとめているだけです。

1. https://github.com/huggingface/transformers/issues/1413
2. https://stackoverflow.com/questions/64669365/huggingface-bert-tokenizer-add-new-token

実装

Google Colab上で検証しているので、必要なライブラリをインストールしておきます。

!apt install aptitude swig
!aptitude install mecab libmecab-dev mecab-ipadic-utf8 git make curl xz-utils file -y
!pip install mecab-python3
!pip install unidic-lite
!pip install transformers
!pip install fugashi
!pip install ipadic

まずは日本語BERTのtokenizerとmodelを用意

import numpy as np
import torch
from transformers import BertModel
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
model = BertModel.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

tokenizer側の設定

今回の検証では、デフォルトではサブワードに分割されてしまう、より細かい形態素に分割されてしまう、UNKトークンとみなされてしまう、ような単語を扱おうと思います。
以下のようにパーセプトロンはパー、##セプト、##ロンに分割され、人工知能は人工と知能に分割され、😁はUNKトークンとなっていることがわかります。

# 対象単語1: パーセプトロン（サブワードに分割されるパターン）
# 対象単語2: 人工知能(細かく形態素に分割されるパターン)
# 対象単語3: 😁(UNKトークン)
text = "パーセプトロンは人工知能の勉強に欠かせない概念ですね😁"

ids = tokenizer.encode(text, return_tensors='pt')
wakati = tokenizer.convert_ids_to_tokens(ids[0])
print(ids)
print(wakati)
#tensor([[    2,  1673,  5048,  1006,     9,  4969, 12588,     5,  8192,     7,
#         20059,   191,    80,  3214,  2992,  1852,     1,     3]])
#['[CLS]', 'パー', '##セプト', '##ロン', 'は', '人工', '知能', 'の', '勉強', 'に', '欠か', 'せ', 'ない', '概念', 'です', 'ね', '[UNK]', '[SEP]']

tokenizerに単語やトークンを追加したいときは、以下のようにtokenizer.add_tokensを使って追加したい単語たちを配列で渡します。
既に登録されている単語をtokenizer.add_tokensで追加しようとしても無視されます。

# 単語追加前のvocab size
print("vocab size", len(tokenizer))
#vocab size 32000

# 単語追加
# 現状special_tokens=Trueを設定しないとエラーコードが返ってくるようです。
tokenizer.add_tokens(["パーセプトロン", "人工知能", "😁"], special_tokens=True)

# vocab sizeが増えていることを確認
print("vocab size", len(tokenizer))
#vocab size 32003

改めて形態素を確認してみると、正しく1語として認識されていることがわかります。追加した順番に新しい単語IDが付与されていることもわかります。

ids = tokenizer.encode(text, return_tensors='pt')
wakati = tokenizer.convert_ids_to_tokens(ids[0])
print(ids)
print(wakati)
#tensor([[    2, 32000,     9, 32001,     5,  8192,     7, 20059,   191,    80,
#          3214,  2992,  1852, 32002,     3]])
#['[CLS]', 'パーセプトロン', 'は', '人工知能', 'の', '勉強', 'に', '欠か', 'せ', 'ない', '概念', 'です', 'ね', '😁', '[SEP]']

model側の設定

上記のようにtokenizer側の設定だけでは、tokenizerに辞書を追加しただけの状態なので、モデル側の単語のベクトルを扱うembeddingの層に、追加した単語のベクトルが存在せず、エラーがでます。

model(ids)

# 〜省略〜
#/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
#   1914         # remove once script supports set_grad_enabled
#   1915         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
#-> 1916     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
#   1917 
#   1918 
#IndexError: index out of range in self

modelのembeddingの層に追加した単語のベクトルを追加します。
model.resize_token_embeddingsで追加語の全単語数（今回のケースでは32003）を指定することでembeddingの行数を拡張することができます。
新しく追加した単語ベクトルはそれぞれ末尾にランダムで生成されるようです。

# 追加前のembeddingのサイズを確認する
print(model.embeddings.word_embeddings)
#Embedding(32000, 768, padding_idx=0)

# ベクトルを追加
model.resize_token_embeddings(len(tokenizer))

print(model.embeddings.word_embeddings)
#Embedding(32003, 768)

# 新しく追加されたベクトルはrandom vectorで生成される
print(model.embeddings.word_embeddings.weight[-1, :]) # 😁
print(model.embeddings.word_embeddings.weight[-2, :]) # 人工知能
print(model.embeddings.word_embeddings.weight[-3, :]) # パーセプトロン
#tensor([-1.1561e-02, -4.2187e-02, -1.6742e-02, -2.4143e-03,  8.0074e-03,
#        -4.1712e-02,  1.5328e-02,  3.6373e-03, -4.9351e-02, -2.5222e-03,
#        -2.5747e-02,  2.7715e-03,  2.3406e-02,  1.9393e-03, -1.2993e-02,
#        -7.8929e-03, -3.5028e-02, -1.1273e-02, -3.1112e-02,  3.3033e-04,
#        -1.2306e-02, -9.9428e-03,  1.4857e-02,  4.4864e-03,  1.5033e-03,
#         6.0422e-02, -1.0403e-02,  2.6659e-02, -4.2973e-02, -6.7658e-03,
#        -1.5566e-02, -1.9462e-02,  7.0844e-03, -1.5666e-02,  4.3054e-02,
#        -8.6333e-03,  1.8615e-02,  5.3128e-03,  2.8191e-03, -2.3805e-02,
# 〜省略〜

おわりに

こうやって追加した新しい単語は、このままだとランダムなベクトルなので、適切な意味のあるベクトルに更新したいところですが、良い方法がまだ良くわかっておりません。。。（←ここが重要でしょ、って感じですがまだ勉強中で申し訳です💦）

上記のように単語をランダムなベクトルとして追加したあとは、

1. タスク個別に（最終層などを）Fine tuningする（i.e. 追加した単語ベクトルはランダムなまま）
2. Fine tuningするときにembeddingも学習対象とする（既存の単語空間がおかしくなったりしないかなー？）
3. はたまた言語モデルとして追加の事前学習をするのか（やり方がよくわかっていない）

などの方針が考えられるのかなーと思いますが、3.はデータ集めるの大変だし現実的じゃないのかなーとか思ったりしてます。無難に1.ないしは2.を試すのが良いんですかね？

参考として挙げたhuggingfaceのGithubのissueでhuggingface内部の人が、以下のように述べています。

Adding tokens adds tokens at the end of the tokenizer's vocabulary, essentially extending the vocabulary. The model's embedding matrix would need to be resized as well to take into account the new tokens, but all the other tokens would keep their representation as-is. Seeing as the new rows in the embedding matrix are randomly initialized, you would still need to fine-tune the model to a dataset containing such tokens.
(DeepLによる翻訳)
トークンを追加すると、トークナイザーの語彙の最後にトークンが追加され、基本的に語彙が拡張されます。モデルの埋め込み行列も、新しいトークンを考慮してサイズを変更する必要がありますが、その他のトークンはそのままの表現を維持します。埋め込み行列の新しい行はランダムに初期化されるため、そのようなトークンを含むデータセットに対してモデルを微調整する必要があります。

モデルを微調整ってのが1.や2.の方法なんだろうと解釈しましたが、はたして。。。
この辺私なりにまとまったら、ここに加筆しようと思います。

おわり