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