train_degree_ed.py (10511B)
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'] 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 } 92 for epoch in range(1, config.max_epoch+1): 93 logger.info(log_path) 94 logger.info(f"Epoch {epoch}") 95 96 # training 97 progress = tqdm.tqdm(total=train_batch_num, ncols=75, desc='Train {}'.format(epoch)) 98 model.train() 99 optimizer.zero_grad() 100 for batch_idx, batch in enumerate(DataLoader(train_set, batch_size=config.train_batch_size // config.accumulate_step, 101 shuffle=True, drop_last=False, collate_fn=train_set.collate_fn)): 102 103 # forard model 104 loss = model(batch) 105 106 # record loss 107 summarizer.scalar_summary('train/loss', loss, summarizer_step) 108 summarizer_step += 1 109 110 loss = loss * (1 / config.accumulate_step) 111 loss.backward() 112 113 if (batch_idx + 1) % config.accumulate_step == 0: 114 progress.update(1) 115 torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clipping) 116 optimizer.step() 117 schedule.step() 118 optimizer.zero_grad() 119 progress.close() 120 121 # eval dev set 122 progress = tqdm.tqdm(total=dev_batch_num, ncols=75, desc='Dev {}'.format(epoch)) 123 model.eval() 124 best_dev_flag = False 125 write_output = [] 126 dev_gold_tri_num, dev_pred_tri_num, dev_match_tri_num = 0, 0, 0 127 128 for batch_idx, batch in enumerate(DataLoader(dev_set, batch_size=config.eval_batch_size, 129 shuffle=False, collate_fn=dev_set.collate_fn)): 130 progress.update(1) 131 pred_text = model.predict(batch, max_length=config.max_output_length) 132 gold_text = batch.target_text 133 input_text = batch.input_text 134 for i_text, g_text, p_text, info in zip(input_text, gold_text, pred_text, batch.infos): 135 theclass = getattr(sys.modules[template_file], info[1].replace(':', '_').replace('-', '_'), False) 136 assert theclass 137 template = theclass(config.input_style, config.output_style, info[2], info[1], info[0]) 138 139 140 # decode predictions 141 pred_object = template.decode(p_text) 142 gold_object = template.trigger_span 143 144 # calculate scores 145 sub_scores = template.evaluate(pred_object) 146 dev_gold_tri_num += sub_scores['gold_tri_num'] 147 dev_pred_tri_num += sub_scores['pred_tri_num'] 148 dev_match_tri_num += sub_scores['match_tri_num'] 149 write_output.append({ 150 'input text': i_text, 151 'gold text': g_text, 152 'pred text': p_text, 153 'gold triggers': gold_object, 154 'pred triggers': pred_object, 155 'score': sub_scores, 156 'gold info': info 157 }) 158 progress.close() 159 160 dev_scores = { 161 'tri_id': compute_f1(dev_pred_tri_num, dev_gold_tri_num, dev_match_tri_num), 162 } 163 164 # print scores 165 print("---------------------------------------------------------------------") 166 print('Trigger I - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 167 dev_scores['tri_id'][0] * 100.0, dev_match_tri_num, dev_pred_tri_num, 168 dev_scores['tri_id'][1] * 100.0, dev_match_tri_num, dev_gold_tri_num, dev_scores['tri_id'][2] * 100.0)) 169 print("---------------------------------------------------------------------") 170 171 # check best dev model 172 if dev_scores['tri_id'][2] > best_dev_scores['tri_id'][2]: 173 best_dev_flag = True 174 175 # if best dev, save model and evaluate test set 176 if best_dev_flag: 177 best_dev_scores = dev_scores 178 best_dev_epoch = epoch 179 180 # save best model 181 logger.info('Saving best model') 182 torch.save(model.state_dict(), best_model_path) 183 184 # save dev result 185 with open(dev_prediction_path, 'w') as fp: 186 json.dump(write_output, fp, indent=4) 187 188 # eval test set 189 progress = tqdm.tqdm(total=test_batch_num, ncols=75, desc='Test {}'.format(epoch)) 190 write_output = [] 191 test_gold_tri_num, test_pred_tri_num, test_match_tri_num = 0, 0, 0 192 for batch_idx, batch in enumerate(DataLoader(test_set, batch_size=config.eval_batch_size, 193 shuffle=False, collate_fn=test_set.collate_fn)): 194 progress.update(1) 195 pred_text = model.predict(batch, max_length=config.max_output_length) 196 gold_text = batch.target_text 197 input_text = batch.input_text 198 for i_text, g_text, p_text, info in zip(input_text, gold_text, pred_text, batch.infos): 199 theclass = getattr(sys.modules[template_file], info[1].replace(':', '_').replace('-', '_'), False) 200 assert theclass 201 template = theclass(config.input_style, config.output_style, info[2], info[1], info[0]) 202 203 204 # decode predictions 205 pred_object = template.decode(p_text) 206 gold_object = template.trigger_span 207 208 # calculate scores 209 sub_scores = template.evaluate(pred_object) 210 test_gold_tri_num += sub_scores['gold_tri_num'] 211 test_pred_tri_num += sub_scores['pred_tri_num'] 212 test_match_tri_num += sub_scores['match_tri_num'] 213 write_output.append({ 214 'input text': i_text, 215 'gold text': g_text, 216 'pred text': p_text, 217 'gold triggers': gold_object, 218 'pred triggers': pred_object, 219 'score': sub_scores, 220 'gold info': info 221 }) 222 progress.close() 223 224 test_scores = { 225 'tri_id': compute_f1(test_pred_tri_num, test_gold_tri_num, test_match_tri_num), 226 } 227 228 # print scores 229 print("---------------------------------------------------------------------") 230 print('Trigger I - P: {:5.2f} ({:4d}/{:4d}), R: {:5.2f} ({:4d}/{:4d}), F: {:5.2f}'.format( 231 test_scores['tri_id'][0] * 100.0, test_match_tri_num, test_pred_tri_num, 232 test_scores['tri_id'][1] * 100.0, test_match_tri_num, test_gold_tri_num, test_scores['tri_id'][2] * 100.0)) 233 print("---------------------------------------------------------------------") 234 235 # save test result 236 with open(test_prediction_path, 'w') as fp: 237 json.dump(write_output, fp, indent=4) 238 239 logger.info({"epoch": epoch, "dev_scores": dev_scores}) 240 if best_dev_flag: 241 logger.info({"epoch": epoch, "test_scores": test_scores}) 242 logger.info("Current best") 243 logger.info({"best_epoch": best_dev_epoch, "best_scores": best_dev_scores}) 244 245 logger.info(log_path) 246 logger.info("Done!") 247