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))
