def _test_reasoning_v6(_st): torch.set_num_threads(multiprocessing.cpu_count()) embedding_size = 20 torch.manual_seed(0) rs = np.random.RandomState(0) triples = [('a', 'p', 'b'), ('b', 'q', 'c'), ('c', 'p', 'd'), ('d', 'q', 'e'), ('e', 'p', 'f'), ('f', 'q', 'g'), ('g', 'p', 'h'), ('h', 'q', 'i'), ('i', 'p', 'l'), ('l', 'q', 'm'), ('m', 'p', 'n'), ('n', 'q', 'o'), ('o', 'p', 'p'), ('p', 'q', 'q'), ('q', 'p', 'r'), ('r', 'q', 's'), ('s', 'p', 't'), ('t', 'q', 'u'), ('u', 'p', 'v'), ('v', 'q', 'w'), ('x', 'r', 'y'), ('x', 's', 'y')] entity_lst = sorted({e for (e, _, _) in triples} | {e for (_, _, e) in triples}) predicate_lst = sorted({p for (_, p, _) in triples}) nb_entities = len(entity_lst) nb_predicates = 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)} for st in [_st]: 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']])) reformulator = SymbolicReformulator(predicate_embeddings, indices) k = 5 rhoppy0 = RecursiveHoppy(model, entity_embeddings, hops=reformulator, depth=0, k=k) rhoppy1 = RecursiveHoppy(model, entity_embeddings, hops=reformulator, depth=1, k=k) rhoppy2 = RecursiveHoppy(model, entity_embeddings, hops=reformulator, depth=2, k=k) rhoppy3 = RecursiveHoppy(model, entity_embeddings, hops=reformulator, depth=3, k=k) rhoppy4 = RecursiveHoppy(model, entity_embeddings, hops=reformulator, depth=4, k=k) xs_np = rs.randint(nb_entities, size=12) xp_np = rs.randint(nb_predicates, size=12) xo_np = rs.randint(nb_entities, size=12) xs_np[0] = entity_to_index['a'] xp_np[0] = predicate_to_index['r'] xo_np[0] = entity_to_index['c'] xs_np[1] = entity_to_index['a'] xp_np[1] = predicate_to_index['r'] xo_np[1] = entity_to_index['e'] xs_np[2] = entity_to_index['a'] xp_np[2] = predicate_to_index['r'] xo_np[2] = entity_to_index['g'] xs_np[3] = entity_to_index['a'] xp_np[3] = predicate_to_index['r'] xo_np[3] = entity_to_index['i'] xs_np[4] = entity_to_index['a'] xp_np[4] = predicate_to_index['r'] xo_np[4] = entity_to_index['m'] xs_np[5] = entity_to_index['a'] xp_np[5] = predicate_to_index['r'] xo_np[5] = entity_to_index['o'] xs_np[6] = entity_to_index['a'] xp_np[6] = predicate_to_index['r'] xo_np[6] = entity_to_index['q'] xs_np[7] = entity_to_index['a'] xp_np[7] = predicate_to_index['r'] xo_np[7] = entity_to_index['s'] xs_np[8] = entity_to_index['a'] xp_np[8] = predicate_to_index['r'] xo_np[8] = entity_to_index['u'] # xs_np[9] = entity_to_index['a'] # xp_np[9] = predicate_to_index['r'] # xo_np[9] = entity_to_index['w'] 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) scores0 = rhoppy0.forward(xp_emb, xs_emb, xo_emb) inf0 = rhoppy0.score(xp_emb, xs_emb, xo_emb) for i in range(xs.shape[0]): scores_sp, scores_po = scores0 inf_np = inf0.cpu().numpy() scores_sp_np = scores_sp.cpu().numpy() scores_po_np = scores_po.cpu().numpy() np.testing.assert_allclose(inf_np[i], scores_sp_np[i, xo[i]], rtol=1e-2, atol=1e-2) np.testing.assert_allclose(inf_np[i], scores_po_np[i, xs[i]], rtol=1e-2, atol=1e-2) scores1 = rhoppy1.forward(xp_emb, xs_emb, xo_emb) inf1 = rhoppy1.score(xp_emb, xs_emb, xo_emb) for i in range(xs.shape[0]): scores_sp, scores_po = scores1 inf_np = inf1.cpu().numpy() scores_sp_np = scores_sp.cpu().numpy() scores_po_np = scores_po.cpu().numpy() np.testing.assert_allclose(inf_np[i], scores_sp_np[i, xo[i]], rtol=1e-2, atol=1e-2) np.testing.assert_allclose(inf_np[i], scores_po_np[i, xs[i]], rtol=1e-2, atol=1e-2) scores2 = rhoppy2.forward(xp_emb, xs_emb, xo_emb) inf2 = rhoppy2.score(xp_emb, xs_emb, xo_emb) for i in range(xs.shape[0]): scores_sp, scores_po = scores2 inf_np = inf2.cpu().numpy() scores_sp_np = scores_sp.cpu().numpy() scores_po_np = scores_po.cpu().numpy() np.testing.assert_allclose(inf_np[i], scores_sp_np[i, xo[i]], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf_np[i], scores_po_np[i, xs[i]], rtol=1e-1, atol=1e-1) scores3 = rhoppy3.forward(xp_emb, xs_emb, xo_emb) inf3 = rhoppy3.score(xp_emb, xs_emb, xo_emb) for i in range(xs.shape[0]): scores_sp, scores_po = scores3 inf_np = inf3.cpu().numpy() scores_sp_np = scores_sp.cpu().numpy() scores_po_np = scores_po.cpu().numpy() np.testing.assert_allclose(inf_np[i], scores_sp_np[i, xo[i]], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf_np[i], scores_po_np[i, xs[i]], rtol=1e-1, atol=1e-1) scores4 = rhoppy4.forward(xp_emb, xs_emb, xo_emb) inf4 = rhoppy4.score(xp_emb, xs_emb, xo_emb) for i in range(xs.shape[0]): scores_sp, scores_po = scores4 inf_np = inf4.cpu().numpy() scores_sp_np = scores_sp.cpu().numpy() scores_po_np = scores_po.cpu().numpy() np.testing.assert_allclose(inf_np[i], scores_sp_np[i, xo[i]], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf_np[i], scores_po_np[i, xs[i]], rtol=1e-1, atol=1e-1) print(inf0) print(inf1) print(inf2) print(inf3) print(inf4) inf0_np = inf0.cpu().numpy() inf1_np = inf1.cpu().numpy() inf2_np = inf2.cpu().numpy() inf3_np = inf3.cpu().numpy() inf4_np = inf4.cpu().numpy() np.testing.assert_allclose(inf0_np, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf1_np, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf2_np, [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf3_np, [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf4_np, [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0], rtol=1e-1, atol=1e-1)
def _test_reasoning_v5(_st): torch.set_num_threads(multiprocessing.cpu_count()) nb_entities = 10 nb_predicates = 5 embedding_size = 20 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 [_st]: 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) hoppy = SimpleHoppy(model, entity_embeddings, hops=reformulator) rhoppy = RecursiveHoppy(model, entity_embeddings, hops=reformulator, depth=1) 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 = hoppy.forward(xp_emb, xs_emb, xo_emb) inf = hoppy.score(xp_emb, xs_emb, xo_emb) scores_h = rhoppy.depth_r_forward(xp_emb, xs_emb, xo_emb, depth=1) inf_h = rhoppy.depth_r_score(xp_emb, xs_emb, xo_emb, depth=1) print(inf) print(inf_h) assert inf[2] > 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-2, atol=1e-2) np.testing.assert_allclose(inf[i], scores_po[i, xs[i]], rtol=1e-2, atol=1e-2) np.testing.assert_allclose(inf_h[i], scores_h_sp[i, xo[i]], rtol=1e-2, atol=1e-2) np.testing.assert_allclose(inf_h[i], scores_h_po[i, xs[i]], rtol=1e-2, atol=1e-2)
def test_masking_v2(): nb_entities = 10 nb_predicates = 5 embedding_size = 10 rs = np.random.RandomState(0) for _ in range(1): for position in [0, 1, 2]: for st in ['min', 'concat']: with torch.no_grad(): triples = [('a', 'p', 'b'), ('b', 'q', 'c'), ('a', 'p', 'c')] entity_to_index = {'a': 0, 'b': 1, 'c': 2, 'd': 3} predicate_to_index = {'p': 0, 'q': 1} kernel = GaussianKernel() entity_emb = nn.Embedding(nb_entities, embedding_size * 2, sparse=True) predicate_emb = 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] base = NeuralKB(entity_embeddings=entity_emb, predicate_embeddings=predicate_emb, kernel=kernel, facts=facts, scoring_type=st) indices = torch.from_numpy( np.array( [predicate_to_index['p'], predicate_to_index['q']])) reformulator = SymbolicReformulator(predicate_emb, indices) model = SimpleHoppy(base, entity_emb, hops=reformulator) 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] = 1 xp_np[1] = 1 xo_np[1] = 2 xs_np[2] = 0 xp_np[2] = 0 xo_np[2] = 2 xs = torch.from_numpy(xs_np) xp = torch.from_numpy(xp_np) xo = torch.from_numpy(xo_np) xi = torch.from_numpy(xi_np) xs_emb = entity_emb(xs) xp_emb = predicate_emb(xp) xo_emb = entity_emb(xo) # xi = None base.mask_indices = xi scores = model.forward(xp_emb, xs_emb, xo_emb) inf = model.score(xp_emb, xs_emb, xo_emb) if position in {0, 1}: assert inf[2] < 0.5 else: assert inf[2] > 0.9 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-2, atol=1e-2) np.testing.assert_allclose(inf[i], scores_po[i, xs[i]], rtol=1e-2, atol=1e-2)
def test_adv_v5(): embedding_size = 20 torch.manual_seed(0) rs = np.random.RandomState(0) triples = [ ('a', 'p', 'b'), ('a', 'p', 'd'), ('c', 'p', 'd'), ('e', 'q', 'f'), ('f', 'p', 'c'), ('x', 'r', 'y') ] entity_lst = sorted({e for (e, _, _) in triples} | {e for (_, _, e) in triples}) predicate_lst = sorted({p for (_, p, _) in triples}) nb_entities = len(entity_lst) nb_predicates = 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)} 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) rel_emb = encode_relation(facts=triples, relation_embeddings=predicate_embeddings, relation_to_idx=predicate_to_index) arg1_emb, arg2_emb = encode_arguments(facts=triples, entity_embeddings=entity_embeddings, entity_to_idx=entity_to_index) batch_size = 6 fact_size = rel_emb.shape[0] entity_size = entity_embeddings.weight.shape[0] rel_emb = rel_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) arg1_emb = arg1_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) arg2_emb = arg2_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) nb_facts = torch.tensor([fact_size for _ in range(batch_size)], dtype=torch.long) emb = entity_embeddings.weight.view(1, entity_size, -1).repeat(batch_size, 1, 1) _nb_entities = torch.tensor([entity_size for _ in range(batch_size)], dtype=torch.long) facts = [rel_emb, arg1_emb, arg2_emb] model = BatchNeuralKB(kernel=kernel) indices = torch.from_numpy(np.array([predicate_to_index['p']])) reformulator = SymbolicReformulator(predicate_embeddings, indices) unary = BatchUnary(model, hops_lst=[(reformulator, True)]) xs_np = rs.randint(nb_entities, size=batch_size) xp_np = rs.randint(nb_predicates, size=batch_size) xo_np = rs.randint(nb_entities, size=batch_size) xs_np[0] = entity_to_index['b'] xp_np[0] = predicate_to_index['r'] xo_np[0] = entity_to_index['a'] xs_np[1] = entity_to_index['b'] xp_np[1] = predicate_to_index['r'] xo_np[1] = entity_to_index['b'] xs_np[2] = entity_to_index['b'] xp_np[2] = predicate_to_index['r'] xo_np[2] = entity_to_index['c'] xs_np[3] = entity_to_index['b'] xp_np[3] = predicate_to_index['r'] xo_np[3] = entity_to_index['d'] xs_np[4] = entity_to_index['b'] xp_np[4] = predicate_to_index['r'] xo_np[4] = entity_to_index['e'] xs_np[5] = entity_to_index['b'] xp_np[5] = predicate_to_index['r'] xo_np[5] = entity_to_index['f'] 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) inf = unary.score(xp_emb, xs_emb, xo_emb, facts=facts, nb_facts=nb_facts, entity_embeddings=emb, nb_entities=_nb_entities) inf_np = inf.cpu().numpy() print(inf_np) np.testing.assert_allclose(inf_np, [1] * batch_size, rtol=1e-2, atol=1e-2)
def test_adv_v1(): embedding_size = 20 torch.manual_seed(0) rs = np.random.RandomState(0) triples = [ ('a', 'p', 'b'), ('b', 'q', 'c'), ('c', 'p', 'd'), ('d', 'q', 'e'), ('e', 'p', 'f'), ('f', 'q', 'g'), ('g', 'p', 'h'), ('h', 'q', 'i'), ('i', 'p', 'l'), ('l', 'q', 'm'), ('m', 'p', 'n'), ('n', 'q', 'o'), ('o', 'p', 'p'), ('p', 'q', 'q'), ('q', 'p', 'r'), ('r', 'q', 's'), ('s', 'p', 't'), ('t', 'q', 'u'), ('u', 'p', 'v'), ('v', 'q', 'w'), ('x', 'r', 'y'), ('x', 's', 'y') ] entity_lst = sorted({e for (e, _, _) in triples} | {e for (_, _, e) in triples}) predicate_lst = sorted({p for (_, p, _) in triples}) nb_entities = len(entity_lst) nb_predicates = 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)} 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) rel_emb = encode_relation(facts=triples, relation_embeddings=predicate_embeddings, relation_to_idx=predicate_to_index) arg1_emb, arg2_emb = encode_arguments(facts=triples, entity_embeddings=entity_embeddings, entity_to_idx=entity_to_index) batch_size = 12 fact_size = rel_emb.shape[0] entity_size = entity_embeddings.weight.shape[0] rel_emb = rel_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) arg1_emb = arg1_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) arg2_emb = arg2_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) nb_facts = torch.tensor([fact_size for _ in range(batch_size)], dtype=torch.long) emb = entity_embeddings.weight.view(1, entity_size, -1).repeat(batch_size, 1, 1) _nb_entities = torch.tensor([entity_size for _ in range(batch_size)], dtype=torch.long) facts = [rel_emb, arg1_emb, arg2_emb] model = BatchNeuralKB(kernel=kernel) indices = torch.from_numpy(np.array([predicate_to_index['p'], predicate_to_index['q']])) reformulator = SymbolicReformulator(predicate_embeddings, indices) hoppy0 = BatchHoppy(model, hops_lst=[(reformulator, False)], depth=0) hoppy1 = BatchHoppy(model, hops_lst=[(reformulator, False)], depth=1) hoppy2 = BatchHoppy(model, hops_lst=[(reformulator, False)], depth=2) hoppy3 = BatchHoppy(model, hops_lst=[(reformulator, False)], depth=3) xs_np = rs.randint(nb_entities, size=batch_size) xp_np = rs.randint(nb_predicates, size=batch_size) xo_np = rs.randint(nb_entities, size=batch_size) xs_np[0] = entity_to_index['a'] xp_np[0] = predicate_to_index['r'] xo_np[0] = entity_to_index['c'] xs_np[1] = entity_to_index['a'] xp_np[1] = predicate_to_index['r'] xo_np[1] = entity_to_index['e'] xs_np[2] = entity_to_index['a'] xp_np[2] = predicate_to_index['r'] xo_np[2] = entity_to_index['g'] xs_np[3] = entity_to_index['a'] xp_np[3] = predicate_to_index['r'] xo_np[3] = entity_to_index['i'] xs_np[4] = entity_to_index['a'] xp_np[4] = predicate_to_index['r'] xo_np[4] = entity_to_index['m'] xs_np[5] = entity_to_index['a'] xp_np[5] = predicate_to_index['r'] xo_np[5] = entity_to_index['o'] xs_np[6] = entity_to_index['a'] xp_np[6] = predicate_to_index['r'] xo_np[6] = entity_to_index['q'] xs_np[7] = entity_to_index['a'] xp_np[7] = predicate_to_index['r'] xo_np[7] = entity_to_index['s'] xs_np[8] = entity_to_index['a'] xp_np[8] = predicate_to_index['r'] xo_np[8] = entity_to_index['u'] 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) inf0 = hoppy0.score(xp_emb, xs_emb, xo_emb, facts=facts, nb_facts=nb_facts, entity_embeddings=emb, nb_entities=_nb_entities) inf1 = hoppy1.score(xp_emb, xs_emb, xo_emb, facts=facts, nb_facts=nb_facts, entity_embeddings=emb, nb_entities=_nb_entities) inf2 = hoppy2.score(xp_emb, xs_emb, xo_emb, facts=facts, nb_facts=nb_facts, entity_embeddings=emb, nb_entities=_nb_entities) inf3 = hoppy3.score(xp_emb, xs_emb, xo_emb, facts=facts, nb_facts=nb_facts, entity_embeddings=emb, nb_entities=_nb_entities) inf0_np = inf0.cpu().numpy() inf1_np = inf1.cpu().numpy() inf2_np = inf2.cpu().numpy() inf3_np = inf3.cpu().numpy() np.testing.assert_allclose(inf0_np, [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf1_np, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf2_np, [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) np.testing.assert_allclose(inf3_np, [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], rtol=1e-1, atol=1e-1) print(inf3_np)
def test_clutrr_v2(): embedding_size = 20 triples, hops = [], [] xxx = [] 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}')] xxx += [(f'a{i}', 'p', f'c{i}'), (f'a{i}', 'q', f'c{i}'), (f'a{i}', 'r', f'c{i}')] entity_lst = sorted({s for (s, _, _) in triples + hops} | {o for (_, _, o) 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, sparse=True) predicate_embeddings = nn.Embedding(nb_predicates, embedding_size, sparse=True) for scoring_type in ['concat']: # ['min', 'concat']: model = BatchNeuralKB(kernel=kernel, scoring_type=scoring_type) indices = torch.from_numpy(np.array([predicate_to_index['p'], predicate_to_index['q']])) _hops = SymbolicReformulator(predicate_embeddings, indices) hoppy = BatchHoppy(model, hops_lst=[(_hops, False)], depth=1) 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.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) rel_emb = encode_relation(facts=triples, relation_embeddings=predicate_embeddings, relation_to_idx=predicate_to_index) arg1_emb, arg2_emb = encode_arguments(facts=triples, entity_embeddings=entity_embeddings, entity_to_idx=entity_to_index) batch_size = xp.shape[0] fact_size = rel_emb.shape[0] entity_size = entity_embeddings.weight.shape[0] rel_emb = rel_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) arg1_emb = arg1_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) arg2_emb = arg2_emb.view(1, fact_size, -1).repeat(batch_size, 1, 1) nb_facts = torch.tensor([fact_size for _ in range(batch_size)], dtype=torch.long) emb = entity_embeddings.weight.view(1, entity_size, -1).repeat(batch_size, 1, 1) _nb_entities = torch.tensor([entity_size for _ in range(batch_size)], dtype=torch.long) facts = [rel_emb, arg1_emb, arg2_emb] inf = hoppy.score(xp_emb, xs_emb, xo_emb, facts=facts, nb_facts=nb_facts, entity_embeddings=emb, nb_entities=_nb_entities) inf_np = inf.cpu().numpy() print(s, p, o, inf_np) assert inf_np[0] > 0.9 if (s, p, o) in (triples + xxx) else inf_np[0] < 0.1