eval_end2endEE.py (19485B)
1 import os, sys, json, logging, pprint, tqdm 2 import numpy as np 3 import torch 4 from torch.utils.data import DataLoader 5 from transformers import AutoTokenizer 6 from model import GenerativeModel 7 from data import GenDataset, EEDataset 8 from utils import compute_f1 9 from argparse import ArgumentParser, Namespace 10 import ipdb 11 12 # configuration 13 parser = ArgumentParser() 14 parser.add_argument('-c', '--e2e_config', required=True) 15 parser.add_argument('-e', '--e2e_model', required=True) 16 parser.add_argument('--no_dev', action='store_true', default=False) 17 parser.add_argument('--eval_batch_size', type=int) 18 parser.add_argument('--write_file', type=str) 19 args = parser.parse_args() 20 with open(args.e2e_config) as fp: 21 config = json.load(fp) 22 config = Namespace(**config) 23 24 if config.dataset == "ace05e" or config.dataset == "ace05ep": 25 from template_generate_ace import eve_template_generator 26 template_file = "template_generate_ace" 27 elif config.dataset == "ere": 28 from template_generate_ere import eve_template_generator 29 template_file = "template_generate_ere" 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 37 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(message)s', datefmt='[%Y-%m-%d %H:%M:%S]') 38 logger = logging.getLogger(__name__) 39 40 def get_span_idx(pieces, token_start_idxs, span, tokenizer, trigger_span=None): 41 """ 42 This function is how we map the generated prediction back to span prediction. 43 44 Detailed Explanation: 45 We will first split our prediction and use tokenizer to tokenize our predicted "span" into pieces. Then, we will find whether we can find a continuous span in the original "pieces" can match tokenized "span". 46 47 If it is an argument/relation extraction task, we will return the one which is closest to the trigger_span. 48 """ 49 words = [] 50 for s in span.split(' '): 51 words.extend(tokenizer.encode(s, add_special_tokens=False)) 52 53 candidates = [] 54 for i in range(len(pieces)): 55 j = 0 56 k = 0 57 while j < len(words) and i+k < len(pieces): 58 if pieces[i+k] == words[j]: 59 j += 1 60 k += 1 61 elif tokenizer.decode(words[j]) == "": 62 j += 1 63 elif tokenizer.decode(pieces[i+k]) == "": 64 k += 1 65 else: 66 break 67 if j == len(words): 68 candidates.append((i, i+k)) 69 70 candidates = [(token_start_idxs.index(c1), token_start_idxs.index(c2)) for c1, c2 in candidates if c1 in token_start_idxs and c2 in token_start_idxs] 71 if len(candidates) < 1: 72 return -1, -1 73 else: 74 if trigger_span is None: 75 return candidates[0] 76 else: 77 return sorted(candidates, key=lambda x: np.abs(trigger_span[0]-x[0]))[0] 78 79 def get_span_idx_tri(pieces, token_start_idxs, span, tokenizer, trigger_span=None): 80 """ 81 This function is how we map the generated prediction back to span prediction. 82 83 Detailed Explanation: 84 We will first split our prediction and use tokenizer to tokenize our predicted "span" into pieces. Then, we will find whether we can find a continuous span in the original "pieces" can match tokenized "span". 85 86 If it is an argument/relation extraction task, we will return the one which is closest to the trigger_span. 87 """ 88 words = [] 89 for s in span.split(' '): 90 words.extend(tokenizer.encode(s, add_special_tokens=False)) 91 92 candidates = [] 93 for i in range(len(pieces)): 94 j = 0 95 k = 0 96 while j < len(words) and i+k < len(pieces): 97 if pieces[i+k] == words[j]: 98 j += 1 99 k += 1 100 elif tokenizer.decode(words[j]) == "": 101 j += 1 102 elif tokenizer.decode(pieces[i+k]) == "": 103 k += 1 104 else: 105 break 106 if j == len(words): 107 candidates.append((i, i+k)) 108 109 candidates = [(token_start_idxs.index(c1), token_start_idxs.index(c2)) for c1, c2 in candidates if c1 in token_start_idxs and c2 in token_start_idxs] 110 if len(candidates) < 1: 111 return [(-1, -1)] 112 else: 113 if trigger_span is None: 114 return candidates 115 else: 116 return sorted(candidates, key=lambda x: np.abs(trigger_span[0]-x[0])) 117 118 def cal_scores(gold_triggers, pred_triggers, gold_roles, pred_roles): 119 assert len(gold_triggers) == len(pred_triggers) 120 assert len(gold_roles) == len(pred_roles) 121 # tri_id 122 gold_tri_id_num, pred_tri_id_num, match_tri_id_num = 0, 0, 0 123 for gold_trigger, pred_trigger in zip(gold_triggers, pred_triggers): 124 gold_set = set([(t[0], t[1]) for t in gold_trigger]) 125 pred_set = set([(t[0], t[1]) for t in pred_trigger]) 126 gold_tri_id_num += len(gold_set) 127 pred_tri_id_num += len(pred_set) 128 match_tri_id_num += len(gold_set & pred_set) 129 130 # tri_cls 131 gold_tri_cls_num, pred_tri_cls_num, match_tri_cls_num = 0, 0, 0 132 for gold_trigger, pred_trigger in zip(gold_triggers, pred_triggers): 133 gold_set = set(gold_trigger) 134 pred_set = set(pred_trigger) 135 gold_tri_cls_num += len(gold_set) 136 pred_tri_cls_num += len(pred_set) 137 match_tri_cls_num += len(gold_set & pred_set) 138 139 # arg_id 140 gold_arg_id_num, pred_arg_id_num, match_arg_id_num = 0, 0, 0 141 for gold_role, pred_role in zip(gold_roles, pred_roles): 142 gold_set = set([(r[0][2],)+r[1][:-1] for r in gold_role]) 143 pred_set = set([(r[0][2],)+r[1][:-1] for r in pred_role]) 144 145 gold_arg_id_num += len(gold_set) 146 pred_arg_id_num += len(pred_set) 147 match_arg_id_num += len(gold_set & pred_set) 148 149 # arg_cls 150 gold_arg_cls_num, pred_arg_cls_num, match_arg_cls_num = 0, 0, 0 151 for gold_role, pred_role in zip(gold_roles, pred_roles): 152 gold_set = set([(r[0][2],)+r[1] for r in gold_role]) 153 pred_set = set([(r[0][2],)+r[1] for r in pred_role]) 154 155 gold_arg_cls_num += len(gold_set) 156 pred_arg_cls_num += len(pred_set) 157 match_arg_cls_num += len(gold_set & pred_set) 158 159 scores = { 160 'tri_id': (gold_tri_id_num, pred_tri_id_num, match_tri_id_num) + compute_f1(pred_tri_id_num, gold_tri_id_num, match_tri_id_num), 161 'tri_cls': (gold_tri_cls_num, pred_tri_cls_num, match_tri_cls_num) + compute_f1(pred_tri_cls_num, gold_tri_cls_num, match_tri_cls_num), 162 'arg_id': (gold_arg_id_num, pred_arg_id_num, match_arg_id_num) + compute_f1(pred_arg_id_num, gold_arg_id_num, match_arg_id_num), 163 'arg_cls': (gold_arg_cls_num, pred_arg_cls_num, match_arg_cls_num) + compute_f1(pred_arg_cls_num, gold_arg_cls_num, match_arg_cls_num), 164 } 165 166 return scores 167 168 # set GPU device 169 torch.cuda.set_device(config.gpu_device) 170 171 # check valid styles 172 assert np.all([style in ['event_type_sent', 'keywords', 'template'] for style in config.input_style]) 173 assert np.all([style in ['trigger:sentence', 'argument:sentence'] for style in config.output_style]) 174 175 # tokenizer 176 tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=config.cache_dir) 177 special_tokens = ['<Trigger>', '<sep>'] 178 tokenizer.add_tokens(special_tokens) 179 180 if args.eval_batch_size: 181 config.eval_batch_size=args.eval_batch_size 182 183 # load data 184 dev_set = EEDataset(tokenizer, config.dev_file, max_length=config.max_length) 185 test_set = EEDataset(tokenizer, config.test_file, max_length=config.max_length) 186 dev_batch_num = len(dev_set) // config.eval_batch_size + (len(dev_set) % config.eval_batch_size != 0) 187 test_batch_num = len(test_set) // config.eval_batch_size + (len(test_set) % config.eval_batch_size != 0) 188 with open(config.vocab_file) as f: 189 vocab = json.load(f) 190 191 # load model 192 logger.info(f"Loading model from {args.e2e_model}") 193 model = GenerativeModel(config, tokenizer) 194 model.load_state_dict(torch.load(args.e2e_model, map_location=f'cuda:{config.gpu_device}')) 195 model.cuda(device=config.gpu_device) 196 model.eval() 197 198 # eval dev set 199 if not args.no_dev: 200 progress = tqdm.tqdm(total=dev_batch_num, ncols=75, desc='Dev') 201 dev_gold_triggers, dev_gold_roles, dev_pred_triggers, dev_pred_roles = [], [], [], [] 202 203 for batch in DataLoader(dev_set, batch_size=config.eval_batch_size, shuffle=False, collate_fn=dev_set.collate_fn): 204 progress.update(1) 205 206 p_triggers = [[] for _ in range(len(batch.tokens))] 207 p_roles = [[] for _ in range(len(batch.tokens))] 208 for event_type in vocab['event_type_itos']: 209 theclass = getattr(sys.modules[template_file], event_type.replace(':', '_').replace('-', '_'), False) 210 211 inputs = [] 212 for tokens in batch.tokens: 213 template = theclass(config.input_style, config.output_style, tokens, event_type) 214 inputs.append(template.generate_input_str('')) 215 216 inputs = tokenizer(inputs, return_tensors='pt', padding=True, max_length=config.max_length) 217 enc_idxs = inputs['input_ids'].cuda() 218 enc_attn = inputs['attention_mask'].cuda() 219 220 outputs = model.model.generate(input_ids=enc_idxs, attention_mask=enc_attn, num_beams=config.beam_size, max_length=config.max_output_length) 221 final_outputs = [tokenizer.decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=True) for output in outputs] 222 223 for bid, (tokens, p_text) in enumerate(zip(batch.tokens, final_outputs)): 224 template = theclass(config.input_style, config.output_style, tokens, event_type) 225 pred_object = template.decode(p_text) 226 227 pred_trigger_object = [] 228 pred_argument_object = [] 229 for obj in pred_object: 230 if obj[1] == event_type: 231 pred_trigger_object.append(obj) 232 else: 233 pred_argument_object.append(obj) 234 235 # decode triggers 236 triggers_ = [mention + (event_type, kwargs) for span, _, kwargs in pred_trigger_object for mention in get_span_idx_tri(batch.piece_idxs[bid], batch.token_start_idxs[bid], span, tokenizer)] 237 triggers_ = [t for t in triggers_ if t[0] != -1] 238 p_triggers_ = [t[:-1] for t in triggers_] 239 p_triggers_ = list(set(p_triggers_)) 240 p_triggers[bid].extend(p_triggers_) 241 242 # decode arguments 243 tri_id2obj = {} 244 for t in triggers_: 245 try: 246 tri_id2obj[t[3]['tri counter']] = (t[0], t[1], t[2]) 247 except: 248 ipdb.set_trace() 249 250 roles_ = [] 251 for span, role_type, kwargs in pred_argument_object: 252 corres_tri_id = kwargs['cor tri cnt'] 253 if corres_tri_id in tri_id2obj.keys(): 254 arg_span = get_span_idx(batch.piece_idxs[bid], batch.token_start_idxs[bid], span, tokenizer, tri_id2obj[corres_tri_id]) 255 if arg_span[0] != -1: 256 roles_.append((tri_id2obj[corres_tri_id], (arg_span[0], arg_span[1], role_type))) 257 else: 258 arg_span = get_span_idx(batch.piece_idxs[bid], batch.token_start_idxs[bid], span, tokenizer) 259 if arg_span[0] != -1: 260 roles_.append(((0, 1, event_type), (arg_span[0], arg_span[1], role_type))) 261 262 p_roles[bid].extend(roles_) 263 264 p_roles = [list(set(role)) for role in p_roles] 265 266 dev_gold_triggers.extend(batch.triggers) 267 dev_gold_roles.extend(batch.roles) 268 dev_pred_triggers.extend(p_triggers) 269 dev_pred_roles.extend(p_roles) 270 271 progress.close() 272 273 # calculate scores 274 dev_scores = cal_scores(dev_gold_triggers, dev_pred_triggers, dev_gold_roles, dev_pred_roles) 275 276 print("---------------------------------------------------------------------") 277 print('Trigger I - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 278 dev_scores['tri_id'][3] * 100.0, dev_scores['tri_id'][2], dev_scores['tri_id'][1], 279 dev_scores['tri_id'][4] * 100.0, dev_scores['tri_id'][2], dev_scores['tri_id'][0], dev_scores['tri_id'][5] * 100.0)) 280 print('Trigger C - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 281 dev_scores['tri_cls'][3] * 100.0, dev_scores['tri_cls'][2], dev_scores['tri_cls'][1], 282 dev_scores['tri_cls'][4] * 100.0, dev_scores['tri_cls'][2], dev_scores['tri_cls'][0], dev_scores['tri_cls'][5] * 100.0)) 283 print("---------------------------------------------------------------------") 284 print('Role I - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 285 dev_scores['arg_id'][3] * 100.0, dev_scores['arg_id'][2], dev_scores['arg_id'][1], 286 dev_scores['arg_id'][4] * 100.0, dev_scores['arg_id'][2], dev_scores['arg_id'][0], dev_scores['arg_id'][5] * 100.0)) 287 print('Role C - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 288 dev_scores['arg_cls'][3] * 100.0, dev_scores['arg_cls'][2], dev_scores['arg_cls'][1], 289 dev_scores['arg_cls'][4] * 100.0, dev_scores['arg_cls'][2], dev_scores['arg_cls'][0], dev_scores['arg_cls'][5] * 100.0)) 290 print("---------------------------------------------------------------------") 291 292 293 # test set 294 progress = tqdm.tqdm(total=test_batch_num, ncols=75, desc='Test') 295 test_gold_triggers, test_gold_roles, test_pred_triggers, test_pred_roles = [], [], [], [] 296 write_object = [] 297 for batch in DataLoader(test_set, batch_size=config.eval_batch_size, shuffle=False, collate_fn=test_set.collate_fn): 298 progress.update(1) 299 p_triggers = [[] for _ in range(len(batch.tokens))] 300 p_roles = [[] for _ in range(len(batch.tokens))] 301 p_texts = [[] for _ in range(len(batch.tokens))] 302 for event_type in vocab['event_type_itos']: 303 theclass = getattr(sys.modules[template_file], event_type.replace(':', '_').replace('-', '_'), False) 304 305 inputs = [] 306 for tokens in batch.tokens: 307 template = theclass(config.input_style, config.output_style, tokens, event_type) 308 inputs.append(template.generate_input_str('')) 309 310 inputs = tokenizer(inputs, return_tensors='pt', padding=True, max_length=config.max_length) 311 enc_idxs = inputs['input_ids'].cuda() 312 enc_attn = inputs['attention_mask'].cuda() 313 314 outputs = model.model.generate(input_ids=enc_idxs, attention_mask=enc_attn, num_beams=config.beam_size, max_length=config.max_output_length) 315 final_outputs = [tokenizer.decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=True) for output in outputs] 316 317 for bid, (tokens, p_text) in enumerate(zip(batch.tokens, final_outputs)): 318 template = theclass(config.input_style, config.output_style, tokens, event_type) 319 pred_object = template.decode(p_text) 320 321 pred_trigger_object = [] 322 pred_argument_object = [] 323 for obj in pred_object: 324 if obj[1] == event_type: 325 pred_trigger_object.append(obj) 326 else: 327 pred_argument_object.append(obj) 328 329 # decode triggers 330 triggers_ = [mention + (event_type, kwargs) for span, _, kwargs in pred_trigger_object for mention in get_span_idx_tri(batch.piece_idxs[bid], batch.token_start_idxs[bid], span, tokenizer)] 331 332 triggers_ = [t for t in triggers_ if t[0] != -1] 333 p_triggers_ = [t[:-1] for t in triggers_] 334 p_triggers_ = list(set(p_triggers_)) 335 336 p_triggers[bid].extend(p_triggers_) 337 338 # decode arguments 339 tri_id2obj = {} 340 for t in triggers_: 341 tri_id2obj[t[3]['tri counter']] = (t[0], t[1], t[2]) 342 343 roles_ = [] 344 for span, role_type, kwargs in pred_argument_object: 345 corres_tri_id = kwargs['cor tri cnt'] 346 if corres_tri_id in tri_id2obj.keys(): 347 arg_span = get_span_idx(batch.piece_idxs[bid], batch.token_start_idxs[bid], span, tokenizer, tri_id2obj[corres_tri_id]) 348 if arg_span[0] != -1: 349 roles_.append((tri_id2obj[corres_tri_id], (arg_span[0], arg_span[1], role_type))) 350 else: 351 arg_span = get_span_idx(batch.piece_idxs[bid], batch.token_start_idxs[bid], span, tokenizer) 352 if arg_span[0] != -1: 353 roles_.append(((0, 1, event_type), (arg_span[0], arg_span[1], role_type))) 354 355 p_roles[bid].extend(roles_) 356 p_texts[bid].append(p_text) 357 358 p_roles = [list(set(role)) for role in p_roles] 359 360 if config.ignore_first_header: 361 for bid, wnd_id in enumerate(batch.wnd_ids): 362 if int(wnd_id.split('-')[-1]) < 4: 363 p_triggers[bid] = [] 364 p_roles[bid] = [] 365 366 test_gold_triggers.extend(batch.triggers) 367 test_gold_roles.extend(batch.roles) 368 test_pred_triggers.extend(p_triggers) 369 test_pred_roles.extend(p_roles) 370 for gt, gr, pt, pr, te in zip(batch.triggers, batch.roles, p_triggers, p_roles, p_texts): 371 write_object.append({ 372 "pred text": te, 373 "pred triggers": pt, 374 "gold triggers": gt, 375 "pred roles": pr, 376 "gold roles": gr 377 }) 378 379 progress.close() 380 381 # calculate scores 382 test_scores = cal_scores(test_gold_triggers, test_pred_triggers, test_gold_roles, test_pred_roles) 383 384 print("---------------------------------------------------------------------") 385 print('Trigger I - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 386 test_scores['tri_id'][3] * 100.0, test_scores['tri_id'][2], test_scores['tri_id'][1], 387 test_scores['tri_id'][4] * 100.0, test_scores['tri_id'][2], test_scores['tri_id'][0], test_scores['tri_id'][5] * 100.0)) 388 print('Trigger C - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 389 test_scores['tri_cls'][3] * 100.0, test_scores['tri_cls'][2], test_scores['tri_cls'][1], 390 test_scores['tri_cls'][4] * 100.0, test_scores['tri_cls'][2], test_scores['tri_cls'][0], test_scores['tri_cls'][5] * 100.0)) 391 print("---------------------------------------------------------------------") 392 print('Role I - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 393 test_scores['arg_id'][3] * 100.0, test_scores['arg_id'][2], test_scores['arg_id'][1], 394 test_scores['arg_id'][4] * 100.0, test_scores['arg_id'][2], test_scores['arg_id'][0], test_scores['arg_id'][5] * 100.0)) 395 print('Role C - P: {:6.2f} ({:4d}/{:4d}), R: {:6.2f} ({:4d}/{:4d}), F: {:6.2f}'.format( 396 test_scores['arg_cls'][3] * 100.0, test_scores['arg_cls'][2], test_scores['arg_cls'][1], 397 test_scores['arg_cls'][4] * 100.0, test_scores['arg_cls'][2], test_scores['arg_cls'][0], test_scores['arg_cls'][5] * 100.0)) 398 print("---------------------------------------------------------------------") 399 400 if args.write_file: 401 with open(args.write_file, 'w') as fw: 402 json.dump(write_object, fw, indent=4)