はじめに

　こんにちは、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など）と同じように、エンティティの特徴量を入力する必要があることがわかります。

LukeForEntitySpanClassificationの入力(Transformersより)

　これらの特徴量は、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は、推論時に再び集約する必要があります。詳細は実装コードのこちらの部分を参照ください。