convert_text_to_tree.py (6077B)
1 #!/usr/bin/env python 2 # -*- coding:utf-8 -*- 3 import os 4 import json 5 from collections import Counter, defaultdict 6 from data_convert.format.text2tree import Text2Tree 7 from data_convert.task_format.event_extraction import Event, DyIEPP 8 from data_convert.utils import read_file, check_output, data_counter_to_table, get_schema, output_schema 9 from nltk.corpus import stopwords 10 11 english_stopwords = set(stopwords.words('english') + ["'s", "'re", "%"]) 12 13 14 def convert_file_tuple(file_tuple, data_class=Event, target_class=Text2Tree, 15 output_folder='data/text2tree/framenet', 16 ignore_nonevent=False, zh=False, 17 mark_tree=False, type_format='subtype'): 18 counter = defaultdict(Counter) 19 data_counter = defaultdict(Counter) 20 21 event_schema_set = set() 22 23 span_output_folder = output_folder + '_span' 24 25 if not os.path.exists(span_output_folder): 26 os.makedirs(span_output_folder) 27 28 for in_filename, output_filename in file_tuple(output_folder): 29 span_output_filename = output_filename.replace( 30 output_folder, span_output_folder) 31 32 event_output = open(output_filename + '.json', 'w') 33 span_event_output = open(span_output_filename + '.json', 'w') 34 35 for line in read_file(in_filename): 36 document = data_class(json.loads(line.strip())) 37 for sentence in document.generate_sentence(type_format=type_format): 38 39 if ignore_nonevent and len(sentence['events']) == 0: 40 continue 41 42 source, target = target_class.annotate_predicate_arguments( 43 tokens=sentence['tokens'], 44 predicate_arguments=sentence['events'], 45 zh=zh 46 ) 47 48 for event in sentence['events']: 49 event_schema_set = event_schema_set | get_schema(event) 50 sep = '' if zh else ' ' 51 predicate = sep.join([sentence['tokens'][index] 52 for index in event['tokens']]) 53 counter['pred'].update([predicate]) 54 counter['type'].update([event['type']]) 55 data_counter[in_filename].update(['event']) 56 for argument in event['arguments']: 57 data_counter[in_filename].update(['argument']) 58 counter['role'].update([argument[0]]) 59 60 data_counter[in_filename].update(['sentence']) 61 62 event_output.write(json.dumps( 63 {'text': source, 'event': target}, ensure_ascii=False) + '\n') 64 65 span_source, span_target = target_class.annotate_span( 66 tokens=sentence['tokens'], 67 predicate_arguments=sentence['events'], 68 zh=zh, 69 mark_tree=mark_tree 70 ) 71 72 span_event_output.write( 73 json.dumps({'text': span_source, 'event': span_target}, ensure_ascii=False) + '\n') 74 75 event_output.close() 76 span_event_output.close() 77 78 check_output(output_filename) 79 check_output(span_output_filename) 80 print('\n') 81 output_schema(event_schema_set, output_file=os.path.join( 82 output_folder, 'event.schema')) 83 output_schema(event_schema_set, output_file=os.path.join( 84 span_output_folder, 'event.schema')) 85 print('Pred:', len(counter['pred']), counter['pred'].most_common(10)) 86 print('Type:', len(counter['type']), counter['type'].most_common(10)) 87 print('Role:', len(counter['role']), counter['role'].most_common(10)) 88 print(data_counter_to_table(data_counter)) 89 print('\n\n\n') 90 91 92 def convert_ace2005_event(output_folder='data/text2tree/ace2005_event', type_format='subtype', 93 ignore_nonevent=False, mark_tree=False): 94 from data_convert.task_format.event_extraction import ace2005_en_file_tuple 95 convert_file_tuple(file_tuple=ace2005_en_file_tuple, 96 output_folder=output_folder, 97 ignore_nonevent=ignore_nonevent, 98 mark_tree=mark_tree, 99 type_format=type_format, 100 ) 101 102 103 def convert_dyiepp_event(output_folder='data/text2tree/ace2005_event', type_format='subtype', 104 ignore_nonevent=False, mark_tree=False): 105 from data_convert.task_format.event_extraction import DyIEPP_ace2005_file_tuple 106 convert_file_tuple(file_tuple=DyIEPP_ace2005_file_tuple, 107 output_folder=output_folder, 108 ignore_nonevent=ignore_nonevent, 109 mark_tree=mark_tree, 110 type_format=type_format, 111 data_class=DyIEPP 112 ) 113 114 115 def convert_ere_event(output_folder='data/text2tree/ere_event', type_format='subtype', 116 ignore_nonevent=False, mark_tree=False): 117 from data_convert.task_format.event_extraction import ere_en_file_tuple 118 convert_file_tuple(file_tuple=ere_en_file_tuple, 119 output_folder=output_folder, 120 ignore_nonevent=ignore_nonevent, 121 mark_tree=mark_tree, 122 type_format=type_format, 123 ) 124 125 126 if __name__ == "__main__": 127 type_format_name = 'subtype' 128 convert_dyiepp_event("data/text2tree/dyiepp_ace2005_%s" % type_format_name, 129 type_format=type_format_name, 130 ignore_nonevent=False, mark_tree=False, 131 ) 132 convert_ace2005_event("data/text2tree/one_ie_ace2005_%s" % type_format_name, 133 type_format=type_format_name, 134 ignore_nonevent=False, 135 mark_tree=False 136 ) 137 convert_ere_event("data/text2tree/one_ie_ere_en_%s" % type_format_name, 138 type_format=type_format_name, 139 ignore_nonevent=False, 140 mark_tree=False)