def test_learning_v1(): embedding_size = 50 triples, hops = [], [] for i in range(16): triples += [(f'a{i}', 'p', f'b{i}'), (f'b{i}', 'q', f'c{i}')] hops += [(f'a{i}', 'r', f'c{i}')] entity_lst = sorted({e for (e, _, _) in triples + hops} | {e for (e, _, e) in triples + hops}) predicate_lst = sorted({p for (_, p, _) in triples + hops}) nb_entities, nb_predicates = len(entity_lst), len(predicate_lst) entity_to_index = {e: i for i, e in enumerate(entity_lst)} predicate_to_index = {p: i for i, p in enumerate(predicate_lst)} kernel = GaussianKernel() entity_embeddings = nn.Embedding(nb_entities, embedding_size * 2, sparse=True) predicate_embeddings = nn.Embedding(nb_predicates, embedding_size * 2, sparse=True) fact_rel = torch.LongTensor( np.array([predicate_to_index[p] for (_, p, _) in triples])) fact_arg1 = torch.LongTensor( np.array([entity_to_index[s] for (s, _, _) in triples])) fact_arg2 = torch.LongTensor( np.array([entity_to_index[o] for (_, _, o) in triples])) facts = [fact_rel, fact_arg1, fact_arg2] for st in ['min', 'concat']: model = NeuralKB(entity_embeddings=entity_embeddings, predicate_embeddings=predicate_embeddings, kernel=kernel, facts=facts, scoring_type=st) for s in entity_lst: for p in predicate_lst: for o in entity_lst: xs_np = np.array([entity_to_index[s]]) xp_np = np.array([predicate_to_index[p]]) xo_np = np.array([entity_to_index[o]]) with torch.no_grad(): xs = torch.LongTensor(xs_np) xp = torch.LongTensor(xp_np) xo = torch.LongTensor(xo_np) xs_emb = entity_embeddings(xs) xp_emb = predicate_embeddings(xp) xo_emb = entity_embeddings(xo) inf = model.score(xp_emb, xs_emb, xo_emb) inf_np = inf.cpu().numpy() if (s, p, o) in triples: assert inf_np[0] > 0.95 else: assert inf_np[0] < 0.01
def test_masking_v1(): nb_entities = 10 nb_predicates = 5 embedding_size = 10 init_size = 1.0 rs = np.random.RandomState(0) for _ in range(1): for position in [0, 1]: for st in ['min', 'concat']: with torch.no_grad(): triples = [ ('a', 'p', 'b'), ('c', 'q', 'd') ] entity_to_index = {'a': 0, 'b': 1, 'c': 2, 'd': 3} predicate_to_index = {'p': 0, 'q': 1} kernel = GaussianKernel() entity_embeddings = nn.Embedding(nb_entities, embedding_size * 2, sparse=True) predicate_embeddings = nn.Embedding(nb_predicates, embedding_size * 2, sparse=True) entity_embeddings.weight.data *= init_size predicate_embeddings.weight.data *= init_size fact_rel = torch.LongTensor(np.array([predicate_to_index[p] for (_, p, _) in triples])) fact_arg1 = torch.LongTensor(np.array([entity_to_index[s] for (s, _, _) in triples])) fact_arg2 = torch.LongTensor(np.array([entity_to_index[o] for (_, _, o) in triples])) facts = [fact_rel, fact_arg1, fact_arg2] model = NeuralKB(entity_embeddings=entity_embeddings, predicate_embeddings=predicate_embeddings, kernel=kernel, facts=facts, scoring_type=st) xs_np = rs.randint(nb_entities, size=32) xp_np = rs.randint(nb_predicates, size=32) xo_np = rs.randint(nb_entities, size=32) xi_np = np.array([position] * xs_np.shape[0]) xs_np[0] = 0 xp_np[0] = 0 xo_np[0] = 1 xs_np[1] = 2 xp_np[1] = 1 xo_np[1] = 3 xs = torch.LongTensor(xs_np) xp = torch.LongTensor(xp_np) xo = torch.LongTensor(xo_np) xi = torch.LongTensor(xi_np) xs_emb = entity_embeddings(xs) xp_emb = predicate_embeddings(xp) xo_emb = entity_embeddings(xo) model.mask_indices = xi scores = model.forward(xp_emb, xs_emb, xo_emb) inf = model.score(xp_emb, xs_emb, xo_emb) if position == 0: assert inf[0] < 0.5 assert inf[1] > 0.9 elif position == 1: assert inf[0] > 0.9 assert inf[1] < 0.5 scores_sp, scores_po = scores inf = inf.cpu().numpy() scores_sp = scores_sp.cpu().numpy() scores_po = scores_po.cpu().numpy() for i in range(xs.shape[0]): np.testing.assert_allclose(inf[i], scores_sp[i, xo[i]], rtol=1e-5, atol=1e-5) np.testing.assert_allclose(inf[i], scores_po[i, xs[i]], rtol=1e-5, atol=1e-5)
def test_reasoning_v4(): torch.set_num_threads(multiprocessing.cpu_count()) nb_entities = 10 nb_predicates = 5 embedding_size = 10 rs = np.random.RandomState(0) triples = [('a', 'p', 'b'), ('b', 'q', 'c'), ('c', 'r', 'd'), ('d', 's', 'e')] entity_to_index = {'a': 0, 'b': 1, 'c': 2, 'd': 3, 'e': 4} predicate_to_index = {'p': 0, 'q': 1, 'r': 2, 's': 3} for st in ['min', 'concat']: with torch.no_grad(): kernel = GaussianKernel() entity_embeddings = nn.Embedding(nb_entities, embedding_size * 2, sparse=True) predicate_embeddings = nn.Embedding(nb_predicates, embedding_size * 2, sparse=True) fact_rel = torch.from_numpy( np.array([predicate_to_index[p] for (_, p, _) in triples])) fact_arg1 = torch.from_numpy( np.array([entity_to_index[s] for (s, _, _) in triples])) fact_arg2 = torch.from_numpy( np.array([entity_to_index[o] for (_, _, o) in triples])) facts = [fact_rel, fact_arg1, fact_arg2] model = NeuralKB(entity_embeddings=entity_embeddings, predicate_embeddings=predicate_embeddings, kernel=kernel, facts=facts, scoring_type=st) indices = torch.from_numpy( np.array([ predicate_to_index['p'], predicate_to_index['q'], predicate_to_index['r'], predicate_to_index['s'] ])) reformulator = SymbolicReformulator(predicate_embeddings, indices) rhoppy = RecursiveHoppy(model, entity_embeddings, hops=reformulator, depth=0) xs_np = rs.randint(nb_entities, size=32) xp_np = rs.randint(nb_predicates, size=32) xo_np = rs.randint(nb_entities, size=32) xs_np[0] = 0 xp_np[0] = 0 xo_np[0] = 1 xs_np[1] = 1 xp_np[1] = 1 xo_np[1] = 2 xs_np[2] = 0 xp_np[2] = 3 xo_np[2] = 4 xs = torch.from_numpy(xs_np) xp = torch.from_numpy(xp_np) xo = torch.from_numpy(xo_np) xs_emb = entity_embeddings(xs) xp_emb = predicate_embeddings(xp) xo_emb = entity_embeddings(xo) scores = model.forward(xp_emb, xs_emb, xo_emb) inf = model.score(xp_emb, xs_emb, xo_emb) scores_h = rhoppy.forward(xp_emb, xs_emb, xo_emb) inf_h = rhoppy.score(xp_emb, xs_emb, xo_emb) print(inf) print(inf_h) assert inf[0] > 0.95 assert inf[1] > 0.95 scores_sp, scores_po = scores scores_h_sp, scores_h_po = scores_h inf = inf.cpu().numpy() scores_sp = scores_sp.cpu().numpy() scores_po = scores_po.cpu().numpy() inf_h = inf_h.cpu().numpy() scores_h_sp = scores_h_sp.cpu().numpy() scores_h_po = scores_h_po.cpu().numpy() np.testing.assert_allclose(inf, inf_h) np.testing.assert_allclose(scores_sp, scores_h_sp) np.testing.assert_allclose(scores_po, scores_h_po) for i in range(xs.shape[0]): np.testing.assert_allclose(inf[i], scores_sp[i, xo[i]], rtol=1e-5, atol=1e-5) np.testing.assert_allclose(inf[i], scores_po[i, xs[i]], rtol=1e-5, atol=1e-5) np.testing.assert_allclose(inf_h[i], scores_h_sp[i, xo[i]], rtol=1e-5, atol=1e-5) np.testing.assert_allclose(inf_h[i], scores_h_po[i, xs[i]], rtol=1e-5, atol=1e-5)