test_ranking_P.py

import os
import random
import argparse
import logging
import json
import time

import multiprocessing as mp
import scipy.sparse as ssp
from tqdm import tqdm
import networkx as nx
import torch
import numpy as np
import dgl

from type import get_ent_types


def process_files(files, saved_relation2id, add_traspose_rels):
    '''
    files: Dictionary map of file paths to read the triplets from.
    saved_relation2id: Saved relation2id (mostly passed from a trained model) which can be used to map relations to pre-defined indices and filter out the unknown ones.
    '''
    entity2id = {}
    relation2id = saved_relation2id

    triplets = {}

    ent = 0
    rel = 0

    for file_type, file_path in files.items():

        data = []
        with open(file_path) as f:
            file_data = [line.split() for line in f.read().split('\n')[:-1]]

        for triplet in file_data:
            if triplet[0] not in entity2id:
                entity2id[triplet[0]] = ent
                ent += 1
            if triplet[2] not in entity2id:
                entity2id[triplet[2]] = ent
                ent += 1

            # Save the triplets corresponding to only the known relations
            if triplet[1] in saved_relation2id:
                data.append([entity2id[triplet[0]], entity2id[triplet[2]], saved_relation2id[triplet[1]]])

        # if file_type == 'links':
        #    data = random.sample(data, 20)
        triplets[file_type] = np.array(data)


    id2entity = {v: k for k, v in entity2id.items()}
    id2relation = {v: k for k, v in relation2id.items()}

    # Construct the list of adjacency matrix each corresponding to eeach relation. Note that this is constructed only from the train data.
    adj_list = []
    for i in range(len(saved_relation2id)):
        idx = np.argwhere(triplets['graph'][:, 2] == i)
        adj_list.append(ssp.csc_matrix((np.ones(len(idx), dtype=np.uint8), (triplets['graph'][:, 0][idx].squeeze(1), triplets['graph'][:, 1][idx].squeeze(1))), shape=(len(entity2id), len(entity2id))))

    # Add transpose matrices to handle both directions of relations.
    adj_list_aug = adj_list
    if add_traspose_rels:
        adj_list_t = [adj.T for adj in adj_list]
        adj_list_aug = adj_list + adj_list_t

    dgl_adj_list = ssp_multigraph_to_dgl(adj_list_aug)

    return adj_list, dgl_adj_list, triplets, entity2id, relation2id, id2entity, id2relation


def intialize_worker(model, adj_list, dgl_adj_list, id2entity, params, ont_scorer,
                     ent2types):
    global model_, adj_list_, dgl_adj_list_, id2entity_, params_, ont_scorer_, ent2types_
    model_, adj_list_, dgl_adj_list_, id2entity_, params_, ont_scorer_, ent2types_ = model, adj_list, dgl_adj_list, id2entity, params, ont_scorer, ent2types



def get_neg_samples_replacing_head_tail(test_links, adj_list, num_samples=50):

    n, r = adj_list[0].shape[0], len(adj_list)
    heads, tails, rels = test_links[:, 0], test_links[:, 1], test_links[:, 2]

    neg_triplets = []
    for i, (head, tail, rel) in enumerate(zip(heads, tails, rels)):
        neg_triplet = {'head': [[], 0], 'tail': [[], 0]}
        neg_triplet['head'][0].append([head, tail, rel])
        while len(neg_triplet['head'][0]) < num_samples:
            neg_head = head
            neg_tail = np.random.choice(n)

            if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
                neg_triplet['head'][0].append([neg_head, neg_tail, rel])

        neg_triplet['tail'][0].append([head, tail, rel])
        while len(neg_triplet['tail'][0]) < num_samples:
            neg_head = np.random.choice(n)
            neg_tail = tail
            # neg_head, neg_tail, rel = np.random.choice(n), np.random.choice(n), np.random.choice(r)

            if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
                neg_triplet['tail'][0].append([neg_head, neg_tail, rel])

        neg_triplet['head'][0] = np.array(neg_triplet['head'][0])
        neg_triplet['tail'][0] = np.array(neg_triplet['tail'][0])

        neg_triplets.append(neg_triplet)

    return neg_triplets


def get_neg_samples_replacing_head_tail_all(test_links, adj_list):

    n, r = adj_list[0].shape[0], len(adj_list)
    heads, tails, rels = test_links[:, 0], test_links[:, 1], test_links[:, 2]

    neg_triplets = []
    print('sampling negative triplets...')
    for i, (head, tail, rel) in tqdm(enumerate(zip(heads, tails, rels)), total=len(heads)):
        neg_triplet = {'head': [[], 0], 'tail': [[], 0]}
        neg_triplet['head'][0].append([head, tail, rel])
        for neg_tail in range(n):
            neg_head = head

            if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
                neg_triplet['head'][0].append([neg_head, neg_tail, rel])

        neg_triplet['tail'][0].append([head, tail, rel])
        for neg_head in range(n):
            neg_tail = tail

            if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
                neg_triplet['tail'][0].append([neg_head, neg_tail, rel])

        neg_triplet['head'][0] = np.array(neg_triplet['head'][0])
        neg_triplet['tail'][0] = np.array(neg_triplet['tail'][0])

        neg_triplets.append(neg_triplet)

    return neg_triplets


# def get_neg_samples_replacing_head_tail_from_ruleN(ruleN_pred_path, entity2id, saved_relation2id):
#     with open(ruleN_pred_path) as f:
#         pred_data = [line.split() for line in f.read().split('\n')[:-1]]
#
#     neg_triplets = []
#     for i in range(len(pred_data) // 3):
#         neg_triplet = {'head': [[], 10000], 'tail': [[], 10000]}
#         if pred_data[3 * i][1] in saved_relation2id:
#             head, rel, tail = entity2id[pred_data[3 * i][0]], saved_relation2id[pred_data[3 * i][1]], entity2id[pred_data[3 * i][2]]
#             for j, new_head in enumerate(pred_data[3 * i + 1][1::2]):
#                 neg_triplet['head'][0].append([entity2id[new_head], tail, rel])
#                 if entity2id[new_head] == head:
#                     neg_triplet['head'][1] = j
#             for j, new_tail in enumerate(pred_data[3 * i + 2][1::2]):
#                 neg_triplet['tail'][0].append([head, entity2id[new_tail], rel])
#                 if entity2id[new_tail] == tail:
#                     neg_triplet['tail'][1] = j
#
#             neg_triplet['head'][0] = np.array(neg_triplet['head'][0])
#             neg_triplet['tail'][0] = np.array(neg_triplet['tail'][0])
#
#             neg_triplets.append(neg_triplet)
#
#     return neg_triplets


def incidence_matrix(adj_list):
    '''
    adj_list: List of sparse adjacency matrices
    '''

    rows, cols, dats = [], [], []
    dim = adj_list[0].shape
    for adj in adj_list:
        adjcoo = adj.tocoo()
        rows += adjcoo.row.tolist()
        cols += adjcoo.col.tolist()
        dats += adjcoo.data.tolist()
    row = np.array(rows)
    col = np.array(cols)
    data = np.array(dats)
    return ssp.csc_matrix((data, (row, col)), shape=dim)


def _bfs_relational(adj, roots, max_nodes_per_hop=None):
    """
    BFS for graphs with multiple edge types. Returns list of level sets.
    Each entry in list corresponds to relation specified by adj_list.
    Modified from dgl.contrib.data.knowledge_graph to node accomodate sampling
    """
    visited = set()
    current_lvl = set(roots)

    next_lvl = set()

    while current_lvl:

        for v in current_lvl:
            visited.add(v)

        next_lvl = _get_neighbors(adj, current_lvl)
        next_lvl -= visited  # set difference

        if max_nodes_per_hop and max_nodes_per_hop < len(next_lvl):
            next_lvl = set(random.sample(next_lvl, max_nodes_per_hop))

        yield next_lvl

        current_lvl = set.union(next_lvl)


def _get_neighbors(adj, nodes):
    """Takes a set of nodes and a graph adjacency matrix and returns a set of neighbors.
    Directly copied from dgl.contrib.data.knowledge_graph"""
    sp_nodes = _sp_row_vec_from_idx_list(list(nodes), adj.shape[1])
    sp_neighbors = sp_nodes.dot(adj)
    neighbors = set(ssp.find(sp_neighbors)[1])  # convert to set of indices
    return neighbors


def _sp_row_vec_from_idx_list(idx_list, dim):
    """Create sparse vector of dimensionality dim from a list of indices."""
    shape = (1, dim)
    data = np.ones(len(idx_list))
    row_ind = np.zeros(len(idx_list))
    col_ind = list(idx_list)
    return ssp.csr_matrix((data, (row_ind, col_ind)), shape=shape)


def get_neighbor_nodes(roots, adj, h=1, max_nodes_per_hop=None):
    bfs_generator = _bfs_relational(adj, roots, max_nodes_per_hop)
    lvls = list()
    for _ in range(h):
        try:
            lvls.append(next(bfs_generator))
        except StopIteration:
            pass
    return set().union(*lvls)


def subgraph_extraction_labeling(ind, rel, A_list, h=1, enclosing_sub_graph=False, max_nodes_per_hop=None, node_information=None, max_node_label_value=None):
    # extract the h-hop enclosing subgraphs around link 'ind'
    A_incidence = incidence_matrix(A_list)
    A_incidence += A_incidence.T

    # could pack these two into a function
    root1_nei = get_neighbor_nodes(set([ind[0]]), A_incidence, h, max_nodes_per_hop)
    root2_nei = get_neighbor_nodes(set([ind[1]]), A_incidence, h, max_nodes_per_hop)

    subgraph_nei_nodes_int = root1_nei.intersection(root2_nei)
    subgraph_nei_nodes_un = root1_nei.union(root2_nei)

    # Extract subgraph | Roots being in the front is essential for labelling and the model to work properly.
    # if enclosing_sub_graph:
    enclosing_subgraph_nodes = list(ind) + list(subgraph_nei_nodes_int)
    disclosing_subgraph_nodes = list(ind) + list(subgraph_nei_nodes_un)

    enclosing_subgraph = [adj[enclosing_subgraph_nodes, :][:, enclosing_subgraph_nodes] for adj in A_list]
    disclosing_subgraph = [adj[disclosing_subgraph_nodes, :][:, disclosing_subgraph_nodes] for adj in A_list]

    # labels, enclosing_subgraph_nodes = node_label_new(incidence_matrix(subgraph), max_distance=h)

    enclosing_labels, enclosing_subgraph_nodes_labeled = node_label_new(incidence_matrix(enclosing_subgraph),
                                                                    max_distance=h, enclosing_flag=True)
    disclosing_labels, disclosing_subgraph_nodes_labeled = node_label_new(incidence_matrix(disclosing_subgraph),
                                                                      max_distance=h, enclosing_flag=False)

    pruned_enclosing_subgraph_nodes = np.array(enclosing_subgraph_nodes)[enclosing_subgraph_nodes_labeled].tolist()
    pruned_enclosing_labels = enclosing_labels[enclosing_subgraph_nodes_labeled]

    pruned_disclosing_subgraph_nodes = np.array(disclosing_subgraph_nodes)[disclosing_subgraph_nodes_labeled].tolist()
    pruned_disclosing_labels = disclosing_labels[disclosing_subgraph_nodes_labeled]

    if max_node_label_value is not None:
        pruned_enclosing_labels = np.array(
            [np.minimum(label, max_node_label_value).tolist() for label in pruned_enclosing_labels])
        pruned_disclosing_labels = np.array(
            [np.minimum(label, max_node_label_value).tolist() for label in pruned_disclosing_labels])

    return pruned_enclosing_subgraph_nodes, pruned_enclosing_labels, pruned_disclosing_subgraph_nodes, pruned_disclosing_labels


def remove_nodes(A_incidence, nodes):
    idxs_wo_nodes = list(set(range(A_incidence.shape[1])) - set(nodes))
    return A_incidence[idxs_wo_nodes, :][:, idxs_wo_nodes]


def node_label_new(subgraph, max_distance=1, enclosing_flag=False):
    # an implementation of the proposed double-radius node labeling (DRNd   L)
    roots = [0, 1]
    sgs_single_root = [remove_nodes(subgraph, [root]) for root in roots]
    dist_to_roots = [np.clip(ssp.csgraph.dijkstra(sg, indices=[0], directed=False, unweighted=True, limit=1e6)[:, 1:], 0, 1e7) for r, sg in enumerate(sgs_single_root)]
    dist_to_roots = np.array(list(zip(dist_to_roots[0][0], dist_to_roots[1][0])), dtype=int)

    # dist_to_roots[np.abs(dist_to_roots) > 1e6] = 0
    # dist_to_roots = dist_to_roots + 1
    target_node_labels = np.array([[0, 1], [1, 0]])
    labels = np.concatenate((target_node_labels, dist_to_roots)) if dist_to_roots.size else target_node_labels

    # enclosing_subgraph_nodes = np.where(np.max(labels, axis=1) <= max_distance)[0]

    if enclosing_flag:

        enclosing_subgraph_nodes = np.where(np.max(labels, axis=1) <= max_distance)[0]
    else:
        # enclosing_subgraph_nodes = np.where(np.max(labels, axis=1) < 1e6)[0]
        # process the unconnected node (neg samples)
        indices_dim0, indices_dim1 = np.where(labels == 1e7)

        indices_dim1_convert = indices_dim1 + 1
        indices_dim1_convert[indices_dim1_convert == 2] = 0
        new_indices = [indices_dim0.tolist(), indices_dim1_convert.tolist()]
        ori_indices = [indices_dim0.tolist(), indices_dim1.tolist()]

        # values = labels[new_indices] + 1
        # labels[ori_indices] = values
        values = labels[tuple(new_indices)] + 1
        labels[tuple(ori_indices)] = values
        # process the unconnected node (neg samples)

        # print(labels)
        enclosing_subgraph_nodes = np.where(np.max(labels, axis=1) <= max_distance)[0]

    # print(len(enclosing_subgraph_nodes))
    return labels, enclosing_subgraph_nodes


def ssp_multigraph_to_dgl(graph, n_feats=None):
    """
    Converting ssp multigraph (i.e. list of adjs) to dgl multigraph.
    """

    g_nx = nx.MultiDiGraph()
    g_nx.add_nodes_from(list(range(graph[0].shape[0])))
    # Add edges
    for rel, adj in enumerate(graph):
        # Convert adjacency matrix to tuples for nx0
        nx_triplets = []
        for src, dst in list(zip(adj.tocoo().row, adj.tocoo().col)):
            nx_triplets.append((src, dst, {'type': rel}))
        g_nx.add_edges_from(nx_triplets)

    # make dgl graph
    g_dgl = dgl.DGLGraph(multigraph=True)
    g_dgl.from_networkx(g_nx, edge_attrs=['type'])
    # add node features
    if n_feats is not None:
        g_dgl.ndata['feat'] = torch.tensor(n_feats)

    return g_dgl


def prepare_features(subgraph, nodes, n_labels, max_n_label, n_feats=None):
    # One hot encode the node label feature and concat to n_featsure
    n_nodes = subgraph.number_of_nodes()
    label_feats = np.zeros((n_nodes, max_n_label[0] + 1 + max_n_label[1] + 1))
    label_feats[np.arange(n_nodes), n_labels[:, 0]] = 1
    label_feats[np.arange(n_nodes), max_n_label[0] + 1 + n_labels[:, 1]] = 1
    n_feats = np.concatenate((label_feats, n_feats), axis=1) if n_feats is not None else label_feats
    subgraph.ndata['feat'] = torch.FloatTensor(n_feats)

    head_id = np.argwhere([label[0] == 0 and label[1] == 1 for label in n_labels])
    tail_id = np.argwhere([label[0] == 1 and label[1] == 0 for label in n_labels])
    n_ids = np.zeros(n_nodes)
    n_ids[head_id] = 1  # head
    n_ids[tail_id] = 2  # tail
    subgraph.ndata['id'] = torch.FloatTensor(n_ids)
    subgraph.ndata['eid'] = torch.LongTensor(nodes)

    return subgraph

def prepare_subgraph(dgl_adj_list, nodes, rel, node_labels, max_node_label_value):
    subgraph = dgl.DGLGraph(dgl_adj_list.subgraph(nodes))
    subgraph.edata['type'] = dgl_adj_list.edata['type'][dgl_adj_list.subgraph(nodes).parent_eid]
    subgraph.edata['label'] = torch.tensor(rel * np.ones(subgraph.edata['type'].shape), dtype=torch.long)

    edges_btw_roots = subgraph.edge_id(0, 1, return_array=True)
    rel_link = np.nonzero(subgraph.edata['type'][edges_btw_roots] == rel)

    # if rel_link.squeeze().nelement() == 0:
    #     # subgraph.add_edge(0, 1, {'type': torch.tensor([rel]), 'label': torch.tensor([rel])})
    #     subgraph.add_edge(0, 1)
    #     subgraph.edata['type'][-1] = torch.tensor(rel).type(torch.LongTensor)
    #     subgraph.edata['label'][-1] = torch.tensor(rel).type(torch.LongTensor)

    if rel_link.squeeze().nelement() == 0:
        subgraph.add_edge(0, 1)
        subgraph.edata['type'][-1] = torch.tensor(rel).type(torch.LongTensor)
        subgraph.edata['label'][-1] = torch.tensor(rel).type(torch.LongTensor)
        e_ids = np.zeros(subgraph.number_of_edges())
        e_ids[-1] = 1  # target edge
    else:
        e_ids = np.zeros(subgraph.number_of_edges())
        e_ids[edges_btw_roots] = 1  # target edge

    subgraph.edata['id'] = torch.FloatTensor(e_ids)

    subgraph = prepare_features(subgraph, nodes, node_labels, max_node_label_value)
    return subgraph

def get_subgraphs(all_links, adj_list, dgl_adj_list, max_node_label_value):
    # dgl_adj_list = ssp_multigraph_to_dgl(adj_list)

    en_subgraphs = []
    dis_subgraphs = []
    r_labels = []

    for link in all_links:
        head, tail, rel = link[0], link[1], link[2]
        en_nodes, en_node_labels, dis_nodes, dis_node_labels = subgraph_extraction_labeling((head, tail), rel, adj_list, h=params_.hop, enclosing_sub_graph=params.enclosing_sub_graph, max_node_label_value=max_node_label_value)

        en_subgraph = prepare_subgraph(dgl_adj_list, en_nodes, rel, en_node_labels, max_node_label_value)
        dis_subgraph = prepare_subgraph(dgl_adj_list, dis_nodes, rel, dis_node_labels, max_node_label_value)


        en_subgraphs.append(en_subgraph)
        dis_subgraphs.append(dis_subgraph)
        r_labels.append(rel)

    batched_en_graph = dgl.batch(en_subgraphs)
    batched_dis_graph = dgl.batch(dis_subgraphs)
    r_labels = torch.LongTensor(r_labels)

    return (batched_en_graph, batched_dis_graph, r_labels)


def get_rank(neg_links):
    head_neg_links = neg_links['head'][0]
    head_target_id = neg_links['head'][1]

    if head_target_id != 10000:
        data = get_subgraphs(head_neg_links, adj_list_, dgl_adj_list_, params.max_label_value)
        head_scores = model_(data).squeeze(1)
        if params.ont:
            head_scores += ont_scorer_(data, ent2types_)
        head_scores = head_scores.detach().numpy()
        head_rank = np.argwhere(np.argsort(head_scores)[::-1] == head_target_id) + 1
    else:
        head_scores = np.array([])
        head_rank = 10000

    tail_neg_links = neg_links['tail'][0]
    tail_target_id = neg_links['tail'][1]

    if tail_target_id != 10000:
        data = get_subgraphs(tail_neg_links, adj_list_, dgl_adj_list_, params.max_label_value)
        tail_scores = model_(data).squeeze(1)
        if params.ont:
            tail_scores += ont_scorer_(data, ent2types_)
        tail_scores = tail_scores.detach().numpy()
        tail_rank = np.argwhere(np.argsort(tail_scores)[::-1] == tail_target_id) + 1
    else:
        tail_scores = np.array([])
        tail_rank = 10000

    return head_scores, head_rank, tail_scores, tail_rank


def save_negative_triples_to_file(neg_triplets, id2entity, id2relation):

    with open(os.path.join('../data', params.dataset, 'ranking_head.txt'), "w") as f:
        for neg_triplet in neg_triplets:
            for s, o, r in neg_triplet['head'][0]:
                f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o]]) + '\n')

    with open(os.path.join('../data', params.dataset, 'ranking_tail.txt'), "w") as f:
        for neg_triplet in neg_triplets:
            for s, o, r in neg_triplet['tail'][0]:
                f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o]]) + '\n')


def save_score_to_file(neg_triplets, all_head_scores, all_tail_scores, id2entity, id2relation):
    with open(os.path.join('../data', params.dataset, f'{params.ont}_ranking_head_predictions.txt'), "w") as f:
        for i, neg_triplet in enumerate(neg_triplets):
            for [s, o, r], head_score in zip(neg_triplet['head'][0], all_head_scores[50 * i:50 * (i + 1)]):
                f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o], str(head_score)]) + '\n')

    with open(os.path.join('../data', params.dataset, f'{params.ont}_ranking_tail_predictions.txt'), "w") as f:
        for i, neg_triplet in enumerate(neg_triplets):
            for [s, o, r], tail_score in zip(neg_triplet['tail'][0], all_tail_scores[50 * i:50 * (i + 1)]):
                f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o], str(tail_score)]) + '\n')


def read_neg_samples(dataset, entity2id, relation2id, num_neg=50):
    neg_triplets = []
    # for i in range(num_neg):
    #     neg_triplet=
    #     neg_triplets.append(neg_triplet)

    with open(os.path.join('../data', dataset, 'ranking_head.txt'), "r") as f1:
        cur = -1
        while True:
            line = f1.readline()
            if line == '':
                break

            neg_triplets.append({'head': [[], 0], 'tail': [[], 0]})
            cur += 1
            s, r, o = line.strip().split()
            neg_triplets[cur]['head'][0].append([entity2id[s], entity2id[o], relation2id[r]])

            for i in range(num_neg - 1):
                line = f1.readline()
                s, r, o = line.strip().split()
                neg_triplets[cur]['head'][0].append([entity2id[s], entity2id[o], relation2id[r]])

    with open(os.path.join('../data', dataset, 'ranking_tail.txt'), "r") as f2:
        cur = -1
        while True:
            line = f2.readline()
            if line == '':
                break

            cur += 1
            s, r, o = line.strip().split()
            neg_triplets[cur]['tail'][0].append([entity2id[s], entity2id[o], relation2id[r]])

            for i in range(num_neg - 1):
                line = f2.readline()
                s, r, o = line.strip().split()
                neg_triplets[cur]['tail'][0].append([entity2id[s], entity2id[o], relation2id[r]])

    return neg_triplets


def main(params):
    model = torch.load(params.model_path, map_location='cpu')

    model.params.gpu = -1
    model.device = torch.device('cpu')
    params.max_label_value = np.array([2, 2])
    adj_list, dgl_adj_list, triplets, entity2id, relation2id, id2entity, id2relation = process_files(params.file_paths, model.relation2id, add_traspose_rels=False)

    ent2types=None
    ont_scorer = None
    if params.ont or params.type_graph:
        dataset_prefix = params.dataset.rsplit('_', maxsplit=2)[0]

        type2id_path = os.path.join(
                                    f'../data/{params.dataset.rsplit("_", maxsplit=1)[0]}/type2id.json')
        with open(type2id_path) as f:
            type2id = json.load(f)
        params.num_types = len(type2id)
        ent2types = get_ent_types(f'../types/{dataset_prefix}_ent2type_top.txt',
                                  entity2id, type2id)
        ent2types = np.array(ent2types, dtype=object)

    if params.ont:
        ont_scorer = torch.load(params.ont_scorer_path, map_location='cpu')
    # id2type = {type2id[t]: t for t in type2id}
    all_mrr = []
    all_hit1 = []
    all_hit5 = []
    all_hit10 = []
    for r in range(1, params.runs + 1):
        # if params.mode == 'sample':
        #     neg_triplets = get_neg_samples_replacing_head_tail(triplets['links'], adj_list)
        #     save_negative_triples_to_file(neg_triplets, id2entity, id2relation)
        # elif params.mode == 'all':
        #     neg_triplets = get_neg_samples_replacing_head_tail_all(triplets['links'], adj_list)
        neg_triplets = read_neg_samples(params.dataset, entity2id, relation2id)

        ranks = []
        all_head_scores = []
        all_tail_scores = []
        with mp.Pool(processes=None, initializer=intialize_worker, initargs=(model, adj_list, dgl_adj_list, id2entity, params, ont_scorer,
                                                                             ent2types)) as p:
            for head_scores, head_rank, tail_scores, tail_rank in tqdm(p.imap(get_rank, neg_triplets), total=len(neg_triplets)):
                ranks.append(head_rank)
                ranks.append(tail_rank)

                all_head_scores += head_scores.tolist()
                all_tail_scores += tail_scores.tolist()

        save_score_to_file(neg_triplets, all_head_scores, all_tail_scores, id2entity, id2relation)

        isHit1List = [x for x in ranks if x <= 1]
        isHit5List = [x for x in ranks if x <= 5]
        isHit10List = [x for x in ranks if x <= 10]
        hits_1 = len(isHit1List) / len(ranks)
        hits_5 = len(isHit5List) / len(ranks)
        hits_10 = len(isHit10List) / len(ranks)

        mrr = np.mean(1 / np.array(ranks))

        all_mrr.append(mrr)
        all_hit1.append(hits_1)
        all_hit5.append(hits_5)
        all_hit10.append(hits_10)

        # print(f'MRR | Hits@1 | Hits@5 | Hits@10 : {mrr} | {hits_1} | {hits_5} | {hits_10}')
        print(f'MRR: {mrr: .4f}')
        print(f'Hits@1  : {hits_1: .4f}')
        print(f'Hits@5  : {hits_5: .4f}')
        print(f'Hits@10 : {hits_10: .4f}')

    avg_mrr = np.mean(all_mrr)
    avg_hit1 = np.mean(all_hit1)
    avg_hit5 = np.mean(all_hit5)
    avg_hit10 = np.mean(all_hit10)

    print('\nAvg test Set Performance ')
    print(f'MRR: {avg_mrr: .4f}')
    print(f'Hits@1  : {avg_hit1: .4f}')
    print(f'Hits@5  : {avg_hit5: .4f}')
    print(f'Hits@10 : {avg_hit10: .4f}')

if __name__ == '__main__':

    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(description='Testing script for hits@10')

    # Experiment setup params
    parser.add_argument("--expri_name", "-e", type=str, default="fb_v2_margin_loss",
                        help="Experiment name. Log file with this name will be created")
    parser.add_argument("--dataset", "-d", type=str, default="FB237_v2", help="Path to dataset")
    parser.add_argument("--mode", "-m", type=str, default="sample", choices=["sample", "all", "ruleN"],
                        help="Negative sampling mode")
    parser.add_argument('--enclosing_sub_graph', '-en', type=bool, default=True, help='whether to only consider enclosing subgraph')
    parser.add_argument("--hop", type=int, default=2, help="How many hops to go while eextracting subgraphs?")
    parser.add_argument('--seed', default=41504, type=int, help='Seed for randomization')
    parser.add_argument('--target2nei_atten', action='store_true', help='apply target-aware attention for 2-hop neighbors')
    parser.add_argument('--conc', action='store_true', help='apply target-aware attention for 2-hop neighbors')
    parser.add_argument('--ablation', type=int, default=0, help='0,1 correspond to base, NE')
    parser.add_argument("--runs", type=int, default=5, help="How many runs to perform for mean and std?")

    parser.add_argument('--ont', action='store_true')
    parser.add_argument('--type_graph', '-tg', action='store_true')
    parser.add_argument("--type_emb_dim", "-t_dim", type=int, default=32,
                    help="Type embedding size")

    params = parser.parse_args()

    params.file_paths = {
        'graph': os.path.join('../data', params.dataset, 'train.txt'),
        'links': os.path.join('../data', params.dataset, 'test.txt')
    }

    params.model_path = os.path.join('expri_save_models', params.expri_name, 'best_graph_classifier.pth')
    params.ont_scorer_path = os.path.join('expri_save_models', params.expri_name, 'best_ont_scorer.pth')

    print('============ Params ============')
    print('\n'.join('%s: %s' % (k, str(v)) for k, v
                          in sorted(dict(vars(params)).items())))
    print('============================================')

    dataset_prefix = params.dataset.rsplit('_', maxsplit=2)[0]

    main(params)