Beispiel #1
0
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 = NeuralKB(kernel=kernel, scoring_type=scoring_type)

        indices = torch.LongTensor(np.array([predicate_to_index['p'], predicate_to_index['q']]))
        _hops = SymbolicReformulator(predicate_embeddings, indices)
        hoppy = Hoppy(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.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)

                        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)

                        facts = [rel_emb, arg1_emb, arg2_emb]

                        inf = hoppy.score(xp_emb, xs_emb, xo_emb, facts=facts,
                                          entity_embeddings=entity_embeddings.weight)
                        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
Beispiel #2
0
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.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]

                    base = NeuralKB(entity_embeddings=entity_emb, predicate_embeddings=predicate_emb,
                                    kernel=kernel, facts=facts, scoring_type=st)

                    indices = torch.LongTensor(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.LongTensor(xs_np)
                    xp = torch.LongTensor(xp_np)
                    xo = torch.LongTensor(xo_np)
                    xi = torch.LongTensor(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-5, atol=1e-5)
                        np.testing.assert_allclose(inf[i], scores_po[i, xs[i]], rtol=1e-5, atol=1e-5)
Beispiel #3
0
def test_clutrr_v7():
    torch.set_num_threads(multiprocessing.cpu_count())

    embedding_size = 50

    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)

        facts = [rel_emb, arg1_emb, arg2_emb]

        k = 5

        model = NeuralKB(kernel=kernel)

        indices = torch.from_numpy(
            np.array([predicate_to_index['p'], predicate_to_index['q']]))
        reformulator = SymbolicReformulator(predicate_embeddings, indices)

        hoppy0 = Hoppy(model, hops_lst=[(reformulator, False)], depth=0)
        hoppy1 = Hoppy(model, hops_lst=[(reformulator, False)], depth=1)
        hoppy2 = Hoppy(model, hops_lst=[(reformulator, False)], depth=2)
        hoppy3 = Hoppy(model, hops_lst=[(reformulator, False)], depth=3)
        hoppy4 = Hoppy(model, hops_lst=[(reformulator, False)], depth=4)

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

        # res0 = hoppy0.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
        inf0 = hoppy0.score(xp_emb,
                            xs_emb,
                            xo_emb,
                            facts=facts,
                            entity_embeddings=entity_embeddings.weight)
        # (scores0_sp, subs0_sp), (scores0_po, subs0_po) = res0

        # res1 = hoppy1.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
        inf1 = hoppy1.score(xp_emb,
                            xs_emb,
                            xo_emb,
                            facts=facts,
                            entity_embeddings=entity_embeddings.weight)
        # (scores1_sp, subs1_sp), (scores1_po, subs1_po) = res1

        # res2 = hoppy2.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
        inf2 = hoppy2.score(xp_emb,
                            xs_emb,
                            xo_emb,
                            facts=facts,
                            entity_embeddings=entity_embeddings.weight)
        # (scores2_sp, subs2_sp), (scores2_po, subs2_po) = res2

        # res3 = hoppy3.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
        inf3 = hoppy3.score(xp_emb,
                            xs_emb,
                            xo_emb,
                            facts=facts,
                            entity_embeddings=entity_embeddings.weight)
        # (scores3_sp, subs3_sp), (scores3_po, subs3_po) = res3

        # res4 = hoppy4.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
        inf4 = hoppy4.score(xp_emb,
                            xs_emb,
                            xo_emb,
                            facts=facts,
                            entity_embeddings=entity_embeddings.weight)
        # (scores4_sp, subs4_sp), (scores4_po, subs4_po) = res4

        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)

        print(inf3_np)

        print(entity_embeddings.weight[entity_to_index['c'], 0].item())
        print(entity_embeddings.weight[entity_to_index['e'], 0].item())
        print(entity_embeddings.weight[entity_to_index['g'], 0].item())
        print(entity_embeddings.weight[entity_to_index['i'], 0].item())
Beispiel #4
0
def test_reasoning_v6():
    torch.set_num_threads(multiprocessing.cpu_count())

    embedding_size = 50

    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 ['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']]))
            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-5,
                                           atol=1e-5)
                np.testing.assert_allclose(inf_np[i],
                                           scores_po_np[i, xs[i]],
                                           rtol=1e-5,
                                           atol=1e-5)

            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-5,
                                           atol=1e-5)
                np.testing.assert_allclose(inf_np[i],
                                           scores_po_np[i, xs[i]],
                                           rtol=1e-5,
                                           atol=1e-5)

            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)
Beispiel #5
0
def test_reasoning_v5():
    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)
            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-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)