utils.py

#!/usr/bin/env python
# coding: utf-8

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import glob,pickle,joblib,os,sys
import seaborn as sns
import zipfile,itertools
from collections import Counter
from tqdm import tqdm
import torch
import torch.nn as nn
from pprint import pprint


def save_pickle(dataset,file_path):
    with open(file_path,'wb') as file:
        pickle.dump(dataset,file)

def load(file_name):
    with open(file_name, "rb") as f:
        return pickle.load(f)
        
def save_checkpoint(state, save_path: str):
    torch.save(state, save_path)

def load_checkpoint(ckpt_path):
    ckpt = torch.load(ckpt_path)
    return ckpt

def eval_metric(predictions_batch,answers_batch,k):

    metrics = {}
    #k=10  # 5.10.20.30
    recall = 0.0
    mrr = 0.0
    ndcg = 0.0

    for i in range(len(answers_batch)):
        pred_list=predictions_batch[i]
        gt=answers_batch[i]
        for j in range(k):
            if pred_list[j] == gt[0]:
                recall = recall + 1.0
                mrr = mrr + 1.0 / (j + 1.0)
                ndcg = ndcg + (1 / np.log2(j + 2))
    metrics['HR@%d' % k] = recall / len(answers_batch)
    metrics['NDCG@%d' % k] = ndcg / len(answers_batch)
    metrics['MRR@%d' % k] = mrr /len(answers_batch)   
    return metrics


def metric_to_list(metrics5, metrics10, metrics20):
    metrics = [metrics5, metrics10, metrics20]
    hr_keys = ['HR@5', 'HR@10', 'HR@20']
    ndcg_keys = ['NDCG@5', 'NDCG@10', 'NDCG@20']
    mrr_keys = ['MRR@5','MRR@10','MRR@20']
    result = [metric[key] for metric, key in zip(metrics * 3, hr_keys + ndcg_keys +mrr_keys)]
    return result

def eval_metric_MM(predictions_batch, answers_batch, scenario_push_batch, source_ads_batch, k):
    preds_push_ads, answers_push_ads = [], []
    preds_push_news, answers_push_news = [], []
    preds_browse_ads, answers_browse_ads = [], []
    preds_browse_news, answers_browse_news = [], []
    
    for prediction, answer, scenario_push, source_ads in zip(predictions_batch, answers_batch, scenario_push_batch, source_ads_batch):
        if scenario_push[0] == 1:
            if source_ads[0] == 1:
                preds_push_ads.append(prediction)
                answers_push_ads.append(answer)
            else:
                preds_push_news.append(prediction)
                answers_push_news.append(answer)
        else:
            if source_ads[0] == 1:
                preds_browse_ads.append(prediction)
                answers_browse_ads.append(answer)
            else:
                preds_browse_news.append(prediction)
                answers_browse_news.append(answer)
    
    metrics = {
        'push_ads': eval_metric(preds_push_ads, answers_push_ads, k),
        'push_news': eval_metric(preds_push_news, answers_push_news, k),
        'browse_ads': eval_metric(preds_browse_ads, answers_browse_ads, k),
        'browse_news': eval_metric(preds_browse_news, answers_browse_news, k)
    }
    return metrics


def progress_bar(count, total, prefix='', suffix=''):
    bar_len = 60
    filled_len = int(round(bar_len * count / float(total)))
    bar = '=' * filled_len + '-' * (bar_len - filled_len)
    sys.stdout.write(prefix + '[%s]-Step [%s/%s]-%s\r' % (bar, count, total, suffix))
    sys.stdout.flush()
    if count == total:
      print("\n")


def experiment_record(*args):
    with open("ckpt/p1_log.txt", 'a') as f:
        print("""=======================================================
UUID:       {}
Time:       {}
Batch size: {}
Lr:         {} 
milestones: {}
use_adam:   {}
Result:
    Epoch:      {}
    Valid ACC:  {}
=======================================================""".format(*args), file=f)