get_doc_statistics.py (2716B)
1 import argparse 2 import json 3 from collections import Counter, defaultdict, OrderedDict 4 import ipdb 5 import os 6 import random 7 8 9 parser = argparse.ArgumentParser() 10 parser.add_argument("-i", "--input_path", type=str, required=True) 11 args = parser.parse_args() 12 13 data = [] 14 for line in open(args.input_path, 'r', encoding='utf-8'): 15 data.append(json.loads(line)) 16 17 class Info: 18 def __init__(self): 19 self.sentence_num = 0 20 self.entity = Counter() 21 self.entity_num = 0 22 self.event = Counter() 23 self.event_num = 0 24 self.role = Counter() 25 self.role_num = 0 26 27 def update(self, instance): 28 self.sentence_num += 1 29 id2entity = dict() 30 for entity in instance['entity_mentions']: 31 id2entity[entity['id']] = entity 32 self.entity[entity['entity_type']] += 1 33 self.entity_num += 1 34 35 for event in instance['event_mentions']: 36 self.event[event['event_type']] += 1 37 self.event_num += 1 38 for argument in event['arguments']: 39 self.role[argument['role']] += 1 40 self.role_num += 1 41 42 def __add__(self, obj): 43 self.sentence_num += obj.sentence_num 44 self.entity = self.entity + obj.entity 45 self.entity_num += obj.entity_num 46 self.event = self.event + obj.event 47 self.event_num += obj.event_num 48 self.role = self.role + obj.role 49 self.role_num += obj.role_num 50 return self 51 52 def get_statistics(list_of_key, d_stat): 53 total_sent = 0 54 total_event = Counter() 55 total_event_num = 0 56 total_role = Counter() 57 total_role_num = 0 58 total_entity_num = 0 59 for l in list_of_key: 60 assert l in d_stat.keys() 61 total_sent += d_stat[l].sentence_num 62 total_event += d_stat[l].event 63 total_role += d_stat[l].role 64 total_event_num += d_stat[l].event_num 65 total_role_num += d_stat[l].role_num 66 total_entity_num += d_stat[l].entity_num 67 68 return len(list_of_key), total_sent, total_event_num, len(total_event), total_role_num, len(total_role) 69 70 def aggregate(list_of_key, d_stat): 71 aggre = Info() 72 for l in list_of_key: 73 assert l in d_stat.keys() 74 aggre = aggre+d_stat[l] 75 return aggre 76 77 # get info 78 doc_statistics = defaultdict(Info) 79 for instance in data: 80 info = doc_statistics[instance['doc_id']] 81 info.update(instance) 82 83 def export_doc_list(lists, filename): 84 with open(filename, 'w') as f: 85 for l in lists: 86 f.write(l + '\n') 87 88 def read_doc_list(filename): 89 lists = [] 90 with open(filename, 'r') as f: 91 for l in f.readlines(): 92 lists.append(l.strip('\n')) 93 return lists