process_ace05e.py (9415B)
1 import json 2 from argparse import ArgumentParser 3 from transformers import BertTokenizer, RobertaTokenizer, AutoTokenizer 4 5 6 def map_index(pieces): 7 idxs = [] 8 for i, piece in enumerate(pieces): 9 if i == 0: 10 idxs.append([0, len(piece)]) 11 else: 12 _, last = idxs[-1] 13 idxs.append([last, last + len(piece)]) 14 return idxs 15 16 def map_decode_back_pieces(encoded_input, ori_tokens, tokenizer): 17 decoded = [tokenizer.decode(x) for x in encoded_input['input_ids']] 18 pieces = [] 19 ori_cnt = 0 20 current = [] 21 for d in decoded: 22 if d == '': 23 continue 24 current.append(d) 25 if ''.join(current) == ori_tokens[ori_cnt]: 26 pieces.append(current) 27 current = [] 28 ori_cnt += 1 29 assert len(pieces) == len(ori_tokens) 30 return pieces 31 32 def convert(input_file, output_file, tokenizer, window_size_=3): 33 with open(input_file, 'r', encoding='utf-8') as r, \ 34 open(output_file, 'w', encoding='utf-8') as w: 35 for line in r: 36 doc = json.loads(line) 37 doc_id = doc['doc_key'] 38 sentences = doc['sentences'] 39 sent_num = len(sentences) 40 total_tokens = sum([len(sent) for sent in sentences]) 41 coref_entities = doc['clusters'] 42 coref_events = doc['event_clusters'] 43 # upper bound on token index for checking index in range 44 sent_starts = doc['_sentence_start'] + [total_tokens] 45 entities = doc.get('ner', [[] for _ in range(sent_num)]) 46 relations = doc.get('relations', [[] for _ in range(sent_num)]) 47 events = doc.get('events', [[] for _ in range(sent_num)]) 48 49 if window_size_ > sent_num: 50 window_size = sent_num 51 else: 52 window_size = window_size_ 53 54 offset = 0 55 for i in range(sent_num - window_size + 1): 56 wnd_sent_starts = sent_starts[i:i+window_size+1] 57 wnd_start, wnd_end = wnd_sent_starts[0], wnd_sent_starts[-1] 58 59 def slice_fn(lst, ind, wnd): 60 return [item for j in range(wnd) for item in lst[ind+j]] 61 wnd_tokens, wnd_entities, wnd_relations, wnd_events = [slice_fn( 62 lst, i, window_size) for lst in [sentences, entities, relations, events]] 63 64 wnd_id = '{}-{}'.format(doc_id, i) 65 pieces = [tokenizer.tokenize(t) for t in wnd_tokens] 66 word_lens = [len(p) for p in pieces] 67 68 wnd_entities_ = [] 69 wnd_entity_map = {} 70 for j, (start, end, entity_type) in enumerate(wnd_entities): 71 start, end = start - offset, end - offset + 1 72 entity_id = '{}-E{}'.format(wnd_id, j) 73 entity = { 74 'id': entity_id, 75 'start': start, 'end': end, 76 'entity_type': entity_type, 77 # Mention types are not included in DyGIE++'s format 78 'mention_type': 'UNK', 79 'text': ' '.join(wnd_tokens[start:end])} 80 wnd_entities_.append(entity) 81 wnd_entity_map[(start, end)] = entity 82 83 wnd_relations_ = [] 84 for j, (start1, end1, start2, end2, rel_type) in enumerate(wnd_relations): 85 start1, end1 = start1 - offset, end1 - offset + 1 86 start2, end2 = start2 - offset, end2 - offset + 1 87 arg1 = wnd_entity_map[(start1, end1)] 88 arg2 = wnd_entity_map[(start2, end2)] 89 relation_id = '{}-R{}'.format(wnd_id, j) 90 rel_type = rel_type.split('.')[0] 91 relation = { 92 'relation_type': rel_type, 93 'id': relation_id, 94 'arguments': [ 95 { 96 'entity_id': arg1['id'], 97 'text': arg1['text'], 98 'role': 'Arg-1' 99 }, 100 { 101 'entity_id': arg2['id'], 102 'text': arg2['text'], 103 'role': 'Arg-2' 104 }, 105 ] 106 } 107 wnd_relations_.append(relation) 108 109 # parse coref entities 110 # for each entity mention in a coref, only look up the obj in the dict if they are in the window 111 wnd_coref_ents = [[wnd_entity_map[(ent[0]-offset, ent[1]-offset+1)] for ent in coref if wnd_start <= ent[0] 112 and ent[1] < wnd_end and (ent[0]-offset, ent[1]-offset+1) in wnd_entity_map] for coref in coref_entities] 113 wnd_coref_ents_ = [] 114 for j, ent_list in enumerate(wnd_coref_ents): 115 if len(ent_list) > 1: 116 wnd_coref_ents_.append({ 117 'id': '{}-CE{}'.format(wnd_id, j), 118 'entities': ent_list 119 }) 120 121 wnd_events_ = [] 122 wnd_event_map = {} 123 for j, event in enumerate(wnd_events): 124 event_id = '{}-EV{}'.format(wnd_id, j) 125 if len(event[0]) == 3: 126 trigger_start, trigger_end, event_type = event[0] 127 elif len(event[0]) == 2: 128 trigger_start, event_type = event[0] 129 trigger_end = trigger_start 130 trigger_start, trigger_end = trigger_start - offset, trigger_end - offset + 1 131 event_type = event_type.replace('.', ':') 132 args = event[1:] 133 args_ = [] 134 for arg_start, arg_end, role in args: 135 arg_start, arg_end = arg_start - offset, arg_end - offset +1 136 arg = wnd_entity_map[(arg_start, arg_end)] 137 args_.append({ 138 'entity_id': arg['id'], 139 'text': arg['text'], 140 'role': role 141 }) 142 event_obj = { 143 'event_type': event_type, 144 'id': event_id, 145 'trigger': { 146 'start': trigger_start, 147 'end': trigger_end, 148 'text': ' '.join(wnd_tokens[trigger_start:trigger_end]) 149 }, 150 'arguments': args_ 151 } 152 wnd_events_.append(event_obj) 153 wnd_event_map[(trigger_start, trigger_end)] = event_obj 154 155 # parse coref events 156 wnd_coref_evts = [[wnd_event_map[(evt[0]-offset, evt[1]-offset+1)] 157 for evt in coref if wnd_start <= evt[0] and evt[1] < wnd_end 158 and (evt[0]-offset, evt[1]-offset+1) in wnd_event_map] for coref in coref_events] 159 wnd_coref_evts_ = [] 160 for j, evt_list in enumerate(wnd_coref_evts): 161 if len(evt_list) > 1: 162 wnd_coref_evts_.append({ 163 'id': '{}-CEV{}'.format(wnd_id, j), 164 'events': evt_list 165 }) 166 167 wnd_ = { 168 'doc_id': doc_id, 169 'wnd_id': wnd_id, 170 'entity_mentions': wnd_entities_, 171 'relation_mentions': wnd_relations_, 172 'event_mentions': wnd_events_, 173 'entity_coreference': wnd_coref_ents_, 174 'event_coreference': wnd_coref_evts_, 175 'tokens': wnd_tokens, 176 'pieces': [p for w in pieces for p in w], 177 'token_lens': word_lens, 178 'sentence': ' '.join(wnd_tokens), 179 'sentence_starts': [x-offset for x in wnd_sent_starts[:-1]], 180 } 181 w.write(json.dumps(wnd_) + '\n') 182 offset += len(sentences[i]) 183 184 185 if __name__ == '__main__': 186 parser = ArgumentParser() 187 parser.add_argument('-i', '--input', help='Path to the input file') 188 parser.add_argument('-o', '--output', help='Path to the output file') 189 parser.add_argument('-b', '--bert', help='BERT model name', default='bert-large-cased') 190 parser.add_argument('-w', '--window', default=1, help='Integer for window size', type=int) 191 args = parser.parse_args() 192 model_name = args.bert 193 if model_name.startswith('bert-'): 194 bert_tokenizer = BertTokenizer.from_pretrained(args.bert, 195 do_lower_case=False) 196 elif model_name.startswith('roberta-'): 197 bert_tokenizer = RobertaTokenizer.from_pretrained(args.bert, 198 do_lower_case=False) 199 else: 200 bert_tokenizer = AutoTokenizer.from_pretrained(args.bert, do_lower_case=False, use_fast=False) 201 202 convert(args.input, args.output, bert_tokenizer, args.window)