def vectorize_action_space(action_space_list, action_space_size): bucket_size = len(action_space_list) r_space = torch.zeros(bucket_size, action_space_size) + self.dummy_r e_space = torch.zeros(bucket_size, action_space_size) + self.dummy_e action_mask = torch.zeros(bucket_size, action_space_size) for i, action_space in enumerate(action_space_list): for j, (r, e) in enumerate(action_space): r_space[i, j] = r e_space[i, j] = e action_mask[i, j] = 1 return (int_var_cuda(r_space), int_var_cuda(e_space)), var_cuda(action_mask)
def build(forks: List[Fork], action_space_size, dummy_r, dummy_e): bucket_size = len(forks) r_space = torch.zeros(bucket_size, action_space_size) + dummy_r e_space = torch.zeros(bucket_size, action_space_size) + dummy_e action_mask = torch.zeros(bucket_size, action_space_size) for i, fork in enumerate(forks): for j, direction in enumerate(fork.directions): r_space[i, j] = direction.rel e_space[i, j] = direction.ent action_mask[i, j] = 1 return ActionSpace(forks, int_var_cuda(r_space), int_var_cuda(e_space), var_cuda(action_mask))
def answers_to_var(d_l): d_v = collections.defaultdict(collections.defaultdict) for x in d_l: for y in d_l[x]: v = torch.LongTensor(list(d_l[x][y])).unsqueeze(1) d_v[x][y] = int_var_cuda(v) return d_v
def top_k_action_r(log_action_dist): """ Get top k relations. - k = beam_size if the beam size is smaller than or equal to the beam action space size - k = beam_action_space_size otherwise :param log_action_dist: [batch_size*k, action_space_size] :return: next_r, log_action_prob, action_offset: [batch_size*new_k] """ full_size = len(log_action_dist) assert (full_size % batch_size == 0) last_k = int(full_size / batch_size) action_space_size = log_action_dist.size()[1] # => [batch_size, k'*action_space_size] log_action_dist = log_action_dist.view(batch_size, -1) beam_action_space_size = log_action_dist.size()[1] k = min(beam_size, beam_action_space_size) # [batch_size, k] log_action_prob, action_ind = torch.topk(log_action_dist, k) next_r = (action_ind % action_space_size).view(-1) # [batch_size, k] => [batch_size*k] log_action_prob = log_action_prob.view(-1) # compute parent offset # [batch_size, k] action_beam_offset = action_ind / action_space_size # [batch_size, 1] action_batch_offset = int_var_cuda(torch.arange(batch_size) * last_k).unsqueeze(1) # [batch_size, k] => [batch_size*k] action_offset = (action_batch_offset + action_beam_offset).view(-1) return next_r, log_action_prob, action_offset
def vectorize_unique_r_space(unique_r_space_list, unique_r_space_size, volatile): bucket_size = len(unique_r_space_list) unique_r_space = torch.zeros(bucket_size, unique_r_space_size) + self.dummy_r for i, u_r_s in enumerate(unique_r_space_list): for j, r in enumerate(u_r_s): unique_r_space[i, j] = r return int_var_cuda(unique_r_space)
def top_k_action(log_action_dist, action_space): """ Get top k actions. - k = beam_size if the beam size is smaller than or equal to the beam action space size - k = beam_action_space_size otherwise :param log_action_dist: [batch_size*k, action_space_size] :param action_space (r_space, e_space): r_space: [batch_size*k, action_space_size] e_space: [batch_size*k, action_space_size] :return: (next_r, next_e), action_prob, action_offset: [batch_size*new_k] """ full_size = len(log_action_dist) assert (full_size % batch_size == 0) last_k = int(full_size / batch_size) (r_space, e_space), _ = action_space action_space_size = r_space.size()[1] # => [batch_size, k'*action_space_size] log_action_dist = log_action_dist.view(batch_size, -1) beam_action_space_size = log_action_dist.size()[1] k = min(beam_size, beam_action_space_size) # [batch_size, k] log_action_prob, action_ind = torch.topk(log_action_dist, k) next_r = ops.batch_lookup(r_space.view(batch_size, -1), action_ind).view(-1) # [batch_size*k] next_e = ops.batch_lookup(e_space.view(batch_size, -1), action_ind).view(-1) # [batch_size*k] # [batch_size, k] => [batch_size*k] log_action_prob = log_action_prob.view(-1) # [batch_size*k] # *** compute parent offset # [batch_size, k] action_beam_offset = action_ind / action_space_size # [batch_size, 1] action_batch_offset = int_var_cuda(torch.arange(batch_size) * last_k).unsqueeze(1) # [batch_size, k] => [batch_size*k] action_offset = (action_batch_offset + action_beam_offset).view(-1) return (next_r, next_e), log_action_prob, action_offset # [batch_size*k]
def top_k_action(log_action_dist, action_space, return_merge_scores=None): """ Get top k actions. - k = beam_size if the beam size is smaller than or equal to the beam action space size - k = beam_action_space_size otherwise :param log_action_dist: [batch_size*k, action_space_size] :param action_space (r_space, e_space): r_space: [batch_size*k, action_space_size] e_space: [batch_size*k, action_space_size] :return: (next_r, next_e), action_prob, action_offset: [batch_size*new_k] """ full_size = len(log_action_dist) assert (full_size % batch_size == 0) last_k = int(full_size / batch_size) (r_space, e_space), _ = action_space action_space_size = r_space.size()[1] # => [batch_size, k'*action_space_size] log_action_dist = log_action_dist.view(batch_size, -1) beam_action_space_size = log_action_dist.size()[1] k = min(beam_size, beam_action_space_size) if return_merge_scores is not None: if return_merge_scores == 'sum': reduce_method = torch.sum elif return_merge_scores == 'mean': reduce_method = torch.mean else: reduce_method = None all_action_ind = torch.LongTensor([range(beam_action_space_size) for _ in range(len(log_action_dist))]).cuda() # _, all_action_ind = torch.topk(all_action_ind, beam_action_space_size, largest=False) # print("all_action_ind:", all_action_ind.shape, all_action_ind) # print ("DEBUG all_action_ind:", all_action_ind.shape, all_action_ind) all_next_r = ops.batch_lookup(r_space.view(batch_size, -1), all_action_ind) all_next_e = ops.batch_lookup(e_space.view(batch_size, -1), all_action_ind) # print ("DEBUG all_next_e:", all_next_e.shape, all_next_e) # print ("DEBUG all_next_r:", all_next_r.shape, all_next_r) real_log_action_prob, real_next_e, real_action_ind, real_next_r = ops.merge_same(log_action_dist, all_next_e, all_next_r, method=reduce_method) # print("DEBUG real_log_action_prob:", real_log_action_prob.shape, real_log_action_prob) # print("DEBUG real_next_e:", real_next_e.shape, real_next_e) next_e_list, next_r_list, action_ind_list, log_action_prob_list = [], [], [], [] for i in range(batch_size): k_prime = min(len(real_log_action_prob[i]), k) top_log_prob, top_ind = torch.topk(real_log_action_prob[i], k_prime) top_next_e, top_next_r, top_ind = real_next_e[i][top_ind], real_next_r[i][top_ind], real_action_ind[i][top_ind] next_e_list.append(top_next_e.unsqueeze(0)) next_r_list.append(top_next_r.unsqueeze(0)) action_ind_list.append(top_ind.unsqueeze(0)) log_action_prob_list.append(top_log_prob.unsqueeze(0)) # print("DEBUG -->next_e_list:", next_e_list, next_e_list) # print("DEBUG -->next_r_list:", next_r_list, next_r_list) next_r = ops.pad_and_cat(next_r_list, padding_value=kg.dummy_r).view(-1) next_e = ops.pad_and_cat(next_e_list, padding_value=kg.dummy_e).view(-1) log_action_prob = ops.pad_and_cat(log_action_prob_list, padding_value=0.0).view(-1) action_ind = ops.pad_and_cat(action_ind_list, padding_value=-1).view(-1) # print("DEBUG next_r, next_e:", next_e.shape, next_r.shape) # next_r = ops.pad_and_cat(next_r_list, padding_value=kg.dummy_r, padding_dim=0).view(-1) # next_e = ops.pad_and_cat(next_e_list, padding_value=kg.dummy_e, padding_dim=0).view(-1) # log_action_prob = ops.pad_and_cat(log_action_prob_list, padding_value=0.0, padding_dim=0).view(-1) # action_ind = ops.pad_and_cat(action_ind_list, padding_value=-1, padding_dim=0).view(-1) else: log_action_prob, action_ind = torch.topk(log_action_dist, k) next_r = ops.batch_lookup(r_space.view(batch_size, -1), action_ind).view(-1) next_e = ops.batch_lookup(e_space.view(batch_size, -1), action_ind).view(-1) # print ("log_action_dist:", log_action_dist) #old start # log_action_prob, action_ind = torch.topk(log_action_dist, k) # next_r = ops.batch_lookup(r_space.view(batch_size, -1), action_ind).view(-1) # next_e = ops.batch_lookup(e_space.view(batch_size, -1), action_ind).view(-1) #old end # [batch_size, k] => [batch_size*k] log_action_prob = log_action_prob.view(-1) # *** compute parent offset # [batch_size, k] action_beam_offset = action_ind / action_space_size # [batch_size, 1] action_batch_offset = int_var_cuda(torch.arange(batch_size) * last_k).unsqueeze(1) # [batch_size, k] => [batch_size*k] action_offset = (action_batch_offset + action_beam_offset).view(-1) return (next_r, next_e), log_action_prob, action_offset
def export_fuzzy_facts(self): """ Export high confidence facts according to the model. """ kg, mdl = self.kg, self.mdl # Gather all possible (subject, relation) and (relation, object) pairs sub_rel, rel_obj = {}, {} for file_name in ['raw.kb', 'train.triples', 'dev.triples', 'test.triples']: with open(os.path.join(self.data_dir, file_name)) as f: for line in f: e1, e2, r = line.strip().split() e1_id, e2_id, r_id = kg.triple2ids((e1, e2, r)) if not e1_id in sub_rel: sub_rel[e1_id] = {} if not r_id in sub_rel[e1_id]: sub_rel[e1_id][r_id] = set() sub_rel[e1_id][r_id].add(e2_id) if not e2_id in rel_obj: rel_obj[e2_id] = {} if not r_id in rel_obj[e2_id]: rel_obj[e2_id][r_id] = set() rel_obj[e2_id][r_id].add(e1_id) o_f = open(os.path.join(self.data_dir, 'train.fuzzy.triples'), 'w') print('Saving fuzzy facts to {}'.format(os.path.join(self.data_dir, 'train.fuzzy.triples'))) count = 0 # Save recovered objects e1_ids, r_ids = [], [] for e1_id in sub_rel: for r_id in sub_rel[e1_id]: e1_ids.append(e1_id) r_ids.append(r_id) for i in range(0, len(e1_ids), self.batch_size): e1_ids_b = e1_ids[i:i+self.batch_size] r_ids_b = r_ids[i:i+self.batch_size] e1 = var_cuda(torch.LongTensor(e1_ids_b)) r = var_cuda(torch.LongTensor(r_ids_b)) pred_scores = mdl.forward(e1, r, kg) for j in range(pred_scores.size(0)): for _e2 in range(pred_scores.size(1)): if _e2 in [NO_OP_ENTITY_ID, DUMMY_ENTITY_ID]: continue if pred_scores[j, _e2] >= self.theta: _e1 = int(e1[j]) _r = int(r[j]) o_f.write('{}\t{}\t{}\t{}\n'.format( kg.id2entity[_e1], kg.id2entity[_e2], kg.id2relation[_r], float(pred_scores[j, _e2]))) count += 1 if count % 1000 == 0: print('{} fuzzy facts exported'.format(count)) # Save recovered subjects e2_ids, r_ids = [], [] for e2_id in rel_obj: for r_id in rel_obj[e2_id]: e2_ids.append(e2_id) r_ids.append(r_id) e1 = int_var_cuda(torch.arange(kg.num_entities)) for i in range(len(e2_ids)): r = int_fill_var_cuda(e1.size(), r_ids[i]) e2 = int_fill_var_cuda(e1.size(), e2_ids[i]) pred_scores = mdl.forward_fact(e1, r, e2, kg) for j in range(pred_scores.size(1)): if pred_scores[j] > self.theta: _e1 = int(e1[j]) if _e1 in [NO_OP_ENTITY_ID, DUMMY_ENTITY_ID]: continue _r = int(r[j]) _e2 = int(e2[j]) if _e1 in sub_rel and _r in sub_rel[_e1]: continue o_f.write('{}\t{}\t{}\t{}\n'.format( kg.id2entity[_e1], kg.id2entity[_e2], kg.id2relation[_r], float(pred_scores[j]))) count += 1 if count % 1000 == 0: print('{} fuzzy facts exported'.format(count))