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