def _check_topk_score2(score_model, g, num_entity, num_rels, exclude_mode):
hidden_dim = 32
num_entity = 40
num_rels = 4
with tempfile.TemporaryDirectory() as tmpdirname:
entity_emb, rel_emb = generate_rand_emb(score_model.model_name, num_entity, num_rels, hidden_dim, 'none')
create_emb_file(Path(tmpdirname), 'entity.npy', entity_emb.numpy())
create_emb_file(Path(tmpdirname), 'relation.npy', rel_emb.numpy())
score_model.load(Path(tmpdirname))
score_model.attach_graph(g)
score_func = score_model._score_func
head = F.arange(0, num_entity // 2)
rel = F.arange(0, num_rels)
tail = F.arange(num_entity // 2, num_entity)
# exec_model==triplet_wise
tw_rel = np.random.randint(0, num_rels, num_entity // 2)
tw_rel = F.tensor(tw_rel)
result1 = score_model.link_predict(head, tw_rel, tail, exec_mode='triplet_wise', exclude_mode=exclude_mode, batch_size=16)
assert len(result1) == 1
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, tw_rel[i], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[i], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(tw_rel[i]))
tail_ids.append(F.asnumpy(tail[i]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
if exclude_mode is None or exclude_mode == 'mask':
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
if exclude_mode == 'mask':
mask = np.zeros((10,))
for i in range(10):
if (head_ids[i] + 1) % num_entity == tail_ids[i] or \
(head_ids[i] - 1) % num_entity == tail_ids[i]:
mask[i] = 1
else:
c_head_idx = []
c_rel_idx = []
c_tail_idx = []
c_score_topk = []
cur_idx = 0
while len(c_head_idx) < 10:
c_idx = idx[cur_idx]
cur_idx += 1
if (head_ids[c_idx] + 1) % num_entity == tail_ids[c_idx] or \
(head_ids[c_idx] - 1) % num_entity == tail_ids[c_idx]:
continue
c_head_idx.append(head_ids[c_idx])
c_tail_idx.append(tail_ids[c_idx])
c_rel_idx.append(rel_ids[c_idx])
c_score_topk.append(scores[c_idx])
head_ids = F.tensor(c_head_idx)
rel_ids = F.tensor(c_rel_idx)
tail_ids = F.tensor(c_tail_idx)
score_topk = F.tensor(c_score_topk)
r1_head, r1_rel, r1_tail, r1_score, r1_mask = result1[0]
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
if exclude_mode == 'mask':
np.testing.assert_allclose(r1_mask, mask)
else:
assert r1_mask is None
# exec_mode==all
result1 = score_model.link_predict(head, rel, tail, topk=20, exclude_mode=exclude_mode, batch_size=16)
result2 = score_model.link_predict(head=head, tail=tail, topk=20, exclude_mode=exclude_mode, batch_size=16)
assert len(result1) == 1
assert len(result2) == 1
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(rel.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
if exclude_mode is None or exclude_mode == 'mask':
idx = idx[:20]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
if exclude_mode == 'mask':
mask = np.zeros((20,))
for i in range(20):
if (head_ids[i] + 1) % num_entity == tail_ids[i] or \
(head_ids[i] - 1) % num_entity == tail_ids[i]:
mask[i] = 1
else:
c_head_idx = []
c_rel_idx = []
c_tail_idx = []
c_score_topk = []
cur_idx = 0
while len(c_head_idx) < 20:
c_idx = idx[cur_idx]
cur_idx += 1
if (head_ids[c_idx] + 1) % num_entity == tail_ids[c_idx] or \
(head_ids[c_idx] - 1) % num_entity == tail_ids[c_idx]:
continue
c_head_idx.append(head_ids[c_idx])
c_tail_idx.append(tail_ids[c_idx])
c_rel_idx.append(rel_ids[c_idx])
c_score_topk.append(scores[c_idx])
head_ids = F.tensor(c_head_idx)
rel_ids = F.tensor(c_rel_idx)
tail_ids = F.tensor(c_tail_idx)
score_topk = F.tensor(c_score_topk)
r1_head, r1_rel, r1_tail, r1_score, r1_mask = result1[0]
r2_head, r2_rel, r2_tail, r2_score, r2_mask = result2[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
if exclude_mode == 'mask':
np.testing.assert_allclose(r1_mask, mask)
np.testing.assert_allclose(r2_mask, mask)
else:
assert r1_mask is None
assert r2_mask is None
result1 = score_model.link_predict(head, rel, tail, exec_mode='batch_rel', exclude_mode=exclude_mode, batch_size=16)
result2 = score_model.link_predict(head=head, tail=tail, exec_mode='batch_rel', exclude_mode=exclude_mode, batch_size=16)
assert len(result1) == num_rels
assert len(result2) == num_rels
for j in range(rel.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
if exclude_mode is None or exclude_mode == 'mask':
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
if exclude_mode == 'mask':
mask = np.full((10,), False)
for i in range(10):
if (head_ids[i] + 1) % num_entity == tail_ids[i] or \
(head_ids[i] - 1) % num_entity == tail_ids[i]:
mask[i] = True
else:
c_head_idx = []
c_rel_idx = []
c_tail_idx = []
c_score_topk = []
cur_idx = 0
while len(c_head_idx) < 10:
c_idx = idx[cur_idx]
cur_idx += 1
if (head_ids[c_idx] + 1) % num_entity == tail_ids[c_idx] or \
(head_ids[c_idx] - 1) % num_entity == tail_ids[c_idx]:
continue
c_head_idx.append(head_ids[c_idx])
c_tail_idx.append(tail_ids[c_idx])
c_rel_idx.append(rel_ids[c_idx])
c_score_topk.append(scores[c_idx])
head_ids = F.tensor(c_head_idx)
rel_ids = F.tensor(c_rel_idx)
tail_ids = F.tensor(c_tail_idx)
score_topk = F.tensor(c_score_topk)
r1_head, r1_rel, r1_tail, r1_score, r1_mask = result1[j]
r2_head, r2_rel, r2_tail, r2_score, r2_mask = result2[j]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
if exclude_mode == 'mask':
np.testing.assert_allclose(r1_mask, mask)
np.testing.assert_allclose(r2_mask, mask)
else:
assert r1_mask is None
assert r2_mask is None
head = F.arange(0, num_entity)
rel = F.arange(0, num_rels)
tail = F.arange(0, num_entity)
result1 = score_model.link_predict(head, rel, tail, exec_mode='batch_head', exclude_mode=exclude_mode, batch_size=16)
result2 = score_model.link_predict(exec_mode='batch_head', exclude_mode=exclude_mode, batch_size=16)
assert len(result1) == num_entity
assert len(result2) == num_entity
for i in range(head.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for j in range(rel.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
if exclude_mode is None or exclude_mode == 'mask':
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
if exclude_mode == 'mask':
mask = np.full((10,), False)
for l in range(10):
if (head_ids[l] + 1) % num_entity == tail_ids[l] or \
(head_ids[l] - 1) % num_entity == tail_ids[l]:
mask[l] = True
else:
c_head_idx = []
c_rel_idx = []
c_tail_idx = []
c_score_topk = []
cur_idx = 0
while len(c_head_idx) < 10:
c_idx = idx[cur_idx]
cur_idx += 1
if (head_ids[c_idx] + 1) % num_entity == tail_ids[c_idx] or \
(head_ids[c_idx] - 1) % num_entity == tail_ids[c_idx]:
continue
c_head_idx.append(head_ids[c_idx])
c_tail_idx.append(tail_ids[c_idx])
c_rel_idx.append(rel_ids[c_idx])
c_score_topk.append(scores[c_idx])
head_ids = F.tensor(c_head_idx)
rel_ids = F.tensor(c_rel_idx)
tail_ids = F.tensor(c_tail_idx)
score_topk = F.tensor(c_score_topk)
r1_head, r1_rel, r1_tail, r1_score, r1_mask = result1[i]
r2_head, r2_rel, r2_tail, r2_score, r2_mask = result2[i]
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
if exclude_mode == 'mask':
np.testing.assert_allclose(r1_mask, mask)
np.testing.assert_allclose(r2_mask, mask)
else:
assert r1_mask is None
assert r2_mask is None
result1 = score_model.link_predict(head, rel, tail, exec_mode='batch_tail', exclude_mode=exclude_mode)
result2 = score_model.link_predict(exec_mode='batch_tail', exclude_mode=exclude_mode)
assert len(result1) == num_entity
assert len(result2) == num_entity
for k in range(tail.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(rel.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
if exclude_mode is None or exclude_mode == 'mask':
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
if exclude_mode == 'mask':
mask = np.full((10,), False)
for l in range(10):
if (head_ids[l] + 1) % num_entity == tail_ids[l] or \
(head_ids[l] - 1) % num_entity == tail_ids[l]:
mask[l] = True
else:
c_head_idx = []
c_rel_idx = []
c_tail_idx = []
c_score_topk = []
cur_idx = 0
while len(c_head_idx) < 10:
c_idx = idx[cur_idx]
cur_idx += 1
if (head_ids[c_idx] + 1) % num_entity == tail_ids[c_idx] or \
(head_ids[c_idx] - 1) % num_entity == tail_ids[c_idx]:
continue
c_head_idx.append(head_ids[c_idx])
c_tail_idx.append(tail_ids[c_idx])
c_rel_idx.append(rel_ids[c_idx])
c_score_topk.append(scores[c_idx])
head_ids = F.tensor(c_head_idx)
rel_ids = F.tensor(c_rel_idx)
tail_ids = F.tensor(c_tail_idx)
score_topk = F.tensor(c_score_topk)
r1_head, r1_rel, r1_tail, r1_score, r1_mask = result1[k]
r2_head, r2_rel, r2_tail, r2_score, r2_mask = result2[k]
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
if exclude_mode == 'mask':
np.testing.assert_allclose(r1_mask, mask)
np.testing.assert_allclose(r2_mask, mask)
else:
assert r1_mask is None
assert r2_mask is None
def run_topk_emb(sfunc, sim_func, create_emb_sim=create_kge_emb_sim):
hidden_dim = 32
num_head = 40
num_tail = 40
num_emb = 80
emb = F.uniform((num_emb, hidden_dim), F.float32, F.cpu(), -1, 1)
head = F.arange(0, num_head)
tail = F.arange(num_head, num_head+num_tail)
sim_infer = create_emb_sim(emb, sfunc)
result1 = sim_infer.topK(head, tail, pair_ws=True)
scores = []
head_ids = []
tail_ids = []
for i in range(head.shape[0]):
j = i
hemb = F.take(emb, head[i], 0)
temb = F.take(emb, tail[j], 0)
score = sim_func(hemb, temb)
scores.append(F.asnumpy(score))
head_ids.append(F.asnumpy(head[i]))
tail_ids.append(F.asnumpy(tail[j]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_tail, r1_score = result1[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
print('pass pair wise')
head = F.arange(0, num_head)
tail = F.arange(num_head, num_head+num_tail)
result1 = sim_infer.topK(head, tail)
assert len(result1) == 1
scores = []
head_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(tail.shape[0]):
hemb = F.take(emb, head[i], 0)
temb = F.take(emb, tail[j], 0)
score = sim_func(hemb, temb)
scores.append(F.asnumpy(score))
head_ids.append(F.asnumpy(head[i]))
tail_ids.append(F.asnumpy(tail[j]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_tail, r1_score = result1[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
emb_ids = F.arange(0, num_emb)
result1 = sim_infer.topK(emb_ids, emb_ids, bcast=True)
result2 = sim_infer.topK(bcast=True)
assert len(result1) == emb_ids.shape[0]
assert len(result2) == emb_ids.shape[0]
for i in range(emb_ids.shape[0]):
scores = []
head_ids = []
tail_ids = []
for j in range(emb_ids.shape[0]):
hemb = F.take(emb, emb_ids[i], 0)
temb = F.take(emb, emb_ids[j], 0)
score = sim_func(hemb, temb)
score = F.asnumpy(score)
scores.append(score)
tail_ids.append(F.asnumpy(emb_ids[j]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = np.full((10,), F.asnumpy(emb_ids[i]))
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_tail, r1_score = result1[i]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
r2_head, r2_tail, r2_score = result2[i]
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
print('pass all')
def run_topk_emb2(sfunc, sim_func, emb_model):
hidden_dim = 32
num_head = 40
num_tail = 40
num_emb = 80
with tempfile.TemporaryDirectory() as tmpdirname:
emb = F.uniform((num_emb, hidden_dim), F.float32, F.cpu(), -1, 1)
create_emb_file(Path(tmpdirname), 'entity.npy', emb.numpy())
create_emb_file(Path(tmpdirname), 'relation.npy', emb.numpy())
emb_model.load(Path(tmpdirname))
head = F.arange(0, num_head)
tail = F.arange(num_head, num_head+num_tail)
result1 = emb_model.embed_sim(head, tail, 'entity', sfunc=sfunc, pair_ws=True)
scores = []
head_ids = []
tail_ids = []
for i in range(head.shape[0]):
j = i
hemb = F.take(emb, head[i], 0)
temb = F.take(emb, tail[j], 0)
score = sim_func(hemb, temb)
scores.append(F.asnumpy(score))
head_ids.append(F.asnumpy(head[i]))
tail_ids.append(F.asnumpy(tail[j]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_tail, r1_score = result1[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
print('pass pair wise')
head = F.arange(0, num_head)
tail = F.arange(num_head, num_head+num_tail)
result1 = emb_model.embed_sim(head, tail, 'entity', sfunc=sfunc)
assert len(result1) == 1
scores = []
head_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(tail.shape[0]):
hemb = F.take(emb, head[i], 0)
temb = F.take(emb, tail[j], 0)
score = sim_func(hemb, temb)
scores.append(F.asnumpy(score))
head_ids.append(F.asnumpy(head[i]))
tail_ids.append(F.asnumpy(tail[j]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_tail, r1_score = result1[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
emb_ids = F.arange(0, num_emb)
result1 = emb_model.embed_sim(emb_ids, emb_ids, 'entity', sfunc=sfunc, bcast=True)
result2 = emb_model.embed_sim(embed_type='entity', sfunc=sfunc, bcast=True)
assert len(result1) == emb_ids.shape[0]
assert len(result2) == emb_ids.shape[0]
for i in range(emb_ids.shape[0]):
scores = []
head_ids = []
tail_ids = []
for j in range(emb_ids.shape[0]):
hemb = F.take(emb, emb_ids[i], 0)
temb = F.take(emb, emb_ids[j], 0)
score = sim_func(hemb, temb)
score = F.asnumpy(score)
scores.append(score)
tail_ids.append(F.asnumpy(emb_ids[j]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = np.full((10,), F.asnumpy(emb_ids[i]))
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_tail, r1_score = result1[i]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
r2_head, r2_tail, r2_score = result2[i]
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
print('pass all')
def check_topk_score(model_name):
hidden_dim = 32
gamma = 12.0
num_entity = 40
num_rels = 4
entity_emb, rel_emb = generate_rand_emb(model_name, num_entity, num_rels, hidden_dim, 'none')
score_model, score_func = create_score_infer(model_name, entity_emb, rel_emb)
head = F.arange(0, num_entity // 2)
rel = F.arange(0, num_rels)
tail = F.arange(num_entity // 2, num_entity)
# exec_model==triplet_wise
tw_rel = np.random.randint(0, num_rels, num_entity // 2)
tw_rel = F.tensor(tw_rel)
result1 = score_model.topK(head, tw_rel, tail, exec_mode='triplet_wise')
assert len(result1) == 1
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, tw_rel[i], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[i], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(tw_rel[i]))
tail_ids.append(F.asnumpy(tail[i]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
# exec_mode==all
result1 = score_model.topK(head, rel, tail, k=20)
result2 = score_model.topK(head=head, tail=tail, k=20)
assert len(result1) == 1
assert len(result2) == 1
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(rel.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:20]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[0]
r2_head, r2_rel, r2_tail, r2_score = result2[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
result1 = score_model.topK(head, rel, tail, exec_mode='batch_rel')
result2 = score_model.topK(head=head, tail=tail, exec_mode='batch_rel')
assert len(result1) == num_rels
assert len(result2) == num_rels
for j in range(rel.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[j]
r2_head, r2_rel, r2_tail, r2_score = result2[j]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
head = F.arange(0, num_entity)
rel = F.arange(0, num_rels)
tail = F.arange(0, num_entity)
result1 = score_model.topK(head, rel, tail, exec_mode='batch_head')
result2 = score_model.topK(exec_mode='batch_head')
assert len(result1) == num_entity
assert len(result2) == num_entity
for i in range(head.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for j in range(rel.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[i]
r2_head, r2_rel, r2_tail, r2_score = result2[i]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
result1 = score_model.topK(head, rel, tail, exec_mode='batch_tail')
result2 = score_model.topK(exec_mode='batch_tail')
assert len(result1) == num_entity
assert len(result2) == num_entity
for k in range(tail.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(rel.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[k]
r2_head, r2_rel, r2_tail, r2_score = result2[k]
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-6, atol=1e-6)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-6, atol=1e-6)
def check_topk_score(model_name):
hidden_dim = 32
gamma = 12.0
num_entity = 40
num_rels = 4
score_model = ScoreInfer(-1, 'config', 'path', 'none')
if model_name == 'TransE' or \
model_name =='TransE_l1' or \
model_name == 'TransE_l2' or \
model_name == 'DistMult' or \
model_name == 'ComplEx':
model = InferModel('cpu', model_name, hidden_dim, batch_size=16)
elif model_name == 'RESCAL':
model = InferModel('cpu', model_name, hidden_dim)
elif model_name == 'RotatE':
model = InferModel('cpu',
model_name,
hidden_dim,
double_entity_emb=True)
entity_emb, rel_emb = generate_rand_emb(model_name, num_entity, num_rels,
hidden_dim, 'none')
model.entity_emb = InferEmbedding('cpu')
model.entity_emb.emb = entity_emb
model.relation_emb = InferEmbedding('cpu')
model.relation_emb.emb = rel_emb
score_model.model = model
score_func = model.score_func
head = F.arange(0, num_entity // 2)
rel = F.arange(0, num_rels)
tail = F.arange(num_entity // 2, num_entity)
# exec_model==triplet_wise
tw_rel = np.random.randint(0, num_rels, num_entity // 2)
tw_rel = F.tensor(tw_rel)
result1 = score_model.topK(head, tw_rel, tail, exec_mode='triplet_wise')
assert len(result1) == 1
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, tw_rel[i], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[i], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(tw_rel[i]))
tail_ids.append(F.asnumpy(tail[i]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
# exec_mode==all
result1 = score_model.topK(head, rel, tail, k=20)
result2 = score_model.topK(head=head, tail=tail, k=20)
assert len(result1) == 1
assert len(result2) == 1
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(rel.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:20]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[0]
r2_head, r2_rel, r2_tail, r2_score = result2[0]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
result1 = score_model.topK(head, rel, tail, exec_mode='batch_rel')
result2 = score_model.topK(head=head, tail=tail, exec_mode='batch_rel')
assert len(result1) == num_rels
assert len(result2) == num_rels
for j in range(rel.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[j]
r2_head, r2_rel, r2_tail, r2_score = result2[j]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
head = F.arange(0, num_entity)
rel = F.arange(0, num_rels)
tail = F.arange(0, num_entity)
result1 = score_model.topK(head, rel, tail, exec_mode='batch_head')
result2 = score_model.topK(exec_mode='batch_head')
assert len(result1) == num_entity
assert len(result2) == num_entity
for i in range(head.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for j in range(rel.shape[0]):
for k in range(tail.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[i]
r2_head, r2_rel, r2_tail, r2_score = result2[i]
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
result1 = score_model.topK(head, rel, tail, exec_mode='batch_tail')
result2 = score_model.topK(exec_mode='batch_tail')
assert len(result1) == num_entity
assert len(result2) == num_entity
for k in range(tail.shape[0]):
scores = []
head_ids = []
rel_ids = []
tail_ids = []
for i in range(head.shape[0]):
for j in range(rel.shape[0]):
hemb = F.take(entity_emb, head[i], 0)
remb = F.take(rel_emb, rel[j], 0)
temb = F.unsqueeze(F.take(entity_emb, tail[k], 0), dim=0)
edge = FakeEdge(hemb, temb, remb)
score = F.asnumpy(score_func.edge_func(edge)['score'])
scores.append(score)
head_ids.append(F.asnumpy(head[i]))
rel_ids.append(F.asnumpy(rel[j]))
tail_ids.append(F.asnumpy(tail[k]))
scores = np.asarray(scores)
scores = scores.reshape(scores.shape[0])
head_ids = np.asarray(head_ids)
rel_ids = np.asarray(rel_ids)
tail_ids = np.asarray(tail_ids)
idx = np.argsort(scores)
idx = idx[::-1]
idx = idx[:10]
head_ids = head_ids[idx]
rel_ids = rel_ids[idx]
tail_ids = tail_ids[idx]
score_topk = scores[idx]
r1_head, r1_rel, r1_tail, r1_score = result1[k]
r2_head, r2_rel, r2_tail, r2_score = result2[k]
np.testing.assert_allclose(r1_head, head_ids)
np.testing.assert_allclose(r2_head, head_ids)
np.testing.assert_allclose(r1_rel, rel_ids)
np.testing.assert_allclose(r2_rel, rel_ids)
np.testing.assert_allclose(r1_tail, tail_ids)
np.testing.assert_allclose(r2_tail, tail_ids)
np.testing.assert_allclose(r1_score, score_topk, rtol=1e-6, atol=1e-6)
np.testing.assert_allclose(r2_score, score_topk, rtol=1e-6, atol=1e-6)
