こんにちは。TSUNADE事業部研究チームリサーチャーの坂田です。

本記事では、Hugging Face 社が開発しているTransformersのPipelinesという仕組みを使って日本語の固有表現抽出を行う方法をご紹介します。

Transformersとは？

TransformersはHuggingFace社が公開しているPython用ライブラリで、BERTを始めとするTransformer系の言語モデルを使用する際のデファクトスタンダードになっています。また、最近では音声系や画像系のモデルも実装されています。ここ1年ほどで文書分類や固有表現抽出などの下流タスクを簡単に試せるようなインターフェースの整備も進んでいます。

日本語学習済み言語モデル

Transformersでは、学習済みの言語モデルを使用出来るようになっています。英語のモデルが多いですが、最近は様々な言語のモデルが追加されてきています。日本語では、東北大学が公開したBERTの学習済みモデルがよく使われています。こちらは、日本語のWikipediaを使ってBERTの学習が行われています。
また、最近ではリクルート社のMegagon Labsから日本語T5モデルが、rinna社からGPT-2とRoBERTaが公開されています。弊社でも、日本語話し言葉BERTを公開しています。

Pipelines

Transformersには、下流タスクを対応するためにPipelinesという仕組みが用意されています。下流タスク名を指定するだけで簡単にタスクの実行が可能です。とりあえず試してみるというときに便利な設計になっています。

例: 固有表現抽出

from transformers import pipeline 

ner_pipeline = pipeline("ner")
print(ner_pipeline("Tokyo is the capital city of Japan"))

出力結果

[{'entity': 'I-LOC',
  'score': 0.99933624,
  'index': 1,
  'word': 'Tokyo',
  'start': 0,
  'end': 5},
 {'entity': 'I-LOC',
  'score': 0.99883103,
  'index': 7,
  'word': 'Japan',
  'start': 29,
  'end': 34}]

ただし、この機能を使用するためには学習済み言語モデルが各タスクで学習済みである必要があります。東北大BERTは固有表現抽出タスクでの学習がされていませんので、別途固有表現抽出タスクでの学習を行う必要があります。

Trainer

Transformersには、言語モデルの学習、下流タスクの学習共に使えるTrainerという仕組みがあります。以前は推論に特化していたものがTrainerの導入によって学習のコードも書きやすくなりました。

固有表現抽出とは？

固有表現抽出とは、文書中から、あらかじめ定められたカテゴリに当てはまる固有表現を抽出するタスクです。 1番有名なデータセットとしては、 CoNLL 2003があります。CoNLL 2003では、文中の固有表現に対して、 persons(人物)、organizations(組織)、 locations(場所)、 miscellaneous(その他)のラベルを付けます。固有表現抽出研究では、よく評価データセットとして使われています。英語とドイツ語が用意されています。

ラベル付けの方法はいくつか存在しますが、代表的なもののひとつにBIOタグがあります。固有表現の始端トークンにB-(固有表現の種類)、それ以降の固有表現内トークンにI-(固有表現の種類)、固有表現以外のトークンにはOのラベルが付与されます。下記の例ではPERは人物を、ORGは組織を表します。

Taro	Yamada	works	as	an	engineer	for	Retrieva	.
B-PER	I-PER	O	O	O	O	O	B-ORG	O

実際に日本語NERTを試してみる

以下で日本語データを使った固有表現抽出を行う方法をご紹介します。

必要な各種依存ライブラリのインストール

Pythonが使用可能な環境ならば、基本的には、依存ライブラリを含めて、pipコマンドを使用することで簡単にインストールが可能です。ただし、Transformersの使用には、Deep Learning用ライブラリのPytorchまたは、Tensorflowの2系のインストールが必要です。どちらを使うかで、使用出来るモデルや使用するAPIが多少異なります。本記事では、Pytorchを使用します。また、日本語では、英語などの単語ごとにスペースで区切られている言語と違い、単語ごとに区切る処理が必要になります。そのため、形態素解析器MeCabのラッパーfugashiとMeCab用辞書のipadicのインストールが必要です。

pip install torch transformers fugashi ipadic sklearn

使用するデータ

今回は、ストックマーク社から公開されている日本語固有表現抽出データセットを使用します。Wikipediaの文書に対して、「人名」、「地名」、「法人名」、「その他の組織」の4種類のラベルが付与されています。後述のPythonコードを実行するディレクトリで git clone してください。

git clone https://github.com/stockmarkteam/ner-wikipedia-dataset.git

日本語固有表現抽出データセットでのFine-tuning

データの前処理をTransformersの公式ドキュメントに準じた形で行います。公式ドキュメント該当ページ公式ドキュメントではWNUT 17という英語固有表現抽出のデータセットを使用していますが、今回は前述の日本語固有表現抽出データセットを使用します。

使用モジュールのimport

from collections import defaultdict
import json

from sklearn.model_selection import train_test_split
import torch
from transformers import BertConfig, BertJapaneseTokenizer
from transformers import BertForTokenClassification
from transformers import pipeline
from transformers import Trainer, TrainingArguments

定数定義

MAX_LENGTH = 256  # 最大文長 
BERT_MODEL = "cl-tohoku/bert-base-japanese"  # 使用する学習済みモデル
DATASET_PATH = "ner-wikipedia-dataset/ner.json"

データセットの読み込みと前処理

こちらで行うことは主に以下の5つです。

各文を単語IDに変換したencoded_text_listを作成
データセットのアノテーションを元に各単語にBIOラベルを付与
各文のラベル列を、ラベルID列に変換したencoded_labels_listを作成
(この際、特殊トークンなどの無視されるべきトークンのラベルIDは -100にする。)
実際に付与されたラベルのリスト unique_labelsを作成
ラベルとラベルIDを対応付けたlabel2id、id2labelの作成

with open(DATASET_PATH) as json_file:  # データセット読み込み
    ner_data_dict_list = json.load(json_file)
text_list = [ner_data_dict["text"] for ner_data_dict in ner_data_dict_list]
encoded_text_list = []  # エンコードされた文のリスト
split_text_list = []  # 単語ごとに分割された文のリスト

for text in text_list:
    encoded_text = tokenizer(text, max_length=MAX_LENGTH, pad_to_max_length=True)
    encoded_text_list.append(encoded_text)
    split_text = tokenizer.decode(encoded_text["input_ids"]).split()
    split_text_list.append(split_text)

entities_list = [ner_data_dict["entities"] for ner_data_dict in ner_data_dict_list]
labels_list = [["O"] * MAX_LENGTH for _ in range(len(text_list))]

unique_labels = set("O")  # 付与された固有表現抽出ラベルの集合
found_named_entity = False

for sample_idx, (encoded_text, split_text, entities) in enumerate(zip(encoded_text_list, split_text_list, entities_list)):
    if len(entities) == 0:
        continue
    target_entity = entities.pop(0)
    entity_name = target_entity["name"]
    entity_label = target_entity["type"]
    for word_idx, word in enumerate(split_text):
        if word == "[CLS]" or word == "[SEP]" or word == "[PAD]":  # 特殊トークンにIGNOREラベルの付与
            labels_list[sample_idx][word_idx] = "IGNORE"
            continue
        
        if len(entities) == 0:
            continue
            
        if entity_name.startswith(word):
            if entity_name.endswith(word):
                label = f"B-{entity_label}"
                labels_list[sample_idx][word_idx] = label
                unique_labels.add(label)
                
                if len(entities) >= 1:
                    target_entity = entities.pop(0)
                    entity_name = target_entity["name"]
                    entity_label = target_entity["type"]
                
            else:
                for word_idx_2 in range(word_idx + 1, len(split_text)):
                    if "".join(split_text[word_idx : word_idx_2 + 1]) not in entity_name.replace(" ", "").replace("　",""):
                        found_named_entity = False
                        break
                    if entity_name.endswith(split_text[word_idx_2]):
                        labels = [f"B-{entity_label}"] + [f"I-{entity_label}"] * (word_idx_2 - word_idx) 
                        unique_labels |= set(labels)
                        labels_list[sample_idx][word_idx : word_idx_2 + 1] = labels
                        found_named_entity = True
                        break
                if found_named_entity:
                    if len(entities) >= 1:
                        target_entity = entities.pop(0)
                        entity_name = target_entity["name"]
                        entity_label = target_entity["type"]
                    
                    
unique_labels = list(unique_labels)
label2id = {label : label_id for label_id, label in enumerate(unique_labels)}  # ラベルにIDを付与
label2id["IGNORE"] = -100  # IGNORE ラベル に ID -100 を付与
encoded_labels_list = [[label2id[label] for label in labels] for labels in labels_list]  # 各文に付与されたラベルをIDで置き換えたもののリストを作成
del label2id["IGNORE"]  # configに渡すときには消す必要あり
id2label= {id: label for label, id in label2id.items()}  # IDをkey、ラベル名をvalueとした逆の辞書を作成

PytorchのDatasetの作成
こちらのイニシャライザの中で、文の単語ID列の形式を、辞書のリストをリストの辞書に変換しています。

class NERDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = defaultdict(list)
        for encoding_dict in encodings:
            for key, value in encoding_dict.items():
                self.encodings[key].append(value)
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

訓練データ、開発データ、検証データそれぞれのDatasetを作成

train_encoded_text_list, test_encoded_text_list, train_encoded_labels_list, test_encoded_labels_list = \
    train_test_split(encoded_text_list, encoded_labels_list, test_size=0.1, random_state=0)
train_encoded_text_list, val_encoded_text_list, train_encoded_labels_list, val_encoded_labels_list = \
    train_test_split(encoded_text_list, encoded_labels_list, test_size=0.2, random_state=0)

train_dataset = NERDataset(train_encoded_text_list, train_encoded_labels_list)
val_dataset = NERDataset(val_encoded_text_list, val_encoded_labels_list)
test_dataset = NERDataset(test_encoded_text_list, test_encoded_labels_list)

東北大BERT の config を読み込み、固有表現ラベル名の部分を書き換え

config = BertConfig.from_pretrained(DATASET_PATH, id2label=id2label, label2id=label2id)

Trainerの設定

training_args = TrainingArguments(
    output_dir='./results',          
    num_train_epochs=3,              
    per_device_train_batch_size=16,  
    per_device_eval_batch_size=64,   
    warmup_steps=500, 
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    save_steps=50,
    do_eval=True,
    eval_steps=50
)

Trainerを使った学習の実行

trainer = Trainer(
    model=model,                         
    args=training_args, 
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)

trainer.train()

Pipelinesのインスタンス作成

ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, config=config)

Pipelinesを使った固有表現抽出実行例

以下のように、日本語固有表現抽出でPipelinesを試すことが出来るようになりました。

例1

ner_pipeline("2012年11月に韓国・社稷野球場で開催されたアジアシリーズ2012に、パース・ヒートの一員として出場した。")

実行結果1

[{'entity': 'B-施設名',
  'score': 0.954906,
  'index': 6,
  'word': '韓国',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.96713436,
  'index': 7,
  'word': '・',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.96969676,
  'index': 8,
  'word': '社',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.9882817,
  'index': 9,
  'word': '##稷',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.9852976,
  'index': 10,
  'word': '野球',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.72150517,
  'index': 11,
  'word': '場',
  'start': None,
  'end': None},
 {'entity': 'B-イベント名',
  'score': 0.4892079,
  'index': 16,
  'word': 'た',
  'start': None,
  'end': None},
 {'entity': 'B-イベント名',
  'score': 0.5815552,
  'index': 17,
  'word': 'アジア',
  'start': None,
  'end': None},
 {'entity': 'I-イベント名',
  'score': 0.9707927,
  'index': 18,
  'word': 'シリーズ',
  'start': None,
  'end': None}]

例2

ner_pipeline("大手町本部ビルは2011年より解体され、跡地には2015年に三井住友銀行本店東館が完成した。")

実行結果2

[{'entity': 'B-施設名',
  'score': 0.98983073,
  'index': 1,
  'word': '大手',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.98958486,
  'index': 2,
  'word': '##町',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.9917477,
  'index': 3,
  'word': '本部',
  'start': None,
  'end': None},
 {'entity': 'I-施設名',
  'score': 0.9465727,
  'index': 4,
  'word': 'ビル',
  'start': None,
  'end': None}]

おわりに

今回は、TransformersのPipelinesで日本語固有表現抽出を試す方法をご紹介しました。
機会があれば別タスクに関してもご紹介したいと思います。

Retrieva TECH BLOG

TransformersのPipelinesで日本語固有表現抽出

Transformersとは？

日本語学習済み言語モデル

Pipelines

Trainer

固有表現抽出とは？

実際に日本語NERTを試してみる

必要な各種依存ライブラリのインストール

使用するデータ

日本語固有表現抽出データセットでのFine-tuning

Pipelinesを使った固有表現抽出実行例

おわりに

参考

Transformersとは？

日本語学習済み言語モデル

Pipelines

Trainer

固有表現抽出とは？

実際に日本語NERTを試してみる

必要な各種依存ライブラリのインストール

使用するデータ

日本語固有表現抽出データセットでのFine-tuning

Pipelinesを使った固有表現抽出 実行例

おわりに

参考

Pipelinesを使った固有表現抽出実行例