helper.py

import sys
import os
import random
import json
import uuid
import time
import argparse
import logging
import logging.config
import numpy as np
from random import randint
from collections import defaultdict as ddict, Counter
from ordered_set import OrderedSet
from pprint import pprint

# PyTorch related imports
import torch
from torch.nn import functional as F
from torch.nn.parameter import Parameter
from torch.nn.init import xavier_normal_, xavier_uniform_
from torch.nn import Parameter as Param
from torch.utils.data import DataLoader

np.set_printoptions(precision=4)


def set_gpu(gpus):
    """
    Sets the GPU to be used for the run

    Parameters
    ----------
    gpus:           List of GPUs to be used for the run

    Returns
    -------

    """
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = gpus


def get_logger(name, log_dir, config_dir):
    """
    Creates a logger object

    Parameters
    ----------
    name:           Name of the logger file
    log_dir:        Directory where logger file needs to be stored
    config_dir:     Directory from where log_config.json needs to be read

    Returns
    -------
    A logger object which writes to both file and stdout

    """
    config_dict = json.load(open(config_dir + 'log_config.json'))
    config_dict['handlers']['file_handler']['filename'] = log_dir + \
        name.replace('/', '-')
    logging.config.dictConfig(config_dict)
    logger = logging.getLogger(name)

    std_out_format = '%(asctime)s - [%(levelname)s] - %(message)s'
    consoleHandler = logging.StreamHandler(sys.stdout)
    consoleHandler.setFormatter(logging.Formatter(std_out_format))
    logger.addHandler(consoleHandler)

    return logger


def get_combined_results(left_results, right_results):
    """
    Computes the average based on head and tail prediction results

    Parameters
    ----------
    left_results:   Head prediction results
    right_results: 	Left prediction results

    Returns
    -------
    Average prediction results

    """

    results = {}
    count = float(left_results['count'])

    results['left_mr'] = round(left_results['mr'] / count, 5)
    results['left_mrr'] = round(left_results['mrr']/count, 5)
    results['right_mr'] = round(right_results['mr'] / count, 5)
    results['right_mrr'] = round(right_results['mrr']/count, 5)
    results['mr'] = round(
        (left_results['mr'] + right_results['mr']) / (2*count), 5)
    results['mrr'] = round(
        (left_results['mrr'] + right_results['mrr'])/(2*count), 5)

    for k in range(10):
        results['left_hits@{}'.format(k+1)] = round(
            left_results['hits@{}'.format(k+1)]/count, 5)
        results['right_hits@{}'.format(k+1)] = round(
            right_results['hits@{}'.format(k+1)]/count, 5)
        results['hits@{}'.format(k+1)] = round((left_results['hits@{}'.format(
            k+1)] + right_results['hits@{}'.format(k+1)])/(2*count), 5)
    return results