data.py (9368B)
1 import json, logging, pickle 2 import torch 3 from torch.utils.data import Dataset 4 from collections import namedtuple 5 6 logger = logging.getLogger(__name__) 7 8 ee_instance_fields = ['doc_id', 'wnd_id', 'tokens', 'pieces', 'piece_idxs', 'token_lens', 'token_start_idxs', 'triggers', 'roles'] 9 ee_batch_fields = ['tokens', 'pieces', 'piece_idxs', 'token_lens', 'token_start_idxs', 'triggers', 'roles', 'wnd_ids'] 10 EEInstance = namedtuple('EEInstance', field_names=ee_instance_fields, defaults=[None] * len(ee_instance_fields)) 11 EEBatch = namedtuple('EEBatch', field_names=ee_batch_fields, defaults=[None] * len(ee_batch_fields)) 12 13 gen_batch_fields = ['input_text', 'target_text', 'enc_idxs', 'enc_attn', 'dec_idxs', 'dec_attn', 'lbl_idxs', 'raw_lbl_idxs', 'infos'] 14 GenBatch = namedtuple('GenBatch', field_names=gen_batch_fields, defaults=[None] * len(gen_batch_fields)) 15 16 def remove_overlap_entities(entities): 17 """There are a few overlapping entities in the data set. We only keep the 18 first one and map others to it. 19 :param entities (list): a list of entity mentions. 20 :return: processed entity mentions and a table of mapped IDs. 21 """ 22 tokens = [None] * 1000 23 entities_ = [] 24 id_map = {} 25 for entity in entities: 26 start, end = entity['start'], entity['end'] 27 break_flag = False 28 for i in range(start, end): 29 if tokens[i]: 30 id_map[entity['id']] = tokens[i] 31 break_flag = True 32 if break_flag: 33 continue 34 entities_.append(entity) 35 for i in range(start, end): 36 tokens[i] = entity['id'] 37 return entities_, id_map 38 39 def get_role_list(entities, events, id_map): 40 entity_idxs = {entity['id']: (i,entity) for i, entity in enumerate(entities)} 41 visited = [[0] * len(entities) for _ in range(len(events))] 42 role_list = [] 43 role_list = [] 44 for i, event in enumerate(events): 45 for arg in event['arguments']: 46 entity_idx = entity_idxs[id_map.get(arg['entity_id'], arg['entity_id'])] 47 48 # This will automatically remove multi role scenario 49 if visited[i][entity_idx[0]] == 0: 50 # ((trigger start, trigger end, trigger type), (argument start, argument end, role type)) 51 temp = ((event['trigger']['start'], event['trigger']['end'], event['event_type']), 52 (entity_idx[1]['start'], entity_idx[1]['end'], arg['role'], entity_idx[1]['text'])) 53 role_list.append(temp) 54 visited[i][entity_idx[0]] = 1 55 role_list.sort(key=lambda x: (x[0][0], x[1][0])) 56 return role_list 57 58 class EEDataset(Dataset): 59 def __init__(self, tokenizer, path, max_length=128, fair_compare=True): 60 self.tokenizer = tokenizer 61 self.path = path 62 self.data = [] 63 self.insts = [] 64 self.max_length = max_length 65 self.fair_compare = fair_compare 66 self.load_data() 67 68 def __len__(self): 69 return len(self.data) 70 71 def __getitem__(self, item): 72 return self.data[item] 73 74 @property 75 def event_type_set(self): 76 type_set = set() 77 for inst in self.insts: 78 for event in inst['event_mentions']: 79 type_set.add(event['event_type']) 80 return type_set 81 82 @property 83 def role_type_set(self): 84 type_set = set() 85 for inst in self.insts: 86 for event in inst['event_mentions']: 87 for arg in event['arguments']: 88 type_set.add(arg['role']) 89 return type_set 90 91 def load_data(self): 92 with open(self.path, 'r', encoding='utf-8') as fp: 93 lines = fp.readlines() 94 self.insts = [] 95 for line in lines: 96 inst = json.loads(line) 97 inst_len = len(inst['pieces']) 98 #if inst_len > self.max_length: 99 # continue 100 self.insts.append(inst) 101 102 for inst in self.insts: 103 doc_id = inst['doc_id'] 104 wnd_id = inst['wnd_id'] 105 tokens = inst['tokens'] 106 pieces = inst['pieces'] 107 108 entities = inst['entity_mentions'] 109 if self.fair_compare: 110 entities, entity_id_map = remove_overlap_entities(entities) 111 else: 112 entities = entities 113 entity_id_map = {} 114 115 events = inst['event_mentions'] 116 events.sort(key=lambda x: x['trigger']['start']) 117 118 token_num = len(tokens) 119 token_lens = inst['token_lens'] 120 121 piece_idxs = self.tokenizer.convert_tokens_to_ids(pieces) 122 assert sum(token_lens) == len(piece_idxs) 123 124 triggers = [(e['trigger']['start'], e['trigger']['end'], e['event_type']) for e in events] 125 roles = get_role_list(entities, events, entity_id_map) 126 127 token_start_idxs = [sum(token_lens[:_]) for _ in range(len(token_lens))] + [sum(token_lens)] 128 129 instance = EEInstance( 130 doc_id=doc_id, 131 wnd_id=wnd_id, 132 tokens=tokens, 133 pieces=pieces, 134 piece_idxs=piece_idxs, 135 token_lens=token_lens, 136 token_start_idxs=token_start_idxs, 137 triggers = triggers, 138 roles = roles, 139 ) 140 self.data.append(instance) 141 142 logger.info(f'Loaded {len(self)}/{len(lines)} instances from {self.path}') 143 144 def collate_fn(self, batch): 145 tokens = [inst.tokens for inst in batch] 146 pieces = [inst.pieces for inst in batch] 147 piece_idxs = [inst.piece_idxs for inst in batch] 148 token_lens = [inst.token_lens for inst in batch] 149 token_start_idxs = [inst.token_start_idxs for inst in batch] 150 triggers = [inst.triggers for inst in batch] 151 roles = [inst.roles for inst in batch] 152 wnd_ids = [inst.wnd_id for inst in batch] 153 154 return EEBatch( 155 tokens=tokens, 156 pieces=pieces, 157 piece_idxs=piece_idxs, 158 token_lens=token_lens, 159 token_start_idxs=token_start_idxs, 160 triggers=triggers, 161 roles=roles, 162 wnd_ids=wnd_ids, 163 ) 164 165 class GenDataset(Dataset): 166 def __init__(self, tokenizer, max_length, path, max_output_length=None, unseen_types=[], no_bos=False): 167 self.tokenizer = tokenizer 168 self.max_length = self.max_output_length = max_length 169 if max_output_length is not None: 170 self.max_output_length = max_output_length 171 self.path = path 172 self.no_bos = no_bos # if you use bart, then this should be False; if you use t5, then this should be True 173 self.data = [] 174 self.load_data(unseen_types) 175 176 def __len__(self): 177 return len(self.data) 178 179 def __getitem__(self, item): 180 return self.data[item] 181 182 def load_data(self, unseen_types): 183 with open(self.path, 'rb') as f: 184 data = pickle.load(f) 185 186 for l_in, l_out, l_info in zip(data['input'], data['target'], data['all']): 187 if len(unseen_types) > 0: 188 if isinstance(l_info, tuple): 189 # instance base 190 if l_info[1] in unseen_types: 191 continue 192 else: 193 # trigger base, used in argument model 194 if l_info['event type'] in unseen_types: 195 continue 196 self.data.append({ 197 'input': l_in, 198 'target': l_out, 199 'info': l_info 200 }) 201 logger.info(f'Loaded {len(self)} instances from {self.path}') 202 203 def collate_fn(self, batch): 204 input_text = [x['input'] for x in batch] 205 target_text = [x['target'] for x in batch] 206 207 # encoder inputs 208 inputs = self.tokenizer(input_text, return_tensors='pt', padding=True, max_length=self.max_length) 209 enc_idxs = inputs['input_ids'] 210 enc_attn = inputs['attention_mask'] 211 212 # decoder inputs 213 targets = self.tokenizer(target_text, return_tensors='pt', padding=True, max_length=self.max_output_length) 214 dec_idxs = targets['input_ids'] 215 batch_size = dec_idxs.size(0) 216 dec_idxs[:, 0] = self.tokenizer.eos_token_id 217 dec_attn = targets['attention_mask'] 218 219 # labels 220 padding = torch.ones((batch_size, 1), dtype=torch.long) 221 padding[:] = self.tokenizer.pad_token_id 222 raw_lbl_idxs = torch.cat((dec_idxs[:, 1:], padding), dim=1) 223 lbl_attn = torch.cat((dec_attn[:, 1:], torch.zeros((batch_size, 1), dtype=torch.long)), dim=1) 224 lbl_idxs = raw_lbl_idxs.masked_fill(lbl_attn==0, -100) # ignore padding 225 226 enc_idxs = enc_idxs.cuda() 227 enc_attn = enc_attn.cuda() 228 dec_idxs = dec_idxs.cuda() 229 dec_attn = dec_attn.cuda() 230 raw_lbl_idxs = raw_lbl_idxs.cuda() 231 lbl_idxs = lbl_idxs.cuda() 232 233 return GenBatch( 234 input_text=input_text, 235 target_text=target_text, 236 enc_idxs=enc_idxs, 237 enc_attn=enc_attn, 238 dec_idxs=dec_idxs, 239 dec_attn=dec_attn, 240 lbl_idxs=lbl_idxs, 241 raw_lbl_idxs=raw_lbl_idxs, 242 infos=[x['info'] for x in batch] 243 )