evaluation.py (8685B)
1 import argparse 2 import json 3 import os 4 import sys 5 import numpy as np 6 from copy import deepcopy 7 from pprint import pprint 8 from extraction.event_schema import EventSchema 9 from extraction.predict_parser.tree_predict_parser import TreePredictParser 10 11 12 def read_file(file_name): 13 return [line.strip() for line in open(file_name).readlines()] 14 15 16 def generate_sentence_dyiepp(filename, type_format='subtype'): 17 for line in open(filename): 18 instance = json.loads(line) 19 sentence = instance['sentence'] 20 sentence_start = instance.get( 21 's_start', instance.get('_sentence_start')) 22 events = instance['event'] 23 24 trigger_list = list() 25 role_list = list() 26 27 for event in events: 28 trigger, event_type = event[0] 29 trigger -= sentence_start 30 31 suptype, subtype = event_type.split('.') 32 33 if type_format == 'subtype': 34 event_type = subtype 35 elif type_format == 'suptype': 36 event_type = suptype 37 else: 38 event_type = suptype + type_format + subtype 39 40 trigger_list += [(event_type, (trigger, trigger))] 41 for start, end, role in event[1:]: 42 start -= sentence_start 43 end -= sentence_start 44 role_list += [(event_type, role, (start, end))] 45 46 yield ' '.join(sentence), trigger_list, role_list 47 48 49 def generate_sentence_oneie(filename, type_format='subtype'): 50 for line in open(filename): 51 instance = json.loads(line) 52 entities = {entity['id']: entity for entity in instance['entity_mentions']} 53 54 trigger_list = list() 55 role_list = list() 56 57 for event in instance['event_mentions']: 58 suptype, subtype = event['event_type'].split(':') 59 60 if type_format == 'subtype': 61 event_type = subtype 62 elif type_format == 'suptype': 63 event_type = suptype 64 else: 65 event_type = suptype + type_format + subtype 66 67 trigger_list += [(event_type, (event['trigger'] 68 ['start'], event['trigger']['end'] - 1))] 69 70 for argument in event['arguments']: 71 argument_entity = entities[argument['entity_id']] 72 role_list += [(event_type, argument['role'], 73 (argument_entity['start'], argument_entity['end'] - 1))] 74 75 yield ' '.join(instance['tokens']), trigger_list, role_list 76 77 78 def match_sublist(the_list, to_match): 79 """ 80 :param the_list: [1, 2, 3, 4, 5, 6, 1, 2, 4, 5] 81 :param to_match: [1, 2] 82 :return: 83 [(0, 1), (6, 7)] 84 """ 85 len_to_match = len(to_match) 86 matched_list = list() 87 for index in range(len(the_list) - len_to_match + 1): 88 if to_match == the_list[index:index + len_to_match]: 89 matched_list += [(index, index + len_to_match - 1)] 90 return matched_list 91 92 93 def record_to_offset(instance): 94 """ 95 Find Role's offset using closest matched with trigger work. 96 :param instance: 97 :return: 98 """ 99 trigger_list = list() 100 role_list = list() 101 102 token_list = instance['text'].split() 103 104 trigger_matched_set = set() 105 for record in instance['pred_record']: 106 event_type = record['type'] 107 trigger = record['trigger'] 108 matched_list = match_sublist(token_list, trigger.split()) 109 110 trigger_offset = None 111 for matched in matched_list: 112 if matched not in trigger_matched_set: 113 trigger_list += [(event_type, matched)] 114 trigger_offset = matched 115 trigger_matched_set.add(matched) 116 break 117 118 # No trigger word, skip the record 119 if trigger_offset is None: 120 break 121 122 for _, role_type, text_str in record['roles']: 123 matched_list = match_sublist(token_list, text_str.split()) 124 if len(matched_list) == 1: 125 role_list += [(event_type, role_type, matched_list[0])] 126 elif len(matched_list) == 0: 127 sys.stderr.write("[Cannot reconstruct]: %s %s\n" % 128 (text_str, token_list)) 129 else: 130 abs_distances = [abs(match[0] - trigger_offset[0]) 131 for match in matched_list] 132 closest_index = np.argmin(abs_distances) 133 role_list += [(event_type, role_type, 134 matched_list[closest_index])] 135 136 return instance['text'], trigger_list, role_list 137 138 139 class Metric: 140 def __init__(self): 141 self.tp = 0. 142 self.gold_num = 0. 143 self.pred_num = 0. 144 145 @staticmethod 146 def safe_div(a, b): 147 if b == 0.: 148 return 0. 149 else: 150 return a / b 151 152 def compute_f1(self, prefix=''): 153 tp = self.tp 154 pred_num = self.pred_num 155 gold_num = self.gold_num 156 p, r = self.safe_div(tp, pred_num), self.safe_div(tp, gold_num) 157 return {prefix + 'tp': tp, 158 prefix + 'gold': gold_num, 159 prefix + 'pred': pred_num, 160 prefix + 'P': p * 100, 161 prefix + 'R': r * 100, 162 prefix + 'F1': self.safe_div(2 * p * r, p + r) * 100 163 } 164 165 def count_instance(self, gold_list, pred_list, verbose=False): 166 if verbose: 167 print("Gold:", gold_list) 168 print("Pred:", pred_list) 169 self.gold_num += len(gold_list) 170 self.pred_num += len(pred_list) 171 172 dup_gold_list = deepcopy(gold_list) 173 for pred in pred_list: 174 if pred in dup_gold_list: 175 self.tp += 1 176 dup_gold_list.remove(pred) 177 178 179 def main(): 180 parser = argparse.ArgumentParser() 181 parser.add_argument('-g', dest='gold_folder') 182 parser.add_argument('-r', dest='offset_folder') 183 parser.add_argument('-p', dest='pred_folder') 184 parser.add_argument('-f', dest='format', default="dyiepp") 185 parser.add_argument('-v', dest='verbose', action='store_true') 186 options = parser.parse_args() 187 188 if options.format == 'dyiepp': 189 data_dict = { 190 'valid': ['eval_preds_seq2seq.txt', 'val.json', "dev_convert.json"], 191 'test': ['test_preds_seq2seq.txt', 'test.json', "test_convert.json"], 192 } 193 generate_sentence = generate_sentence_dyiepp 194 elif options.format == 'oneie': 195 data_dict = { 196 'valid': ['eval_preds_seq2seq.txt', 'val.json', "dev.oneie.json"], 197 'test': ['test_preds_seq2seq.txt', 'test.json', "test.oneie.json"], 198 } 199 generate_sentence = generate_sentence_oneie 200 else: 201 raise NotImplementedError('%s not support' % options.format) 202 203 label_schema = EventSchema.read_from_file( 204 filename=os.path.join(options.gold_folder, 'event.schema') 205 ) 206 207 pred_folder = options.pred_folder 208 gold_folder = options.gold_folder 209 offset_folder = options.offset_folder 210 211 for data_key, (generation, text_file, offset_file) in data_dict.items(): 212 213 trigger_metric = Metric() 214 argument_metric = Metric() 215 216 # Reconstruct the offset of predicted event records. 217 text_filename = os.path.join(gold_folder, text_file) 218 pred_filename = os.path.join(pred_folder, generation) 219 220 print("pred:", pred_filename) 221 pred_reader = TreePredictParser(label_constraint=label_schema) 222 event_list, _ = pred_reader.decode( 223 gold_list=[], 224 pred_list=read_file(pred_filename), 225 text_list=[json.loads(line)['text'] 226 for line in read_file(text_filename)], 227 ) 228 pred_list = [ 229 record_to_offset(decoding_instance) for decoding_instance in event_list 230 ] 231 232 # Read gold event annotation with offsets. 233 gold_filename = os.path.join(offset_folder, offset_file) 234 print("gold:", gold_filename) 235 gold_list = [event for event in generate_sentence(gold_filename)] 236 237 for pred, gold in zip(pred_list, gold_list): 238 assert pred[0] == gold[0] 239 trigger_metric.count_instance( 240 gold_list=gold[1], 241 pred_list=pred[1], 242 verbose=options.verbose, 243 ) 244 argument_metric.count_instance( 245 gold_list=gold[2], 246 pred_list=pred[2], 247 verbose=options.verbose, 248 ) 249 250 trigger_result = trigger_metric.compute_f1(prefix=data_key + '-trig-') 251 role_result = argument_metric.compute_f1(prefix=data_key + '-role-') 252 253 pprint(trigger_result) 254 pprint(role_result) 255 256 257 if __name__ == "__main__": 258 main()