text2tree.py (5268B)
1 #!/usr/bin/env python 2 # -*- coding:utf-8 -*- 3 from data_convert.format.target_format import TargetFormat 4 5 6 def get_str_from_tokens(tokens, sentence, separator=' '): 7 start, end_exclude = tokens[0], tokens[-1] + 1 8 return separator.join(sentence[start:end_exclude]) 9 10 11 type_start = '<extra_id_0>' 12 type_end = '<extra_id_1>' 13 role_start = '<extra_id_2>' 14 role_end = '<extra_id_3>' 15 16 17 class Text2Tree(TargetFormat): 18 19 @staticmethod 20 def annotate_predicate_arguments(tokens, predicate_arguments, mark_tree=False, multi_tree=False, zh=False): 21 """ 22 23 :param tokens: 24 US President George W. Bush told Canadian Prime Minister Jean Chretien by telephone Monday that he looked forward 25 to seeing him at the upcoming summit of major industrialized nations and Russia , the White House said Tuesday . 26 :param predicate_arguments: 27 28 :param mark_tree False 29 ( 30 (Meet summit (Entity Russia) (Entity major industrialized nations)) 31 (Phone-Write told (Entity US President George W. Bush) (Entity Canadian Prime Minister Jean Chretien) (Time-Within Monday)) 32 ) 33 :param multi_tree True 34 (Meet summit [Entity Russia] [Entity major industrialized nations]) 35 (Phone-Write told [Entity US President George W. Bush] [Entity Canadian Prime Minister Jean Chretien] [Time-Within Monday]) 36 37 :return: 38 """ 39 token_separator = '' if zh else ' ' 40 41 event_str_rep_list = list() 42 43 for predicate_argument in predicate_arguments: 44 event_type = predicate_argument['type'] 45 46 predicate_text = get_str_from_tokens(predicate_argument['tokens'], tokens, separator=token_separator) 47 48 # prefix_tokens[predicate_argument['tokens'][0]] = ['[ '] 49 # suffix_tokens[predicate_argument['tokens'][-1]] = [' ]'] 50 51 role_str_list = list() 52 for role_name, role_tokens in predicate_argument['arguments']: 53 # if role_name == 'Place' or role_name.startswith('Time'): 54 if role_name == event_type: 55 continue 56 57 role_text = get_str_from_tokens(role_tokens, tokens, separator=token_separator) 58 if mark_tree: 59 role_str = ' '.join([role_start, role_name, role_text, role_end]) 60 else: 61 role_str = ' '.join([type_start, role_name, role_text, type_end]) 62 role_str_list += [role_str] 63 role_str_list_str = ' '.join(role_str_list) 64 event_str_rep = f"{type_start} {event_type} {predicate_text} {role_str_list_str} {type_end}" 65 event_str_rep_list += [event_str_rep] 66 67 source_text = token_separator.join(tokens) 68 target_text = ' '.join(event_str_rep_list) 69 70 if not multi_tree: 71 target_text = f'{type_start} ' + \ 72 ' '.join(event_str_rep_list) + f' {type_end}' 73 74 return source_text, target_text 75 76 @staticmethod 77 def annotate_span(tokens, predicate_arguments, mark_tree=False, zh=False): 78 """ 79 80 :param tokens: 81 US President George W. Bush told Canadian Prime Minister Jean Chretien by telephone Monday that he looked forward 82 to seeing him at the upcoming summit of major industrialized nations and Russia , the White House said Tuesday . 83 :param predicate_arguments: 84 85 :return: 86 mark_tree False 87 ( 88 (Meet summit (Entity Russia) (Entity major industrialized nations)) 89 (Phone-Write told (Entity US President George W. Bush) (Entity Canadian Prime Minister Jean Chretien) (Time-Within Monday)) 90 ) 91 mark_tree True 92 ( 93 (Meet summit [Entity Russia] [Entity major industrialized nations]) 94 (Phone-Write told [Entity US President George W. Bush] [Entity Canadian Prime Minister Jean Chretien] [Time-Within Monday]) 95 ) 96 """ 97 98 token_separator = '' if zh else ' ' 99 100 event_str_rep_list = list() 101 102 for predicate_argument in predicate_arguments: 103 event_type = predicate_argument['type'] 104 105 predicate_text = get_str_from_tokens(predicate_argument['tokens'], tokens, separator=token_separator) 106 107 span_str_list = [' '.join([type_start, event_type, predicate_text, type_end])] 108 109 for role_name, role_tokens in predicate_argument['arguments']: 110 # if role_name == 'Place' or role_name.startswith('Time'): 111 if role_name == event_type: 112 continue 113 114 role_text = get_str_from_tokens(role_tokens, tokens, separator=token_separator) 115 if mark_tree: 116 role_str = ' '.join([role_start, role_name, role_text, role_end]) 117 else: 118 role_str = ' '.join([type_start, role_name, role_text, type_end]) 119 span_str_list += [role_str] 120 event_str_rep_list += [' '.join(span_str_list)] 121 122 source_text = token_separator.join(tokens) 123 target_text = f'{type_start} ' + ' '.join(event_str_rep_list) + f' {type_end}' 124 125 return source_text, target_text 126 127 128 if __name__ == "__main__": 129 pass