# %% import itertools from tqdm import tqdm import numpy as np import torch from transformers import BertJapaneseTokenizer, BertForTokenClassification import pytorch_lightning as pl # from torch.utils.data import DataLoader # import from_XML_to_json as XtC # import random # import json # import unicodedata # import pandas as pd # %% # 8-16 # PyTorch Lightningのモデル class BertForTokenClassification_pl(pl.LightningModule): def __init__(self, model_name, num_labels, lr): super().__init__() self.save_hyperparameters() self.bert_tc = BertForTokenClassification.from_pretrained( model_name, num_labels=num_labels ) 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.hparams.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