NER_medNLP.py

# %%

import itertools

import numpy as np
import pytorch_lightning as pl
import torch
from transformers import BertForTokenClassification, BertJapaneseTokenizer


# %%
# 8-16
# PyTorch Lightningのモデル
class BertForTokenClassification_pl(pl.LightningModule):

    def __init__(self, num_labels, model='sociocom/MedNERN-CR-JA', lr=1e-4):
        super().__init__()
        self.lr = lr
        self.save_hyperparameters()
        self.bert_tc = BertForTokenClassification.from_pretrained(
            model,
            num_labels=num_labels,
            ignore_mismatched_sizes=True
        )

    @classmethod
    def from_pretrained_bin(cls, model_path, num_labels):
        model = cls(num_labels)
        model.load_state_dict(torch.load(model_path))
        return model

    def training_step(self, batch, batch_idx):
        output = self.bert_tc(**batch)
        loss = output.loss
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        output = self.bert_tc(**batch)
        val_loss = output.loss
        self.log('val_loss', val_loss)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)


# %%
class NER_tokenizer_BIO(BertJapaneseTokenizer):

    # 初期化時に固有表現のカテゴリーの数`num_entity_type`を
    # 受け入れるようにする。
    def __init__(self, *args, **kwargs):
        self.num_entity_type = kwargs.pop('num_entity_type')
        super().__init__(*args, **kwargs)

    def encode_plus_tagged(self, text, entities, max_length):
        """
        文章とそれに含まれる固有表現が与えられた時に、
        符号化とラベル列の作成を行う。
        """
        # 固有表現の前後でtextを分割し、それぞれのラベルをつけておく。
        splitted = []  # 分割後の文字列を追加していく
        position = 0

        for entity in entities:
            start = entity['span'][0]
            end = entity['span'][1]
            label = entity['type_id']
            splitted.append({'text': text[position:start], 'label': 0})
            splitted.append({'text': text[start:end], 'label': label})
            position = end
        splitted.append({'text': text[position:], 'label': 0})
        splitted = [s for s in splitted if s['text']]

        # 分割されたそれぞれの文章をトークン化し、ラベルをつける。
        tokens = []  # トークンを追加していく
        labels = []  # ラベルを追加していく
        for s in splitted:
            tokens_splitted = self.tokenize(s['text'])
            label = s['label']
            if label > 0:  # 固有表現
                # まずトークン全てにI-タグを付与
                # 番号順O-tag:0, B-tag:1~タグの数，I-tag:タグの数〜
                labels_splitted = \
                    [label + self.num_entity_type] * len(tokens_splitted)
                # 先頭のトークンをB-タグにする
                labels_splitted[0] = label
            else:  # それ以外
                labels_splitted = [0] * len(tokens_splitted)

            tokens.extend(tokens_splitted)
            labels.extend(labels_splitted)

        # 符号化を行いBERTに入力できる形式にする。
        input_ids = self.convert_tokens_to_ids(tokens)
        encoding = self.prepare_for_model(
            input_ids,
            max_length=max_length,
            padding='max_length',
            truncation=True
        )

        # ラベルに特殊トークンを追加
        # max_lengthで切り取って，その前後に[CLS]と[SEP]を追加するためのラベルを入れる
        labels = [0] + labels[:max_length - 2] + [0]
        # max_lengthに満たない場合は，満たない分を後ろ側に追加する
        labels = labels + [0] * (max_length - len(labels))
        encoding['labels'] = labels

        return encoding

    def encode_plus_untagged(
            self, text, max_length=None, return_tensors=None
    ):
        """
        文章をトークン化し、それぞれのトークンの文章中の位置も特定しておく。
        IO法のトークナイザのencode_plus_untaggedと同じ
        """
        # 文章のトークン化を行い、
        # それぞれのトークンと文章中の文字列を対応づける。
        tokens = []  # トークンを追加していく。
        tokens_original = []  # トークンに対応する文章中の文字列を追加していく。
        words = self.word_tokenizer.tokenize(text)  # MeCabで単語に分割
        for word in words:
            # 単語をサブワードに分割
            tokens_word = self.subword_tokenizer.tokenize(word)
            tokens.extend(tokens_word)
            if tokens_word[0] == '[UNK]':  # 未知語への対応
                tokens_original.append(word)
            else:
                tokens_original.extend([
                    token.replace('##', '') for token in tokens_word
                ])

        # 各トークンの文章中での位置を調べる。（空白の位置を考慮する）
        position = 0
        spans = []  # トークンの位置を追加していく。
        for token in tokens_original:
            l = len(token)
            while 1:
                if token != text[position:position + l]:
                    position += 1
                else:
                    spans.append([position, position + l])
                    position += l
                    break

        # 符号化を行いBERTに入力できる形式にする。
        input_ids = self.convert_tokens_to_ids(tokens)
        encoding = self.prepare_for_model(
            input_ids,
            max_length=max_length,
            padding='max_length' if max_length else False,
            truncation=True if max_length else False
        )
        sequence_length = len(encoding['input_ids'])
        # 特殊トークン[CLS]に対するダミーのspanを追加。
        spans = [[-1, -1]] + spans[:sequence_length - 2]
        # 特殊トークン[SEP]、[PAD]に対するダミーのspanを追加。
        spans = spans + [[-1, -1]] * (sequence_length - len(spans))

        # 必要に応じてtorch.Tensorにする。
        if return_tensors == 'pt':
            encoding = {k: torch.tensor([v]) for k, v in encoding.items()}

        return encoding, spans

    @staticmethod
    def Viterbi(scores_bert, num_entity_type, penalty=10000):
        """
        Viterbiアルゴリズムで最適解を求める。
        """
        m = 2 * num_entity_type + 1
        penalty_matrix = np.zeros([m, m])
        for i in range(m):
            for j in range(1 + num_entity_type, m):
                if not ((i == j) or (i + num_entity_type == j)):
                    penalty_matrix[i, j] = penalty
        path = [[i] for i in range(m)]
        scores_path = scores_bert[0] - penalty_matrix[0, :]
        scores_bert = scores_bert[1:]

        for scores in scores_bert:
            assert len(scores) == 2 * num_entity_type + 1
            score_matrix = np.array(scores_path).reshape(-1, 1) \
                           + np.array(scores).reshape(1, -1) \
                           - penalty_matrix
            scores_path = score_matrix.max(axis=0)
            argmax = score_matrix.argmax(axis=0)
            path_new = []
            for i, idx in enumerate(argmax):
                path_new.append(path[idx] + [i])
            path = path_new

        labels_optimal = path[np.argmax(scores_path)]
        return labels_optimal

    def convert_bert_output_to_entities(self, text, scores, spans):
        """
        文章、分類スコア、各トークンの位置から固有表現を得る。
        分類スコアはサイズが（系列長、ラベル数）の2次元配列
        """
        assert len(spans) == len(scores)
        num_entity_type = self.num_entity_type

        # 特殊トークンに対応する部分を取り除く
        scores = [score for score, span in zip(scores, spans) if span[0] != -1]
        spans = [span for span in spans if span[0] != -1]

        # Viterbiアルゴリズムでラベルの予測値を決める。
        labels = self.Viterbi(scores, num_entity_type)

        # 同じラベルが連続するトークンをまとめて、固有表現を抽出する。
        entities = []
        for label, group \
                in itertools.groupby(enumerate(labels), key=lambda x: x[1]):

            group = list(group)
            start = spans[group[0][0]][0]
            end = spans[group[-1][0]][1]

            if label != 0:  # 固有表現であれば
                if 1 <= label <= num_entity_type:
                    # ラベルが`B-`ならば、新しいentityを追加
                    entity = {
                        "name": text[start:end],
                        "span": [start, end],
                        "type_id": label
                    }
                    entities.append(entity)
                else:
                    # ラベルが`I-`ならば、直近のentityを更新
                    entity['span'][1] = end
                    entity['name'] = text[entity['span'][0]:entity['span'][1]]

        return entities