はじめに
こんにちは、MNTSQ(モンテスキュー)のアルゴリズムエンジニアの清水です。本記事では事前学習済み言語モデルの一つであるLUKEを用いた固有表現抽出の実装方法について紹介します。
LUKEとは
LUKEは、LUKE: Deep Contextualized Entity Representations with Entity-aware Self-attentionにおいて提案された言語モデルです。LUKEは、単語とエンティティの文脈付きベクトルを出力する知識拡張型(knowledge-enhanced)の言語モデルであることが大きな特徴です。単語とエンティティの双方を独立したトークンとして扱い、(他のBERT系モデルと同様に)マスクされた単語を予測する訓練を行うと同時に、マスクされたエンティティを予測する訓練を行うことで、エンティティを考慮した文脈表現を獲得できます。また、固有表現抽出などのエンティティの知識を使うことが重要なタスクにおいて、当時のSoTAを達成しています。
詳しくは論文著者のスライドで分かりやすく解説されているので、ぜひご覧ください。
TransformersのLukeForEntitySpanClassificationによる固有表現抽出
LUKEはTransformersから簡単に利用することができます。LUKEではエンティティ表現を使って各種NLPタスクを解くことが可能であり、TransformersではLukeForEntitySpanClassification
を用いることで固有表現抽出を行うことができます。ただし、少し独特な実装が必要であり、公式のドキュメントにも多くの記載がないため、本記事で実装方法を紹介します。また、GitHubに実装コードを公開しているので、そちらも合わせてご参照ください。
モデルへ入力
(詳しくは論文著者のスライドをご覧いただければと思いますが)LUKEの事前学習タスクは以下の通り、テキストとは別に入力されるエンティティを復元するタスクになっています。 (本記事では図中の”Words”を「単語」、”Entities”を「エンティティ」と呼称します。)
TransformersのドキュメントのLUKEのページを参照すると、通常のBERTに入力する各単語の特徴量(input_ids
、attention_mask
など)と同じように、エンティティの特徴量を入力する必要があることがわかります。
これらの特徴量は、entity_spans
という特徴量(後述)をLukeTokenizer
(もしくはその多言語版であるMLukeTokenizer
)に入力することによって自動で作成されるため、実装上は特に意識する必要がありません。ただentity_position_ids
の計算方法を把握しておくと、モデル理解の助けになるでしょう。entity_position_ids
はエンティティがどの単語に対応するかを示すことで、テキスト中のエンティティの位置を表現します。エンティティが複数の単語に対応する場合、下記スライドのようにポジションエンベディングからの出力を平均し、モデル本体に入力します。
入力の作成方法
以下のようにTransformersのLukeTokenizer
(もしくはMLukeTokenizer
)にentity_spans
を与えることで、前述したモデルに必要な入力を自動で作成することができます。引数のtask
には "entity_classification" 、 "entity_pair_classification" 、 "entity_span_classification" のいずれかを指定する必要があり、本記事では固有表現抽出を行うことができる"entity_span_classification"を指定します。
tokenizer = MLukeTokenizer.from_pretrained( "studio-ousia/luke-japanese-base-lite", task="entity_spans_classification" ) text: str = "MNTSQ株式会社は全ての合意をフェアにします。" # entity_spansの具体的な作成方法については後述 entity_spans: list[tuple[int, int]] = [(0, 5), (0, 9), (0, 10), (0, 12)...] encoding = tokenizer(text, entity_spans=entity_spans, return_tensors="pt")
このentity_spans
が、通常の固有表現抽出から連想される形式と異なるため注意が必要です。テキストを任意の方法(MeCab・Sudachiなどによる分かち書き)で分割し、連続するトークンの組み合わせをentity_spans
として入力します。より正確には、テキスト中のエンティティの始点位置・終点位置を示す二つの整数のタプルのリストを渡します。入力イメージとしては以下の通りです。
text = "MNTSQ株式会社は全ての合意をフェアにします。" # 任意の方法で分割し、連続するトークンの組み合わせを`entity_spans`として与える。 # MNTSQ/株式会社/は/全て/の/合意/を/フェア/に/し/ます/。(Sudachiによる分かち書き) entity_spans = [ (0, 5), # "MNTSQ", (0, 9), # "MNTSQ株式会社", (0, 10), # "MNTSQ株式会社は", (0, 12), # "MNTSQ株式会社は全て", (5, 9), # "株式会社", (5, 10), # "株式会社は", # 省略 (20, 21), # "し", (20, 23), # "します", (20, 24), # "します。", (21, 23), # "ます", (21, 24), # "ます。", (23, 24), # "。" ]
このentity_spans
を使用して、entity_position_ids
などのモデルに必要な特徴量がLukeTokenizer
によって作成されます。また、正解ラベルはこのentity_spans
内のスパン一つ一つに対して付与されます。上記の例の場合は上から2つ目の"MNTSQ株式会社"のスパンに対してのみ"ORG"などのラベルが付与されることが期待されます*1。また、(LukeTokenizer
によって作成される)entity_ids
には一律して"[MASK]"という特殊トークンのIDが割り振られます。*2*3。
その他実装上の注意
LUKEと他のBERT系のモデル間で最も異なる点はentity_spans
を作成してトークナイザに入力する必要がある点です。具体的な実装は実装コードをご覧いただければと思いますが、本記事でも実装上注意する必要がある点をご紹介しておきます。
一つのentity_spansにおける最大長
全ての連続するトークンの組み合わせをentity_spans
として作成すると、スパンの数が膨大になりますし、また明らかにエンティティではないスパン(例えばテキストの最初から最後までのスパン)も作成されることになります。そのため、スパンの最大長を決めて、一つ一つのスパンがその最大長を超えないようにentity_spans
を作成する必要があります。また、この最大長はトークナイザによって分割される単語(上記論文著者のスライドの18ページの”Words”に当たる)の数を基準に設定する必要があります。つまり、一つのエンティティに定めた個数以上の単語が含まれないように制限します。また、その最大長をLukeTokenizer
にmax_mention_length
という引数として与える必要があります。詳細は実装コードのこちらの部分を参照ください。
各データに対するentity_spans内のスパンの個数
また、メモリーエラーを防ぐために、一つのデータ(一つのテキスト)に対するentity_spans
内のスパンの個数も制限する必要があります。溢れてしまったスパンは、テキストを複製し別のデータとしてモデルに入力します。この一つのデータに対するスパンの最大個数はLukeTokenizer
にmax_entity_length
という引数として与える必要があります。詳細は実装コードのこちらの部分を参照ください。
また、このように別データとして入力された同一テキストに対するentity_spans
は、推論時に再び集約する必要があります。詳細は実装コードのこちらの部分を参照ください。
推論時の処理
推論時にも独特の処理が必要です。具体的には、予測確率の高い順に、テキストの各文字に対して予測結果を反映していく処理が必要になります。詳細は実装コードのこちらの部分 を参照ください。
精度評価
ストックマーク株式会社が公開している日本語の固有表現抽出データセットで評価*4したところ、Accuracyが0.96、F1スコアが0.89となりました。詳しい結果はWandBのログをご覧ください。
終わりに
LUKEによる固有表現抽出は通常の言語モデルによる固有表現抽出とかなり性質が異なります。私は通常の固有表現抽出の先入観に引っ張られ、当初は誤った実装をしてしまいました。LUKEにご興味があり、実際に使ってみようと考えている方の参考になれば幸いです。
MNTSQ, Ltd.では一緒に働く仲間を募集しています
MNTSQでは自然言語処理やLLMにご興味のあるエンジニアを募集しています。是非、下記リンクからカジュアル面談をお申し込みください!
*1:一般的な固有表現抽出の"B-ORG"、"I-ORG"のようなBIO形式でないことに注意してください。一方で、推論時にはseqevalを使ってスコアを計算するためにBIOタグに変換しています。
*2:事前訓練時に使った"[MASK]"トークンを入力することで、入力テキスト中からエンティティに関する情報を集約した表現が得られます。詳しくはこちら。
*3:entity_idsは全て"[MASK]"トークンであるため、エンティティのエンベディング層が不要です。そのためエンティティのエンべディング層を持たないliteモデルを使用します。
*4:全データのうち20%をテストデータとして使用しました。