Sansan Builders Blog

Sansanのものづくりを支えるメンバーの技術やデザイン、プロダクトマネジメントの情報を発信

【Techの道も一歩から】第27回「BERTで作ってみた日本語固有表現抽出器の推論部分を書く」

f:id:s_yuka:20200424095303j:plain こんにちは。DSOC 研究開発部の高橋寛治です。

今回は、前回の 「BERTで日本語固有表現抽出器を作ってみた」 に続き、作った固有表現抽出器をWebAPI化します。

モデルを把握する

transformers ライブラリの 固有表現抽出のサンプル を流用してモデルを作成しました。

こちらのコードをもとに学習を実行すると、コマンドライン引数で指定したディクレトリにモデルファイルが出力されます。

model_dir
├── config.json
├── eval_results.txt
├── pytorch_model.bin
├── special_tokens_map.json
├── test_predictions.txt
├── test_results.txt
├── test_test_predictions.tsv
├── tokenizer_config.json
├── training_args.bin
└── vocab.txt

transformers ライブラリでは、固有表現抽出器で利用する BertTokenizerBertForTokenClassification クラスがファイルから読み込むメソッド(from_pretrained)を提供しています。 このメソッドに、上記のモデルファイルのディレクトリのパスを渡すことで、モデルを読み込みます。

ここから紹介するコードではやや冗長な箇所もありますが、サンプルで提供されている utils_ner.py を最大限に利用しています。

推論部を記述する

サンプルコードでは推論をし評価を行う evaluate メソッドが提供されています。 このメソッドに次の引数を渡しています。

  • モデルインスタンス
  • トークナイザインスタンス
  • ラベルマップ
  • パディング用トークンのID
  • モード(dev, test)

この部分から抜粋して、次のように利用可能なクラスとして実装します。

bert_ner = BertNER(
    model_path="janer-model",
    label_path="./data/preprocessed/labels.txt"
)
result = bert_ner.predict("解析対象の文書")

evaluate メソッドに渡していた引数を保持

上記の evaluate メソッドが持っていた状態をうまくインスタンス変数で保有するようにします。

import unicodedata

from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from transformers import (
    BertConfig,
    BertForTokenClassification,
    BertTokenizer
)
from utils_ner import (
    InputExample,
    InputFeatures,
    convert_examples_to_features,
    get_labels
)


class BertNER:
    def __init__(self, model_path, label_path):
        # モデル
        self.model = BertForTokenClassification.from_pretrained(model_path)
        self.model.to("cpu")
        self.model.eval()

        # トークナイザ
        self.tokenizer = BertTokenizer.from_pretrained(
            model_path,
            do_lower_case=False
        )
        # 形態素解析用
        self.tagger = MeCab.Tagger("-Owakati")

        # ラベルマップ
        self.labels = get_labels(label_path)
        self.label_map = {i: label for i, label in enumerate(self.labels)}

        # パディング用トークン
        self._pad_token_label_id = CrossEntropyLoss().ignore_index

        # モード
        self._mode = "test"

テキストを読み込む

日本語処理を前提にして、句点「。」で文として分割します。 文は、utils_ner.py で提供される InputExample クラスを利用して保持します。

BertNER クラスに、次のメソッドを追加します。

def load_document(self, text):
    examples = []
    normalized_text = unicodedata.normalize("NFKC", text)
    sentences = [sentence.strip() + "。" for sentence in normalized_text.split("。") if sentence.strip()]
    for guid, sentence in enumerate(sentences, start=1):
        words = self.tagger.parse(sentence).strip().split()
        dummy_labels = ['O' for _ in words]
        examples.append(InputExample(guid="{}-{}".format(self._mode, guid), words=words, labels=dummy_labels))

    return examples

推論の流れを記述する

文書を受け取り、トークナイズされた文集合へと分割します。 この文集合をPyTorchで取り扱える形式へと変換したのちに、推論し結果を返します。

BertNER クラスに、次のメソッドを追加します。

def predict(self, text):
    examples = self.load_document(text)
    dataloader = self.convert_examples_to_dataloader(examples)
    preds_list = self._predict(dataloader)
    return self._to_dict(examples, preds_list)

InputExample クラスから DataLoader

PyTorchで扱える形式のデータへと変換します。 utils_ner.py で提供される convert_examples_to_features メソッドを利用します。 各種パラメータは学習時と同様にします。

BertNER クラスに、次のメソッドを追加します。

def convert_examples_to_dataloader(self, examples):
    features = convert_examples_to_features(
        examples,
        self.labels,
        128,
        self.tokenizer,
        cls_token_at_end=False,
        cls_token=self.tokenizer.cls_token,
        cls_token_segment_id=0,
        sep_token=self.tokenizer.sep_token,
        sep_token_extra=False,
        pad_on_left=False,
        pad_token=self.tokenizer.convert_tokens_to_ids([self.tokenizer.pad_token])[0],
        pad_token_segment_id=0,
        pad_token_label_id=self._pad_token_label_id,
    )

    dataset = TensorDataset(
        torch.tensor([f.input_ids for f in features], dtype=torch.long),
        torch.tensor([f.input_mask for f in features], dtype=torch.long),
        torch.tensor([f.segment_ids for f in features], dtype=torch.long),
        torch.tensor([f.label_ids for f in features], dtype=torch.long)
    )

    return DataLoader(dataset, batch_size=32)

読み込んだデータに対しての推論

DataLoader からバッチを取得し、モデルに入力します。 CPUでの推論を行うため、それに必要な設定や、勾配を計算しない設定を行います。

コード後半部分は、モデルの出力列を数値列からラベル列に書き換えます。

BertNER クラスに、次のメソッドを追加します。

def _predict(self, dataloader):
    preds = None
    out_label_ids = None

    for batch in dataloader:
        batch = tuple(t.to('cpu') for t in batch)

        with torch.no_grad():
            inputs = {
                "input_ids": batch[0],
                "attention_mask": batch[1],
                "token_type_ids": batch[2],
                "labels": batch[3]
            }
            outputs = self.model(**inputs)
            _, logits = outputs[:2]

            if preds is None:
                preds = logits.detach().cpu().numpy()
                out_label_ids = inputs["labels"].detach().cpu().numpy()
            else:
                preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
                out_label_ids = np.append(out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)

    preds = np.argmax(preds, axis=2)

    out_label_list = [[] for _ in range(out_label_ids.shape[0])]
    preds_list = [[] for _ in range(out_label_ids.shape[0])]

    for i in range(out_label_ids.shape[0]):
        for j in range(out_label_ids.shape[1]):
            if out_label_ids[i, j] != self._pad_token_label_id:
                out_label_list[i].append(self.label_map[out_label_ids[i][j]])
                preds_list[i].append(self.label_map[preds[i][j]])
    return preds_list

推論されたタグと単語を紐付ける

推論結果と単語を紐付けて、トークンを辞書で表現します。

BertNER クラスに、次のメソッドを追加します。

def _to_dict(self, examples, preds_list):
    sentences = []
    for example, pred_list in zip(examples, preds_list):
        sentence = [
            {"Token": word, "Type": tag}
            for word, tag in zip(example.words, pred_list)
        ]
        sentences.append(sentence)
    return sentences

推論を試す

完成した BertNER クラスを ipython 上で試します。

$ ipython
In [1]: from ner import BertNER

In [2]: bert_ner = BertNER()

In [3]: bert_ner = BertNER(
   ...:     model_path="model_path",
   ...:     label_path="./path/to/labels.txt"
   ...: )

In [4]: result = bert_ner.predict("Sansan株式会社はブログ「Sansan Builders Box」を運営しています。")

In [5]: result
Out[5]:
[[{'Token': 'Sansan', 'Type': 'B-ORGANIZATION'},
  {'Token': '株式会社', 'Type': 'I-ORGANIZATION'},
  {'Token': 'は', 'Type': 'O'},
  {'Token': 'ブログ', 'Type': 'O'},
  {'Token': '「', 'Type': 'O'},
  {'Token': 'Sansan', 'Type': 'O'},
  {'Token': 'Builders', 'Type': 'O'},
  {'Token': 'Box', 'Type': 'O'},
  {'Token': '」', 'Type': 'O'},
  {'Token': 'を', 'Type': 'O'},
  {'Token': '運営', 'Type': 'O'},
  {'Token': 'し', 'Type': 'O'},
  {'Token': 'て', 'Type': 'O'},
  {'Token': 'い', 'Type': 'O'},
  {'Token': 'ます', 'Type': 'O'},
  {'Token': '。', 'Type': 'O'}]]

このようにして推論することができました。 後段の処理で、固有表現のトークンを結合するなど処理するといいでしょう。

まとめ

transformers が公開しているサンプルを元に、日本語の固有表現抽出器の推論部分を記述しました。 サンプルコードでは、推論と評価部分がセットになっているため、実際に使う際には、切り分けが必要です。 コードの役割や実装意図を紹介しました。

BERTを利用したモデルには可能性を感じるため、いろいろなところで気軽に適用していきたいです。

▼本連載のほかの記事はこちら

buildersbox.corp-sansan.com

執筆者プロフィール

高橋寛治 Sansan株式会社 DSOC (Data Strategy & Operation Center) R&Dグループ研究員

阿南工業高等専門学校卒業後に、長岡技術科学大学に編入学。同大学大学院電気電子情報工学専攻修了。在学中は、自然言語処理の研究に取り組み、解析ツールの開発や機械翻訳に関連する研究を行う。大学院を卒業後、2017年にSansan株式会社に入社。キーワード抽出など自然言語処理を生かした研究開発に取り組む。

© Sansan, Inc.