utils.py (1903B)
1 #!/usr/bin/env python 2 # -*- coding:utf-8 -*- 3 import json 4 from collections import defaultdict 5 6 from tabulate import tabulate 7 8 9 def read_file(filename): 10 from tqdm import tqdm 11 num_lines = sum(1 for _ in open(filename, 'r')) 12 with open(filename, 'r') as f: 13 for line in tqdm(f, total=num_lines): 14 yield line 15 16 17 def check_output(filename, line_num=2): 18 import os 19 os.system('tail -n %s %s*' % (line_num, filename)) 20 21 22 def data_counter_to_table(data_counter): 23 table = list() 24 for filename, file_counter in data_counter.items(): 25 table += [[filename, file_counter['sentence'], 26 file_counter['event'], file_counter['argument']]] 27 return tabulate(table, headers=['file', '#sent', '#event', '#arg']) 28 29 30 def get_schema(event): 31 event_type = event['type'] 32 if len(event['arguments']) == 0: 33 return {(event_type, None)} 34 return set([(event_type, argument[0]) for argument in event['arguments']]) 35 36 37 def output_schema(event_schema_set, output_file): 38 event_type_list = list(set([schema[0] for schema in event_schema_set])) 39 argument_role_list = list(set([schema[1] for schema in event_schema_set])) 40 41 if None in argument_role_list: 42 # Same Event only Type without argument 43 argument_role_list.remove(None) 44 45 event_type_set_dict = defaultdict(set) 46 47 for event_type, arg_role in event_schema_set: 48 if arg_role is None: 49 continue 50 event_type_set_dict[event_type].add(arg_role) 51 52 event_type_list_dict = defaultdict(list) 53 54 for event_type in event_type_set_dict: 55 event_type_list_dict[event_type] = list( 56 event_type_set_dict[event_type]) 57 58 with open(output_file, 'w') as output: 59 output.write(json.dumps(event_type_list) + '\n') 60 output.write(json.dumps(argument_role_list) + '\n') 61 output.write(json.dumps(event_type_list_dict) + '\n')