extract_constraint.py (16775B)
1 #!/usr/bin/env python 2 # -*- coding:utf-8 -*- 3 from typing import List, Dict 4 5 from data_convert.format.text2tree import type_start, type_end 6 from extraction.label_tree import get_label_name_tree 7 8 import os 9 10 debug = True if 'DEBUG' in os.environ else False 11 debug_step = True if 'DEBUG_STEP' in os.environ else False 12 13 14 def match_sublist(the_list, to_match): 15 """ 16 17 :param the_list: [1, 2, 3, 4, 5, 6, 1, 2, 4, 5] 18 :param to_match: [1, 2] 19 :return: 20 [(0, 1), (6, 7)] 21 """ 22 len_to_match = len(to_match) 23 matched_list = list() 24 for index in range(len(the_list) - len_to_match + 1): 25 if to_match == the_list[index:index + len_to_match]: 26 matched_list += [(index, index + len_to_match - 1)] 27 return matched_list 28 29 30 def find_bracket_position(generated_text, _type_start, _type_end): 31 bracket_position = {_type_start: list(), _type_end: list()} 32 for index, char in enumerate(generated_text): 33 if char in bracket_position: 34 bracket_position[char] += [index] 35 return bracket_position 36 37 38 def generated_search_src_sequence(generated, src_sequence, end_sequence_search_tokens=None): 39 print(generated, src_sequence) if debug else None 40 41 if len(generated) == 0: 42 # It has not been generated yet. All SRC are valid. 43 return src_sequence 44 45 matched_tuples = match_sublist(the_list=src_sequence, to_match=generated) 46 47 valid_token = list() 48 for _, end in matched_tuples: 49 next_index = end + 1 50 if next_index < len(src_sequence): 51 valid_token += [src_sequence[next_index]] 52 53 if end_sequence_search_tokens: 54 valid_token += end_sequence_search_tokens 55 56 return valid_token 57 58 59 def get_constraint_decoder(tokenizer, type_schema, decoding_schema, source_prefix=None): 60 if decoding_schema == 'tree': 61 return TreeConstraintDecoder(tokenizer=tokenizer, type_schema=type_schema, source_prefix=source_prefix) 62 elif decoding_schema == 'list': 63 return ListConstraintDecoder(tokenizer=tokenizer, type_schema=type_schema, source_prefix=source_prefix) 64 elif decoding_schema == 'treespan': 65 return SpanConstraintDecoder(tokenizer=tokenizer, type_schema=type_schema, source_prefix=source_prefix) 66 else: 67 raise NotImplementedError( 68 'Type Schema %s, Decoding Schema %s do not map to constraint decoder.' % ( 69 decoding_schema, decoding_schema) 70 ) 71 72 73 class ConstraintDecoder: 74 def __init__(self, tokenizer, source_prefix): 75 self.tokenizer = tokenizer 76 self.source_prefix = source_prefix 77 self.source_prefix_tokenized = tokenizer.encode(source_prefix, 78 add_special_tokens=False) if source_prefix else [] 79 80 def get_state_valid_tokens(self, src_sentence: List[str], tgt_generated: List[str]) -> List[str]: 81 pass 82 83 def constraint_decoding(self, src_sentence, tgt_generated): 84 if self.source_prefix_tokenized: 85 # Remove Source Prefix for Generation 86 src_sentence = src_sentence[len(self.source_prefix_tokenized):] 87 88 if debug: 89 print("Src:", self.tokenizer.convert_ids_to_tokens(src_sentence)) 90 print("Tgt:", self.tokenizer.convert_ids_to_tokens(tgt_generated)) 91 92 valid_token_ids = self.get_state_valid_tokens( 93 src_sentence.tolist(), 94 tgt_generated.tolist() 95 ) 96 97 if debug: 98 print('========================================') 99 print('valid tokens:', self.tokenizer.convert_ids_to_tokens( 100 valid_token_ids), valid_token_ids) 101 if debug_step: 102 input() 103 104 # return self.tokenizer.convert_tokens_to_ids(valid_tokens) 105 return valid_token_ids 106 107 108 class TreeConstraintDecoder(ConstraintDecoder): 109 def __init__(self, tokenizer, type_schema, *args, **kwargs): 110 super().__init__(tokenizer, *args, **kwargs) 111 self.tree_end = '<tree-end>' 112 self.type_tree = get_label_name_tree( 113 type_schema.type_list, self.tokenizer, end_symbol=self.tree_end) 114 self.role_tree = get_label_name_tree( 115 type_schema.role_list, self.tokenizer, end_symbol=self.tree_end) 116 self.type_start = self.tokenizer.convert_tokens_to_ids([type_start])[0] 117 self.type_end = self.tokenizer.convert_tokens_to_ids([type_end])[0] 118 119 def check_state(self, tgt_generated): 120 if tgt_generated[-1] == self.tokenizer.pad_token_id: 121 return 'start', -1 122 123 special_token_set = {self.type_start, self.type_end} 124 special_index_token = list( 125 filter(lambda x: x[1] in special_token_set, list(enumerate(tgt_generated)))) 126 127 last_special_index, last_special_token = special_index_token[-1] 128 129 if len(special_index_token) == 1: 130 if last_special_token != self.type_start: 131 return 'error', 0 132 133 bracket_position = find_bracket_position( 134 tgt_generated, _type_start=self.type_start, _type_end=self.type_end) 135 start_number, end_number = len(bracket_position[self.type_start]), len( 136 bracket_position[self.type_end]) 137 138 if start_number == end_number: 139 return 'end_generate', -1 140 if start_number == end_number + 1: 141 state = 'start_first_generation' 142 elif start_number == end_number + 2: 143 state = 'generate_trigger' 144 elif start_number == end_number + 3: 145 state = 'generate_role' 146 else: 147 state = 'error' 148 return state, last_special_index 149 150 def search_prefix_tree_and_sequence(self, generated: List[str], prefix_tree: Dict, src_sentence: List[str], 151 end_sequence_search_tokens: List[str] = None): 152 """ 153 Generate Type Name + Text Span 154 :param generated: 155 :param prefix_tree: 156 :param src_sentence: 157 :param end_sequence_search_tokens: 158 :return: 159 """ 160 tree = prefix_tree 161 for index, token in enumerate(generated): 162 tree = tree[token] 163 is_tree_end = len(tree) == 1 and self.tree_end in tree 164 165 if is_tree_end: 166 valid_token = generated_search_src_sequence( 167 generated=generated[index + 1:], 168 src_sequence=src_sentence, 169 end_sequence_search_tokens=end_sequence_search_tokens, 170 ) 171 return valid_token 172 173 if self.tree_end in tree: 174 try: 175 valid_token = generated_search_src_sequence( 176 generated=generated[index + 1:], 177 src_sequence=src_sentence, 178 end_sequence_search_tokens=end_sequence_search_tokens, 179 ) 180 return valid_token 181 except IndexError: 182 # Still search tree 183 continue 184 185 valid_token = list(tree.keys()) 186 return valid_token 187 188 def get_state_valid_tokens(self, src_sentence, tgt_generated): 189 """ 190 191 :param src_sentence: 192 :param tgt_generated: 193 :return: 194 List[str], valid token list 195 """ 196 if self.tokenizer.eos_token_id in src_sentence: 197 src_sentence = src_sentence[:src_sentence.index( 198 self.tokenizer.eos_token_id)] 199 200 state, index = self.check_state(tgt_generated) 201 202 print("State: %s" % state) if debug else None 203 204 if state == 'error': 205 print("Error:") 206 print("Src:", src_sentence) 207 print("Tgt:", tgt_generated) 208 valid_tokens = [self.tokenizer.eos_token_id] 209 210 elif state == 'start': 211 valid_tokens = [self.type_start] 212 213 elif state == 'start_first_generation': 214 valid_tokens = [self.type_start, self.type_end] 215 216 elif state == 'generate_trigger': 217 218 if tgt_generated[-1] == self.type_start: 219 # Start Event Label 220 return list(self.type_tree.keys()) 221 222 elif tgt_generated[-1] == self.type_end: 223 # EVENT_TYPE_LEFT: Start a new role 224 # EVENT_TYPE_RIGHT: End this event 225 return [self.type_start, self.type_end] 226 else: 227 valid_tokens = self.search_prefix_tree_and_sequence( 228 generated=tgt_generated[index + 1:], 229 prefix_tree=self.type_tree, 230 src_sentence=src_sentence, 231 end_sequence_search_tokens=[self.type_start, self.type_end] 232 ) 233 234 elif state == 'generate_role': 235 236 if tgt_generated[-1] == self.type_start: 237 # Start Role Label 238 return list(self.role_tree.keys()) 239 240 generated = tgt_generated[index + 1:] 241 valid_tokens = self.search_prefix_tree_and_sequence( 242 generated=generated, 243 prefix_tree=self.role_tree, 244 src_sentence=src_sentence, 245 end_sequence_search_tokens=[self.type_end] 246 ) 247 248 elif state == 'end_generate': 249 valid_tokens = [self.tokenizer.eos_token_id] 250 251 else: 252 raise NotImplementedError( 253 'State `%s` for %s is not implemented.' % (state, self.__class__)) 254 255 print("Valid: %s" % self.tokenizer.convert_ids_to_tokens( 256 valid_tokens)) if debug else None 257 return valid_tokens 258 259 260 class ListConstraintDecoder(TreeConstraintDecoder): 261 def check_state(self, tgt_generated): 262 if tgt_generated[-1] == self.tokenizer.pad_token_id: 263 return 'start', -1 264 265 special_token_set = {self.type_start, self.type_end} 266 special_index_token = list( 267 filter(lambda x: x[1] in special_token_set, list(enumerate(tgt_generated)))) 268 269 last_special_index, last_special_token = special_index_token[-1] 270 271 if len(special_index_token) == 1: 272 if last_special_token != self.type_start: 273 return 'error', 0 274 275 bracket_position = find_bracket_position( 276 tgt_generated, _type_start=self.type_start, _type_end=self.type_end) 277 start_number, end_number = len(bracket_position[self.type_start]), len( 278 bracket_position[self.type_end]) 279 280 if start_number == end_number: 281 return 'end_generate', -1 282 if start_number == end_number + 1: 283 state = 'start_first_generation' 284 elif start_number == end_number + 2: 285 state = 'generate_role' 286 else: 287 state = 'error' 288 return state, last_special_index 289 290 def search_prefix_tree_and_sequence(self, generated: List[str], prefix_tree: Dict, src_sentence: List[str], 291 end_sequence_search_tokens: List[str] = None): 292 """ 293 Generate Type Name + Text Span 294 :param generated: 295 :param prefix_tree: 296 :param src_sentence: 297 :param end_sequence_search_tokens: 298 :return: 299 """ 300 tree = prefix_tree 301 for index, token in enumerate(generated): 302 tree = tree[token] 303 is_tree_end = len(tree) == 1 and self.tree_end in tree 304 305 if is_tree_end: 306 valid_token = generated_search_src_sequence( 307 generated=generated[index + 1:], 308 src_sequence=src_sentence, 309 end_sequence_search_tokens=end_sequence_search_tokens, 310 ) 311 return valid_token 312 313 if self.tree_end in tree: 314 # Generate freely 315 return list(self.tokenizer.get_vocab().values()) 316 317 valid_token = list(tree.keys()) 318 return valid_token 319 320 321 class SpanConstraintDecoder(ConstraintDecoder): 322 def __init__(self, tokenizer, type_schema, *args, **kwargs): 323 super().__init__(tokenizer, *args, **kwargs) 324 self.tree_end = '<tree-end>' 325 self.type_tree = get_label_name_tree(type_schema.type_list + type_schema.role_list, 326 tokenizer=self.tokenizer, 327 end_symbol=self.tree_end) 328 self.type_start = self.tokenizer.convert_tokens_to_ids([type_start])[0] 329 self.type_end = self.tokenizer.convert_tokens_to_ids([type_end])[0] 330 331 def check_state(self, tgt_generated): 332 if tgt_generated[-1] == self.tokenizer.pad_token_id: 333 return 'start', -1 334 335 special_token_set = {self.type_start, self.type_end} 336 special_index_token = list( 337 filter(lambda x: x[1] in special_token_set, list(enumerate(tgt_generated)))) 338 339 last_special_index, last_special_token = special_index_token[-1] 340 341 if len(special_index_token) == 1: 342 if last_special_token != self.type_start: 343 return 'error', 0 344 345 bracket_position = find_bracket_position( 346 tgt_generated, _type_start=self.type_start, _type_end=self.type_end) 347 start_number, end_number = len(bracket_position[self.type_start]), len( 348 bracket_position[self.type_end]) 349 350 if start_number == end_number: 351 return 'end_generate', -1 352 if start_number == end_number + 1: 353 state = 'start_first_generation' 354 elif start_number == end_number + 2: 355 state = 'generate_span' 356 else: 357 state = 'error' 358 return state, last_special_index 359 360 def search_prefix_tree_and_sequence(self, generated: List[str], prefix_tree: Dict, src_sentence: List[str], 361 end_sequence_search_tokens: List[str] = None): 362 """ 363 Generate Type Name + Text Span 364 :param generated: 365 :param prefix_tree: 366 :param src_sentence: 367 :param end_sequence_search_tokens: 368 :return: 369 """ 370 tree = prefix_tree 371 for index, token in enumerate(generated): 372 tree = tree[token] 373 is_tree_end = len(tree) == 1 and self.tree_end in tree 374 375 if is_tree_end: 376 valid_token = generated_search_src_sequence( 377 generated=generated[index + 1:], 378 src_sequence=src_sentence, 379 end_sequence_search_tokens=end_sequence_search_tokens, 380 ) 381 return valid_token 382 383 if self.tree_end in tree: 384 try: 385 valid_token = generated_search_src_sequence( 386 generated=generated[index + 1:], 387 src_sequence=src_sentence, 388 end_sequence_search_tokens=end_sequence_search_tokens, 389 ) 390 return valid_token 391 except IndexError: 392 # Still search tree 393 continue 394 395 valid_token = list(tree.keys()) 396 return valid_token 397 398 def get_state_valid_tokens(self, src_sentence, tgt_generated): 399 """ 400 401 :param src_sentence: 402 :param tgt_generated: 403 :return: 404 List[str], valid token list 405 """ 406 if self.tokenizer.eos_token_id in src_sentence: 407 src_sentence = src_sentence[:src_sentence.index( 408 self.tokenizer.eos_token_id)] 409 410 state, index = self.check_state(tgt_generated) 411 412 print("State: %s" % state) if debug else None 413 414 if state == 'error': 415 print("Error:") 416 print("Src:", src_sentence) 417 print("Tgt:", tgt_generated) 418 valid_tokens = [self.tokenizer.eos_token_id] 419 420 elif state == 'start': 421 valid_tokens = [self.type_start] 422 423 elif state == 'start_first_generation': 424 valid_tokens = [self.type_start, self.type_end] 425 426 elif state == 'generate_span': 427 428 if tgt_generated[-1] == self.type_start: 429 # Start Event Label 430 return list(self.type_tree.keys()) 431 432 elif tgt_generated[-1] == self.type_end: 433 raise RuntimeError('Invalid %s in %s' % 434 (self.type_end, tgt_generated)) 435 436 else: 437 valid_tokens = self.search_prefix_tree_and_sequence( 438 generated=tgt_generated[index + 1:], 439 prefix_tree=self.type_tree, 440 src_sentence=src_sentence, 441 end_sequence_search_tokens=[self.type_end] 442 ) 443 444 elif state == 'end_generate': 445 valid_tokens = [self.tokenizer.eos_token_id] 446 447 else: 448 raise NotImplementedError( 449 'State `%s` for %s is not implemented.' % (state, self.__class__)) 450 451 print("Valid: %s" % valid_tokens) if debug else None 452 return valid_tokens