from collections import defaultdict import json import os import random import logging import time import torch import torch.nn as nn import numpy as np from transformers import BertTokenizer, AutoTokenizer, AdamW, get_linear_schedule_with_warmup from utils.argparse import ConfigurationParer from inputs.vocabulary import Vocabulary from inputs.fields.token_field import TokenField from inputs.fields.raw_token_field import RawTokenField from inputs.instance import Instance from inputs.datasets.dataset import Dataset from inputs.dataset_readers.oie_reader_for_relation_detection import ReaderForRelationDecoding from models.relation_decoding.relation_decoder import RelDecoder from utils.nn_utils import get_n_trainable_parameters logger = logging.getLogger(__name__) def step(cfg, model, batch_inputs, device): batch_inputs["tokens"] = torch.LongTensor(batch_inputs["tokens"]) batch_inputs["label_ids"] = torch.LongTensor(batch_inputs["label_ids"]) batch_inputs["label_ids_mask"] = torch.BoolTensor(batch_inputs["relation_ids_mask"]) batch_inputs["relation_ids"] = torch.LongTensor(batch_inputs["relation_ids"]) batch_inputs["relation_ids_mask"] = torch.BoolTensor(batch_inputs["relation_ids_mask"]) batch_inputs["argument_ids"] = torch.LongTensor(batch_inputs["argument_ids"]) batch_inputs["argument_ids_mask"] = torch.BoolTensor(batch_inputs["argument_ids_mask"]) batch_inputs["wordpiece_tokens"] = torch.LongTensor(batch_inputs["wordpiece_tokens"]) batch_inputs["wordpiece_tokens_mask"] = torch.BoolTensor(batch_inputs["wordpiece_tokens_mask"]) batch_inputs["wordpiece_tokens_index"] = torch.LongTensor(batch_inputs["wordpiece_tokens_index"]) batch_inputs["wordpiece_segment_ids"] = torch.LongTensor(batch_inputs["wordpiece_segment_ids"]) if device > -1: batch_inputs["tokens"] = batch_inputs["tokens"].cuda(device=device, non_blocking=True) batch_inputs["label_ids"] = batch_inputs["label_ids"].cuda(device=device, non_blocking=True) batch_inputs["label_ids_mask"] = batch_inputs["label_ids_mask"].cuda(device=device, non_blocking=True) batch_inputs["relation_ids"] = batch_inputs["relation_ids"].cuda(device=device,non_blocking=True) batch_inputs["relation_ids_mask"] = batch_inputs["relation_ids_mask"].cuda(device=device, non_blocking=True) batch_inputs["argument_ids"] = batch_inputs["argument_ids"].cuda(device=device,non_blocking=True) batch_inputs["argument_ids_mask"] = batch_inputs["argument_ids_mask"].cuda(device=device, non_blocking=True) batch_inputs["wordpiece_tokens"] = batch_inputs["wordpiece_tokens"].cuda(device=device, non_blocking=True) batch_inputs["wordpiece_tokens_mask"] = batch_inputs["wordpiece_tokens_mask"].cuda(device=device, non_blocking=True) batch_inputs["wordpiece_tokens_index"] = batch_inputs["wordpiece_tokens_index"].cuda(device=device, non_blocking=True) batch_inputs["wordpiece_segment_ids"] = batch_inputs["wordpiece_segment_ids"].cuda(device=device, non_blocking=True) outputs = model(batch_inputs) batch_outputs = [] if not model.training: for sent_idx in range(len(batch_inputs['tokens_lens'])): sent_output = dict() sent_output['tokens'] = batch_inputs['tokens'][sent_idx].cpu().numpy() sent_output['label_ids'] = batch_inputs['label_ids'][sent_idx].cpu().numpy() sent_output['relation_ids'] = batch_inputs['relation_ids'][sent_idx].cpu().numpy() sent_output['argument_ids'] = batch_inputs['argument_ids'][sent_idx].cpu().numpy() sent_output['seq_len'] = batch_inputs['tokens_lens'][sent_idx] sent_output['label_preds'] = outputs['label_preds'][sent_idx].cpu().numpy() batch_outputs.append(sent_output) return batch_outputs return outputs['loss'] def train(cfg, dataset, model): logger.info("Training starting...") for name, param in model.named_parameters(): logger.info("{!r}: size: {} requires_grad: {}.".format(name, param.size(), param.requires_grad)) logger.info("Trainable parameters size: {}.".format(get_n_trainable_parameters(model))) parameters = [(name, param) for name, param in model.named_parameters() if param.requires_grad] no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] bert_layer_lr = {} base_lr = cfg.bert_learning_rate for i in range(11, -1, -1): bert_layer_lr['.' + str(i) + '.'] = base_lr base_lr *= cfg.lr_decay_rate optimizer_grouped_parameters = [] for name, param in parameters: params = {'params': [param], 'lr': cfg.learning_rate} if any(item in name for item in no_decay): params['weight_decay_rate'] = 0.0 else: if 'bert' in name: params['weight_decay_rate'] = cfg.adam_bert_weight_decay_rate else: params['weight_decay_rate'] = cfg.adam_weight_decay_rate for bert_layer_name, lr in bert_layer_lr.items(): if bert_layer_name in name: params['lr'] = lr break optimizer_grouped_parameters.append(params) optimizer = AdamW(optimizer_grouped_parameters, betas=(cfg.adam_beta1, cfg.adam_beta2), lr=cfg.learning_rate, eps=cfg.adam_epsilon, weight_decay=cfg.adam_weight_decay_rate, correct_bias=False) total_train_steps = (dataset.get_dataset_size("train") + cfg.train_batch_size * cfg.gradient_accumulation_steps - 1) / (cfg.train_batch_size * cfg.gradient_accumulation_steps) * cfg.epochs num_warmup_steps = int(cfg.warmup_rate * total_train_steps) + 1 scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=total_train_steps) last_epoch = 1 batch_id = 0 best_f1 = 0.0 early_stop_cnt = 0 accumulation_steps = 0 model.zero_grad() for epoch, batch in dataset.get_batch('train', cfg.train_batch_size, None): if last_epoch != epoch or (batch_id != 0 and batch_id % cfg.validate_every == 0): if accumulation_steps != 0: optimizer.step() scheduler.step() model.zero_grad() if epoch > cfg.pretrain_epochs: dev_f1 = dev(cfg, dataset, model) if dev_f1 > best_f1: early_stop_cnt = 0 best_f1 = dev_f1 logger.info("Save model...") torch.save(model.state_dict(), open(cfg.best_model_path, "wb")) elif last_epoch != epoch: early_stop_cnt += 1 if early_stop_cnt > cfg.early_stop: logger.info("Early Stop: best F1 score: {:6.2f}%".format(100 * best_f1)) break if epoch > cfg.epochs: torch.save(model.state_dict(), open(cfg.last_model_path, "wb")) logger.info("Training Stop: best F1 score: {:6.2f}%".format(100 * best_f1)) break if last_epoch != epoch: batch_id = 0 last_epoch = epoch model.train() batch_id += len(batch['tokens_lens']) batch['epoch'] = (epoch - 1) loss = step(cfg, model, batch, cfg.device) if batch_id % cfg.logging_steps == 0: logger.info("Epoch: {} Batch: {} Loss: {})".format(epoch, batch_id, loss.item())) if cfg.gradient_accumulation_steps > 1: loss /= cfg.gradient_accumulation_steps loss.backward() accumulation_steps = (accumulation_steps + 1) % cfg.gradient_accumulation_steps if accumulation_steps == 0: nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=cfg.gradient_clipping) optimizer.step() scheduler.step() model.zero_grad() state_dict = torch.load(open(cfg.best_model_path, "rb"), map_location=lambda storage, loc: storage) model.load_state_dict(state_dict) test(cfg, dataset, model) def simple_accuracy(preds, labels): return (preds == labels).mean() def compute_f1(output): n_gold = n_pred = n_correct = 0 for sent in output: for pred, label in zip(sent["label_preds"], sent["label_ids"]): if pred != 0: n_pred += 1 if label != 0: n_gold += 1 if (pred != 0) and (label != 0) and (pred == label): n_correct += 1 if n_correct == 0: return {'precision': 0.0, 'recall': 0.0, 'f1': 0.0} else: prec = n_correct * 1.0 / n_pred recall = n_correct * 1.0 / n_gold if prec + recall > 0: f1 = 2.0 * prec * recall / (prec + recall) else: f1 = 0.0 return {'precision': prec, 'task_recall': recall, 'task_f1': f1, 'n_correct': n_correct, 'n_pred': n_pred, 'task_ngold': n_gold} def evaluate(outputs): result = compute_f1(outputs) logger.info("Validation F1: {}, Accuracy: {}, Recall: {}".format(result["task_f1"], result["precision"], result["task_recall"])) return result["task_f1"] def dev(cfg, dataset, model): logger.info("Validate starting...") model.zero_grad() all_outputs = [] cost_time = 0 for _, batch in dataset.get_batch('dev', cfg.test_batch_size, None): model.eval() with torch.no_grad(): cost_time -= time.time() batch_outpus = step(cfg, model, batch, cfg.device) cost_time += time.time() all_outputs.extend(batch_outpus) logger.info(f"Cost time: {cost_time}s") f1 = evaluate(all_outputs) return f1 def test(cfg, dataset, model): logger.info("Testing starting...") model.zero_grad() all_outputs = [] cost_time = 0 for _, batch in dataset.get_batch('test', cfg.test_batch_size, None): model.eval() with torch.no_grad(): cost_time -= time.time() batch_outpus = step(cfg, model, batch, cfg.device) cost_time += time.time() all_outputs.extend(batch_outpus) logger.info(f"Cost time: {cost_time}s") f1 = evaluate(all_outputs) print("test F1: ", f1) def main(): # config settings parser = ConfigurationParer() parser.add_save_cfgs() parser.add_data_cfgs() parser.add_model_cfgs() parser.add_optimizer_cfgs() parser.add_run_cfgs() cfg = parser.parse_args() logger.info(parser.format_values()) # set random seed random.seed(cfg.seed) torch.manual_seed(cfg.seed) np.random.seed(cfg.seed) if cfg.device > -1 and not torch.cuda.is_available(): logger.error('config conflicts: no gpu available, use cpu for training.') cfg.device = -1 if cfg.device > -1: torch.cuda.manual_seed(cfg.seed) # define fields tokens = TokenField("tokens", "tokens", "tokens", True) label_ids = RawTokenField("label_ids", "label_ids") relation_ids = RawTokenField("relation_ids", "relation_ids") argument_ids = RawTokenField("argument_ids", "argument_ids") wordpiece_tokens = TokenField("wordpiece_tokens", "wordpiece", "wordpiece_tokens", False) wordpiece_tokens_index = RawTokenField("wordpiece_tokens_index", "wordpiece_tokens_index") wordpiece_segment_ids = RawTokenField("wordpiece_segment_ids", "wordpiece_segment_ids") fields = [tokens, label_ids, relation_ids, argument_ids] if cfg.embedding_model in ['bert', 'pretrained']: fields.extend([wordpiece_tokens, wordpiece_tokens_index, wordpiece_segment_ids]) # define counter and vocabulary counter = defaultdict(lambda: defaultdict(int)) vocab = Vocabulary() # define instance (data sets) train_instance = Instance(fields) dev_instance = Instance(fields) test_instance = Instance(fields) # define dataset reader max_len = {'tokens': cfg.max_sent_len, 'wordpiece_tokens': cfg.max_wordpiece_len} ent_rel_file = json.load(open(cfg.ent_rel_file, 'r', encoding='utf-8')) pretrained_vocab = {'ent_rel_id': ent_rel_file["id"]} if cfg.embedding_model == 'bert': tokenizer = BertTokenizer.from_pretrained(cfg.bert_model_name) logger.info("Load bert tokenizer successfully.") pretrained_vocab['wordpiece'] = tokenizer.get_vocab() elif cfg.embedding_model == 'pretrained': tokenizer = AutoTokenizer.from_pretrained(cfg.pretrained_model_name) logger.info("Load {} tokenizer successfully.".format(cfg.pretrained_model_name)) pretrained_vocab['wordpiece'] = tokenizer.get_vocab() train_reader = ReaderForRelationDecoding(cfg.train_file, False, max_len) dev_reader = ReaderForRelationDecoding(cfg.dev_file, False, max_len) test_reader = ReaderForRelationDecoding(cfg.test_file, False, max_len) # define dataset oie_dataset = Dataset("OIE4") oie_dataset.add_instance("train", train_instance, train_reader, is_count=True, is_train=True) oie_dataset.add_instance("dev", dev_instance, dev_reader, is_count=True, is_train=False) oie_dataset.add_instance("test", test_instance, test_reader, is_count=True, is_train=False) min_count = {"tokens": 1} no_pad_namespace = ["ent_rel_id"] no_unk_namespace = ["ent_rel_id"] contain_pad_namespace = {"wordpiece": tokenizer.pad_token} contain_unk_namespace = {"wordpiece": tokenizer.unk_token} oie_dataset.build_dataset(vocab=vocab, counter=counter, min_count=min_count, pretrained_vocab=pretrained_vocab, no_pad_namespace=no_pad_namespace, no_unk_namespace=no_unk_namespace, contain_pad_namespace=contain_pad_namespace, contain_unk_namespace=contain_unk_namespace) oie_dataset.set_wo_padding_namespace(wo_padding_namespace=[]) if cfg.test: vocab = Vocabulary.load(cfg.relation_vocab) else: vocab.save(cfg.relation_vocab) # rel model model = RelDecoder(cfg=cfg, vocab=vocab, ent_rel_file=ent_rel_file) if cfg.test and os.path.exists(cfg.best_model_path): state_dict = torch.load(open(cfg.best_model_path, 'rb'), map_location=lambda storage, loc: storage) model.load_state_dict(state_dict) logger.info("Loading best training model {} successfully for testing.".format(cfg.best_model_path)) if cfg.device > -1: model.cuda(device=cfg.device) if cfg.test: dev(cfg, oie_dataset, model) test(cfg, oie_dataset, model) else: train(cfg, oie_dataset, model) if __name__ == '__main__': main()