BERTによる日本語固有表現抽出

{
    "curid": "473536",
    "text": "イギリスはリーマンショック直後の2008年10月にイングランド銀行のバランスシートを一気に3倍近く増やした後、2008年11月から2009年3月にかけて段階的に縮小させていった。",
    "entities": [
        {
            "name": "イギリス",
            "span": [0,4],
            "type": "地名"
        },
        {
            "name": "リーマンショック",
            "span": [5,13],
            "type": "イベント名"
        },
        {
            "name": "イングランド銀行",
            "span": [25,33],
            "type": "政治的組織名"
        }
    ]
}

import json
from transformers import BertJapaneseTokenizer
from label import label2id

MAX_LENGTH = 128  # 一文あたりの最大トークン数
BERT_MODEL = "cl-tohoku/bert-base-japanese-v2"  # 使用する学習済みモデル
DATASET_PATH = "./dataset/ner.json"
TAGGED_DATASET_PATH = "./dataset/ner_tagged.json"

# 1. データ読み込み

with open(DATASET_PATH) as f:
  ner_data_list = json.load(f)

# 2. 固有表現タグづけ

# 半角スペースによってエンティティの開始位置がずれるので、エンティティの開始・終了位置を調整する（トークナイズで半角スペースが削除されるため）
def adjust_entity_span(text, entities):
  white_spece_posisions = [i for i, c in enumerate(text) if c == " "]
  for entity in entities:
    start_pos = entity["span"][0]
    end_pos = entity["span"][1]
    start_diff = sum(white_spece_pos < start_pos for white_spece_pos in white_spece_posisions)
    end_diff = sum(white_spece_pos < end_pos for white_spece_pos in white_spece_posisions)
    entity["span"] = [start_pos - start_diff, end_pos - end_diff]

for ner_data in ner_data_list:
  adjust_entity_span(ner_data["text"], ner_data["entities"])

sentence_list = [ner_data["text"] for ner_data in ner_data_list]

tokenizer = BertJapaneseTokenizer.from_pretrained(BERT_MODEL)

encoded_sentence_list = [tokenizer(sentence, max_length=MAX_LENGTH, padding="max_length", truncation=True) for sentence in sentence_list]

def calc_token_length(token):
  return len(token) -2 if token.startswith("##") else len(token)

def warn_start_pos(pos, token, entity, curid):
  print("[WARN] トークンの開始位置がエンティティの開始位置を超えました。エンティティの開始=<" + str(entity["span"][0]) + "> トークンの開始=<" + str(pos) + "> curid=<" + curid + "> token=<" + token + "> entity=<" + entity["name"] + ">")

def warn_end_pos(pos, token, entity, curid):
  token_length = calc_token_length(token)
  print("[WARN] トークンの終了位置がエンティティの終了位置を超えました。エンティティの終了=<" + str(entity["span"][1]) + "> トークンの終了=<" + str(pos + token_length) + "> curid=<" + curid + "> token=<" + token + "> entity=<" + entity["name"] + ">")

def search_tokens(tokens, entity, curid):
  ret = {}

  entity_type = entity["type"]
  entity_span = entity["span"]
  entity_start_pos = entity_span[0]
  entity_end_pos = entity_span[1]

  pos = 0
  is_inside_entity = False
  for i, token in enumerate(tokens):
    if token in ["[UNK]", "[SEP]", "[PAD]"]:
      break
    elif token == "[CLS]":
      continue

    token_length = calc_token_length(token)
    if not is_inside_entity: # まだエンティティの中に入っていない場合
      if pos == entity_start_pos: # トークンの開始がエンティティの開始に一致した場合
        ret[i] = "B-" + entity_type
        if pos + token_length == entity_end_pos: # トークンの終了がエンティティの終了に一致した場合
          break
        elif pos + token_length < entity_end_pos:
          is_inside_entity = True
        else: # [WARN]トークンの終了がエンティティの終了を超えた場合
          warn_end_pos(pos, token, entity, curid)
          print(tokens)
      elif pos > entity_start_pos: # [WARN]トークンの開始がエンティティの開始を超えた場合
        warn_start_pos(pos, token, entity, curid)
        print(tokens)
        break
    else: # エンティティの中に入っている場合
      if pos + token_length == entity_end_pos: # トークンの終わりがエンティティの終わりに一致した場合
        ret[i] = "I-" + entity_type
        is_inside_entity = False
        break
      elif pos + token_length < entity_end_pos: # トークンがまだエンティティの終わりに達していない場合
        ret[i] = "I-" + entity_type
      else: # [WARN]トークンがエンティティの終わりを超えた場合
        warn_end_pos(pos, token, entity, curid)
        print(tokens)
        ret.clear()
        is_inside_entity = False
        break
    pos += token_length

  return ret

# トークンにタグ付けをする
tags_list = []
for i, encoded_sentence in enumerate(encoded_sentence_list):
  tokens = tokenizer.convert_ids_to_tokens(encoded_sentence["input_ids"])

  tags = ["O"] * MAX_LENGTH

  ner_data = ner_data_list[i]
  curid = ner_data["curid"]

  entities = ner_data["entities"]

  for entity in entities:
    found_token_pos_tags = search_tokens(tokens, entity, curid)
    for pos, tag in found_token_pos_tags.items():
      tags[pos] = tag

  tags_list.append(tags)

  # 固有表現タグをIDに変換
  encoded_tags_list = [[label2id[tag] for tag in tags] for tags in tags_list] # 学習で利用

# タグづけしたデータの保存
tagged_sentence_list = []
for encoded_sentence, encoded_tags in zip(encoded_sentence_list, encoded_tags_list):
  tagged_sentence = {}
  tagged_sentence['input_ids'] = encoded_sentence['input_ids']
  tagged_sentence['token_type_ids'] = encoded_sentence['token_type_ids']
  tagged_sentence['attention_mask'] = encoded_sentence['attention_mask']
  tagged_sentence['labels'] = encoded_tags
  tagged_sentence_list.append(tagged_sentence)

with open(TAGGED_DATASET_PATH, 'w') as f:
  json.dump(tagged_sentence_list, f)
  

import json
import torch
from transformers import BertJapaneseTokenizer, BertForTokenClassification, BertConfig
from label import label2id, id2label

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

MAX_LENGTH = 128  # 一文あたりの最大トークン数
BERT_MODEL = "cl-tohoku/bert-base-japanese-v2"  # 使用する学習済みモデル
TAGGED_DATASET_PATH = "./dataset/ner_tagged.json"
MODEL_DIR = "./model"
LOG_DIR = "./logs"

# データの読み込み
with open(TAGGED_DATASET_PATH, 'r') as f:
  encoded_tagged_sentence_list = json.load(f)

# 3. データセットの作成
from sklearn.model_selection import train_test_split

class NERDataset(torch.utils.data.Dataset):
  def __init__(self, encoded_tagged_sentence_list):
    self.encoded_tagged_sentence_list = encoded_tagged_sentence_list

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

  def __getitem__(self, idx):
    # 辞書の値をTensorに変換
    item = {key: torch.tensor(val).to(device) for key, val in self.encoded_tagged_sentence_list[idx].items()}
    return item

# データを学習用、検証用に分割
train_encoded_tagged_sentence_list, eval_encoded_tagged_sentence_list = train_test_split(encoded_tagged_sentence_list)
# データセットに変換
train_data = NERDataset(train_encoded_tagged_sentence_list)
eval_data = NERDataset(eval_encoded_tagged_sentence_list)

# 4. Trainerの作成
from transformers import Trainer, TrainingArguments

import numpy as np
from datasets import load_metric

# 事前学習モデル
config = BertConfig.from_pretrained(BERT_MODEL, id2label=id2label, label2id=label2id)
model = BertForTokenClassification.from_pretrained(BERT_MODEL, config=config).to(device)
tokenizer = BertJapaneseTokenizer.from_pretrained(BERT_MODEL)

# 学習用パラメーター
training_args = TrainingArguments(
    output_dir = MODEL_DIR,
    num_train_epochs = 2,
    per_device_train_batch_size = 8,
    per_device_eval_batch_size = 32,
    warmup_steps = 500,  # 学習係数が0からこのステップ数で上昇
    weight_decay = 0.01,  # 重みの減衰率
    logging_dir = LOG_DIR,
)

metric = load_metric("seqeval")

def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    # ラベルのIDをラベルに変換
    predictions = [
        [id2label[p] for p in prediction] for prediction in predictions
    ]
    labels = [
        [id2label[l] for l in label] for label in labels
    ]

    results = metric.compute(predictions=predictions, references=labels)
    print(results)
    return {
        "precision": results["overall_precision"],
        "recall": results["overall_recall"],
        "f1": results["overall_f1"],
        "accuracy": results["overall_accuracy"],
    }

# Trainerの初期化
trainer = Trainer(
    model = model, # 学習対象のモデル
    args = training_args, # 学習用パラメーター
    compute_metrics = compute_metrics, # 評価用関数
    train_dataset = train_data, # 学習用データ
    eval_dataset = eval_data, # 検証用データ
    tokenizer = tokenizer
)

# 5. 学習
trainer.train()
trainer.evaluate()

trainer.save_model(MODEL_DIR)

from transformers import pipeline
from transformers import BertJapaneseTokenizer, BertForTokenClassification

MODEL_DIR = "./model"

model = BertForTokenClassification.from_pretrained(MODEL_DIR)
tokenizer = BertJapaneseTokenizer.from_pretrained(MODEL_DIR)

ner_pipeline = pipeline('ner', model=model, tokenizer=tokenizer)

ner_pipeline("株式会社はJurabi、東京都台東区に本社を置くIT企業である。")

[
  {'entity': 'B-法人名', 'score': 0.9811771, 'index': 4, 'word': 'Ju', 'start': None, 'end': None
  },
  {'entity': 'I-法人名', 'score': 0.9945182, 'index': 5, 'word': '##ra', 'start': None, 'end': None
  },
  {'entity': 'I-法人名', 'score': 0.9943316, 'index': 6, 'word': '##b', 'start': None, 'end': None
  },
  {'entity': 'I-法人名', 'score': 0.9925209, 'index': 7, 'word': '##i', 'start': None, 'end': None
  },
  {'entity': 'B-地名', 'score': 0.99534696, 'index': 9, 'word': '東京', 'start': None, 'end': None
  },
  {'entity': 'I-地名', 'score': 0.9967154, 'index': 10, 'word': '都', 'start': None, 'end': None
  },
  {'entity': 'I-地名', 'score': 0.996228, 'index': 11, 'word': '台東', 'start': None, 'end': None
  },
  {'entity': 'I-地名', 'score': 0.9965228, 'index': 12, 'word': '区', 'start': None, 'end': None
  }
]