train_degree_e2e.py (13059B)
1 import os, sys, json, logging, time, pprint, tqdm 2 import numpy as np 3 import torch 4 from torch.utils.data import DataLoader 5 from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup 6 from model import GenerativeModel 7 from data import GenDataset 8 from utils import Summarizer, compute_f1 9 from argparse import ArgumentParser, Namespace 10 import ipdb 11 12 # configuration 13 parser = ArgumentParser() 14 parser.add_argument('-c', '--config', required=True) 15 args = parser.parse_args() 16 with open(args.config) as fp: 17 config = json.load(fp) 18 config.update(args.__dict__) 19 config = Namespace(**config) 20 21 if config.dataset == "ace05e" or config.dataset == "ace05ep": 22 from template_generate_ace import eve_template_generator 23 template_file = "template_generate_ace" 24 elif config.dataset == "ere": 25 from template_generate_ere import eve_template_generator 26 template_file = "template_generate_ere" 27 28 # fix random seed 29 np.random.seed(config.seed) 30 torch.manual_seed(config.seed) 31 torch.backends.cudnn.enabled = False 32 33 # logger and summarizer 34 timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) 35 output_dir = os.path.join(config.output_dir, timestamp) 36 if not os.path.exists(output_dir): 37 os.makedirs(output_dir) 38 log_path = os.path.join(output_dir, "train.log") 39 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(message)s', datefmt='[%Y-%m-%d %H:%M:%S]', 40 handlers=[logging.FileHandler(os.path.join(output_dir, "train.log")), logging.StreamHandler()]) 41 logger = logging.getLogger(__name__) 42 logger.info(f"\n{pprint.pformat(vars(config), indent=4)}") 43 summarizer = Summarizer(output_dir) 44 45 # set GPU device 46 torch.cuda.set_device(config.gpu_device) 47 48 # check valid styles 49 assert np.all([style in ['event_type_sent', 'keywords', 'template'] for style in config.input_style]) 50 assert np.all([style in ['trigger:sentence', 'argument:sentence'] for style in config.output_style]) 51 52 # output 53 with open(os.path.join(output_dir, 'config.json'), 'w') as fp: 54 json.dump(vars(config), fp, indent=4) 55 best_model_path = os.path.join(output_dir, 'best_model.mdl') 56 dev_prediction_path = os.path.join(output_dir, 'pred.dev.json') 57 test_prediction_path = os.path.join(output_dir, 'pred.test.json') 58 59 # tokenizer 60 tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=config.cache_dir) 61 special_tokens = ['<Trigger>', '<sep>'] 62 tokenizer.add_tokens(special_tokens) 63 64 # load data 65 train_set = GenDataset(tokenizer, config.max_length, config.train_finetune_file, config.max_output_length) 66 dev_set = GenDataset(tokenizer, config.max_length, config.dev_finetune_file, config.max_output_length) 67 test_set = GenDataset(tokenizer, config.max_length, config.test_finetune_file, config.max_output_length) 68 train_batch_num = len(train_set) // config.train_batch_size + (len(train_set) % config.train_batch_size != 0) 69 dev_batch_num = len(dev_set) // config.eval_batch_size + (len(dev_set) % config.eval_batch_size != 0) 70 test_batch_num = len(test_set) // config.eval_batch_size + (len(test_set) % config.eval_batch_size != 0) 71 72 # initialize the model 73 model = GenerativeModel(config, tokenizer) 74 model.cuda(device=config.gpu_device) 75 76 # optimizer 77 param_groups = [{'params': model.parameters(), 'lr': config.learning_rate, 'weight_decay': config.weight_decay}] 78 optimizer = AdamW(params=param_groups) 79 schedule = get_linear_schedule_with_warmup(optimizer, 80 num_warmup_steps=train_batch_num*config.warmup_epoch, 81 num_training_steps=train_batch_num*config.max_epoch) 82 83 84 85 # start training 86 logger.info("Start training ...") 87 summarizer_step = 0 88 best_dev_epoch = -1 89 best_dev_scores = { 90 'tri_id': (0.0, 0.0, 0.0), 91 'arg_id': (0.0, 0.0, 0.0), 92 'arg_cls': (0.0, 0.0, 0.0) 93 } 94 for epoch in range(1, config.max_epoch+1): 95 logger.info(log_path) 96 logger.info(f"Epoch {epoch}") 97 98 # training 99 progress = tqdm.tqdm(total=train_batch_num, ncols=75, desc='Train {}'.format(epoch)) 100 model.train() 101 optimizer.zero_grad() 102 for batch_idx, batch in enumerate(DataLoader(train_set, batch_size=config.train_batch_size // config.accumulate_step, 103 shuffle=True, drop_last=False, collate_fn=train_set.collate_fn)): 104 105 # forard model 106 loss = model(batch) 107 108 # record loss 109 summarizer.scalar_summary('train/loss', loss, summarizer_step) 110 summarizer_step += 1 111 112 loss = loss * (1 / config.accumulate_step) 113 loss.backward() 114 115 if (batch_idx + 1) % config.accumulate_step == 0: 116 progress.update(1) 117 torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clipping) 118 optimizer.step() 119 schedule.step() 120 optimizer.zero_grad() 121 progress.close() 122 123 # eval dev set 124 progress = tqdm.tqdm(total=dev_batch_num, ncols=75, desc='Dev {}'.format(epoch)) 125 model.eval() 126 best_dev_flag = False 127 write_output = [] 128 dev_gold_tri_num, dev_pred_tri_num, dev_match_tri_num = 0, 0, 0 129 dev_gold_arg_num, dev_pred_arg_num, dev_match_arg_id, dev_match_arg_cls = 0, 0, 0, 0 130 131 for batch_idx, batch in enumerate(DataLoader(dev_set, batch_size=config.eval_batch_size, 132 shuffle=False, collate_fn=dev_set.collate_fn)): 133 progress.update(1) 134 pred_text = model.predict(batch, max_length=config.max_output_length) 135 gold_text = batch.target_text 136 input_text = batch.input_text 137 for i_text, g_text, p_text, info in zip(input_text, gold_text, pred_text, batch.infos): 138 theclass = getattr(sys.modules[template_file], info[1].replace(':', '_').replace('-', '_'), False) 139 assert theclass 140 template = theclass(config.input_style, config.output_style, info[2], info[1], info[0]) 141 142 # decode predictions 143 pred_object = template.decode(p_text) 144 gold_object = template.trigger_span + [_ for _ in template.get_converted_gold()] 145 146 # calculate scores 147 sub_scores = template.evaluate(pred_object) 148 dev_gold_tri_num += sub_scores['gold_tri_num'] 149 dev_pred_tri_num += sub_scores['pred_tri_num'] 150 dev_match_tri_num += sub_scores['match_tri_num'] 151 dev_gold_arg_num += sub_scores['gold_arg_num'] 152 dev_pred_arg_num += sub_scores['pred_arg_num'] 153 dev_match_arg_id += sub_scores['match_arg_id'] 154 dev_match_arg_cls += sub_scores['match_arg_cls'] 155 write_output.append({ 156 'input text': i_text, 157 'gold text': g_text, 158 'pred text': p_text, 159 'gold triggers': gold_object, 160 'pred triggers': pred_object, 161 'score': sub_scores, 162 'gold info': info 163 }) 164 progress.close() 165 166 dev_scores = { 167 'tri_id': compute_f1(dev_pred_tri_num, dev_gold_tri_num, dev_match_tri_num), 168 'arg_id': compute_f1(dev_pred_arg_num, dev_gold_arg_num, dev_match_arg_id), 169 'arg_cls': compute_f1(dev_pred_arg_num, dev_gold_arg_num, dev_match_arg_cls) 170 } 171 172 # print scores 173 print("---------------------------------------------------------------------") 174 print('Trigger I - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 175 dev_scores['tri_id'][0] * 100.0, dev_match_tri_num, dev_pred_tri_num, 176 dev_scores['tri_id'][1] * 100.0, dev_match_tri_num, dev_gold_tri_num, dev_scores['tri_id'][2] * 100.0)) 177 print("---------------------------------------------------------------------") 178 print('Role I - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 179 dev_scores['arg_id'][0] * 100.0, dev_match_arg_id, dev_pred_arg_num, 180 dev_scores['arg_id'][1] * 100.0, dev_match_arg_id, dev_gold_arg_num, dev_scores['arg_id'][2] * 100.0)) 181 print('Role C - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 182 dev_scores['arg_cls'][0] * 100.0, dev_match_arg_cls, dev_pred_arg_num, 183 dev_scores['arg_cls'][1] * 100.0, dev_match_arg_cls, dev_gold_arg_num, dev_scores['arg_cls'][2] * 100.0)) 184 print("---------------------------------------------------------------------") 185 186 # check best dev model 187 if dev_scores['arg_cls'][2] > best_dev_scores['arg_cls'][2]: 188 best_dev_flag = True 189 190 # if best dev, save model and evaluate test set 191 if best_dev_flag: 192 best_dev_scores = dev_scores 193 best_dev_epoch = epoch 194 195 # save best model 196 logger.info('Saving best model') 197 torch.save(model.state_dict(), best_model_path) 198 199 # save dev result 200 with open(dev_prediction_path, 'w') as fp: 201 json.dump(write_output, fp, indent=4) 202 203 # eval test set 204 progress = tqdm.tqdm(total=test_batch_num, ncols=75, desc='Test {}'.format(epoch)) 205 write_output = [] 206 test_gold_tri_num, test_pred_tri_num, test_match_tri_num = 0, 0, 0 207 test_gold_arg_num, test_pred_arg_num, test_match_arg_id, test_match_arg_cls = 0, 0, 0, 0 208 for batch_idx, batch in enumerate(DataLoader(test_set, batch_size=config.eval_batch_size, 209 shuffle=False, collate_fn=test_set.collate_fn)): 210 progress.update(1) 211 pred_text = model.predict(batch, max_length=config.max_output_length) 212 gold_text = batch.target_text 213 input_text = batch.input_text 214 for i_text, g_text, p_text, info in zip(input_text, gold_text, pred_text, batch.infos): 215 theclass = getattr(sys.modules[template_file], info[1].replace(':', '_').replace('-', '_'), False) 216 assert theclass 217 template = theclass(config.input_style, config.output_style, info[2], info[1], info[0]) 218 219 # decode predictions 220 pred_object = template.decode(p_text) 221 gold_object = template.trigger_span + [_ for _ in template.get_converted_gold()] 222 223 # calculate scores 224 sub_scores = template.evaluate(pred_object) 225 test_gold_tri_num += sub_scores['gold_tri_num'] 226 test_pred_tri_num += sub_scores['pred_tri_num'] 227 test_match_tri_num += sub_scores['match_tri_num'] 228 test_gold_arg_num += sub_scores['gold_arg_num'] 229 test_pred_arg_num += sub_scores['pred_arg_num'] 230 test_match_arg_id += sub_scores['match_arg_id'] 231 test_match_arg_cls += sub_scores['match_arg_cls'] 232 write_output.append({ 233 'input text': i_text, 234 'gold text': g_text, 235 'pred text': p_text, 236 'gold triggers': gold_object, 237 'pred triggers': pred_object, 238 'score': sub_scores, 239 'gold info': info 240 }) 241 progress.close() 242 243 test_scores = { 244 'tri_id': compute_f1(test_pred_tri_num, test_gold_tri_num, test_match_tri_num), 245 'arg_id': compute_f1(test_pred_arg_num, test_gold_arg_num, test_match_arg_id), 246 'arg_cls': compute_f1(test_pred_arg_num, test_gold_arg_num, test_match_arg_cls) 247 } 248 249 # print scores 250 print("---------------------------------------------------------------------") 251 print('Trigger I - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 252 test_scores['tri_id'][0] * 100.0, test_match_tri_num, test_pred_tri_num, 253 test_scores['tri_id'][1] * 100.0, test_match_tri_num, test_gold_tri_num, test_scores['tri_id'][2] * 100.0)) 254 print("---------------------------------------------------------------------") 255 print('Role I - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 256 test_scores['arg_id'][0] * 100.0, test_match_arg_id, test_pred_arg_num, 257 test_scores['arg_id'][1] * 100.0, test_match_arg_id, test_gold_arg_num, test_scores['arg_id'][2] * 100.0)) 258 print('Role C - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 259 test_scores['arg_cls'][0] * 100.0, test_match_arg_cls, test_pred_arg_num, 260 test_scores['arg_cls'][1] * 100.0, test_match_arg_cls, test_gold_arg_num, test_scores['arg_cls'][2] * 100.0)) 261 print("---------------------------------------------------------------------") 262 263 # save test result 264 with open(test_prediction_path, 'w') as fp: 265 json.dump(write_output, fp, indent=4) 266 267 logger.info({"epoch": epoch, "dev_scores": dev_scores}) 268 if best_dev_flag: 269 logger.info({"epoch": epoch, "test_scores": test_scores}) 270 logger.info("Current best") 271 logger.info({"best_epoch": best_dev_epoch, "best_scores": best_dev_scores}) 272 273 logger.info(log_path) 274 logger.info("Done!") 275