tree_predict_parser.py (9063B)
1 from collections import Counter 2 from typing import Tuple, List, Dict 3 4 from nltk.tree import ParentedTree 5 import re 6 7 from data_convert.ucdp_fields import UCDP_FIELDS 8 from extraction.predict_parser.predict_parser import PredictParser 9 10 type_start = '<extra_id_0>' 11 type_end = '<extra_id_1>' 12 role_start = '<extra_id_2>' 13 role_end = '<extra_id_3>' 14 left_bracket = '【' 15 right_bracket = '】' 16 brackets = left_bracket + right_bracket 17 18 split_bracket = re.compile(r"<extra_id_\d>") 19 20 21 def add_space(text): 22 """ 23 add space between special token 24 :param text: 25 :return: 26 """ 27 new_text_list = list() 28 for item in zip(split_bracket.findall(text), split_bracket.split(text)[1:]): 29 new_text_list += item 30 return ' '.join(new_text_list) 31 32 33 def find_bracket_num(tree_str): 34 """ 35 Count Bracket Number, 0 indicate num_left = num_right 36 :param tree_str: 37 :return: 38 """ 39 count = 0 40 for char in tree_str: 41 if char == left_bracket: 42 count += 1 43 elif char == right_bracket: 44 count -= 1 45 else: 46 pass 47 return count 48 49 50 def check_well_form(tree_str): 51 return find_bracket_num(tree_str) == 0 52 53 54 def clean_text(tree_str): 55 count = 0 56 sum_count = 0 57 58 tree_str_list = tree_str.split() 59 # bracket_num = find_bracket_num(tree_str_list) 60 # bracket_num = find_bracket_num(tree_str_list) 61 62 for index, char in enumerate(tree_str_list): 63 if char == left_bracket: 64 count += 1 65 sum_count += 1 66 elif char == right_bracket: 67 count -= 1 68 sum_count += 1 69 else: 70 pass 71 if count == 0 and sum_count > 0: 72 return ' '.join(tree_str_list[:index + 1]) 73 return ' '.join(tree_str_list) 74 75 76 def add_bracket(tree_str): 77 """ 78 add right bracket to fill ill-formed 79 :param tree_str: 80 :return: 81 """ 82 tree_str_list = tree_str.split() 83 bracket_num = find_bracket_num(tree_str_list) 84 tree_str_list += [right_bracket] * bracket_num 85 return ' '.join(tree_str_list) 86 87 88 def get_tree_str(tree): 89 """ 90 get str from event tree 91 :param tree: 92 :return: 93 """ 94 str_list = list() 95 for element in tree: 96 if isinstance(element, str): 97 str_list += [element] 98 return ' '.join(str_list) 99 100 101 class TreePredictParser(PredictParser): 102 103 def decode(self, gold_list, pred_list, text_list=None, raw_list=None) -> Tuple[List[Dict], Counter]: 104 """ 105 106 :param gold_list: 107 :param pred_list: 108 :param text_list: 109 :param raw_list: 110 :return: 111 dict: 112 pred_event -> [(type1, trigger1), (type2, trigger2), ...] 113 gold_event -> [(type1, trigger1), (type2, trigger2), ...] 114 pred_role -> [(type1, role1, argument1), (type2, role2, argument2), ...] 115 gold_role -> [(type1, role1, argument1), (type2, role2, argument2), ...] 116 Counter: 117 """ 118 counter = Counter() 119 well_formed_list = [] 120 121 def convert_bracket(_text): 122 _text = add_space(_text) 123 for start in [role_start, type_start]: 124 _text = _text.replace(start, left_bracket) 125 for end in [role_end, type_end]: 126 _text = _text.replace(end, right_bracket) 127 return _text 128 129 if gold_list is None or len(gold_list) == 0: 130 gold_list = ["%s%s" % (type_start, type_end)] * len(pred_list) 131 132 if text_list is None: 133 text_list = [None] * len(pred_list) 134 135 if raw_list is None: 136 raw_list = [None] * len(pred_list) 137 138 for gold, pred, text, raw_data in zip(gold_list, pred_list, text_list, raw_list): 139 gold = convert_bracket(gold) 140 pred = convert_bracket(pred) 141 142 pred = clean_text(pred) 143 144 try: 145 gold_tree = ParentedTree.fromstring(gold, brackets=brackets) 146 except ValueError: 147 print(gold) 148 print(add_bracket(gold)) 149 gold_tree = ParentedTree.fromstring(add_bracket(gold), brackets=brackets) 150 counter.update(['gold_tree add_bracket']) 151 152 instance = {'gold': gold, 153 'pred': pred, 154 'gold_tree': gold_tree, 155 'text': text, 156 'raw_data': raw_data 157 } 158 159 counter.update(['gold_tree' for _ in gold_tree]) 160 161 try: 162 pred_tree = ParentedTree.fromstring(pred, brackets=brackets) 163 counter.update(['pred_tree' for _ in pred_tree]) 164 165 instance['pred_tree'] = pred_tree 166 counter.update(['well-formed']) 167 168 except ValueError: 169 170 counter.update(['ill-formed']) 171 instance['pred_tree'] = ParentedTree.fromstring(left_bracket + right_bracket, brackets=brackets) 172 173 instance['pred_event'], instance['pred_role'], instance['pred_record'] = self.get_event_list( 174 tree=instance["pred_tree"], 175 text=instance['text'] 176 ) 177 instance['gold_event'], instance['gold_role'], instance['gold_record'] = self.get_event_list( 178 tree=instance["gold_tree"], 179 text=instance['text'] 180 ) 181 182 self.count_multi_event_role_in_instance(instance=instance, counter=counter) 183 184 well_formed_list += [instance] 185 186 return well_formed_list, counter 187 188 def get_event_list(self, tree, text=None): 189 190 event_list = list() 191 role_list = list() 192 record_list = list() 193 194 for event_tree in tree: 195 196 if isinstance(event_tree, str): 197 continue 198 if len(event_tree) == 0: 199 continue 200 201 event_type = event_tree.label() 202 event_trigger = get_tree_str(event_tree) 203 204 # Invalid Event Type 205 if self.predicate_set and event_type not in self.predicate_set: 206 continue 207 # Invalid Text Span 208 if text is not None and event_trigger not in text: 209 continue 210 211 record = {'roles': list(), 212 'type': event_type, 213 'trigger': event_trigger} 214 for role_tree in event_tree[1:]: 215 role_text = get_tree_str(role_tree) 216 217 if isinstance(role_tree, str) or len(role_tree) < 1: 218 continue 219 220 # Invalid Role Type 221 if self.role_set and role_tree.label() not in self.role_set: 222 continue 223 # Invalid Text Span 224 if text is not None and role_text not in text: 225 continue 226 227 role_list += [(event_type, role_tree.label(), role_text)] 228 record['roles'] += [(event_type, role_tree.label(), role_text)] 229 230 event_list += [(event_type, event_trigger)] 231 record_list += [record] 232 233 return event_list, role_list, record_list 234 235 236 class ListPredictParser(TreePredictParser): 237 def decode(self, gold_list, pred_list, text_list=None, raw_list=None) -> Tuple[List[Dict], Counter]: 238 counter = Counter() 239 well_formed_list = [] 240 241 for gold, pred in zip(gold_list, pred_list): 242 gold_tree, gold_roles = self.tree_of_text(gold) 243 pred_tree, pred_roles = self.tree_of_text(pred) 244 245 counter["gold_tree"] += 1 246 counter["pred_tree"] += 1 247 if len(pred_roles) == len(gold_roles) == len(UCDP_FIELDS): 248 counter["well-formed"] += 1 249 else: 250 counter["ill-formed"] += 1 251 252 well_formed_list.append({ 253 "text": None, 254 "raw_data": None, 255 "gold": gold, 256 "pred": pred, 257 "gold_tree": gold_tree, 258 "pred_tree": pred_tree, 259 "gold_event": [("UCDP", "")], 260 "pred_event": [("UCDP", "")], 261 "gold_role": gold_roles, 262 "pred_role": pred_roles, 263 "gold_regord": [{"roles": gold_roles.copy(), "type": "UCDP", "trigger": ""}], 264 "pred_regord": [{"roles": pred_roles.copy(), "type": "UCDP", "trigger": ""}], 265 }) 266 267 return well_formed_list, counter 268 269 def tree_of_text(self, text): 270 tree = ParentedTree("UCDP", []) 271 roles = [] 272 for pair in text.split("<extra_id_1> <extra_id_0>"): 273 pair = pair.replace("<extra_id_0>", "").replace("<extra_id_1>", "").strip() 274 field, role = None, None 275 for crole, cfield, _ in sorted(UCDP_FIELDS, key=lambda field: len(field[1]), reverse=True): 276 if pair.startswith(cfield): 277 field, role = cfield, crole 278 break 279 if field is not None: 280 argument = pair[len(field):].strip() 281 roles.append(("UCDP", role, argument)) 282 tree.append(ParentedTree(role, [argument])) 283 return tree, roles