memory_transfer.py

# GRAPH_AGG3

from torch.nn import Linear
from torch_geometric.data import TemporalData
from torch_geometric.loader import TemporalDataLoader
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.model_selection import train_test_split
import torch.nn.functional as F
import torch
from torch import Tensor
import pandas as pd
import numpy as np
from torch_geometric.nn import Node2Vec
from collections import Counter
from sklearn.metrics.pairwise import cosine_similarity
from torch_geometric.utils import negative_sampling
from sklearn.metrics import roc_auc_score
import sklearn

import pickle
import os
from tqdm import tqdm
from numpy.linalg import norm
print(torch.__version__)
import random
import json
import argparse

from utils import get_global_cats,seed_everything,save_dataframe,returnEdgeIndexes,save_dataframe_wiki, save_dataframe_yelp
from model import GraphAggregator

SEED = 2019
seed_everything(SEED)

import pathlib
pwd = str(pathlib.Path(__file__).parent.resolve())

print("Current working directory", pwd)
import sys
sys.path.append(pwd)

def create_arg_parser():
    parser = argparse.ArgumentParser(description="Your description here.")

    parser.add_argument('--dataset', type=str, default="4square",choices=["4square","wiki","yelp"], help='dataset name')
    parser.add_argument('--trainfrac', type=float, default=0.10, help='training data fraction')
    parser.add_argument('--src', type=str, help='Source CITY')
    parser.add_argument('--tgt', type=str, help='Target CITY')
    parser.add_argument('--gpu', type=int, choices=[0, 1, 2, 3], default=0, help='GPU #')
    parser.add_argument('--layers', type=int, default=1, choices=[0, 1, 2, 3], help='number of GAT layers')
    parser.add_argument('--mode', type=str, default="GAT2",choices=["GAT2"], help='Type of Mapping')
    return parser

parser = create_arg_parser()
parser = parser.parse_args()
device = parser.gpu
src_city = parser.src
tgt_city = parser.tgt
data_frac = parser.trainfrac
num_layers_gat = parser.layers
dataset = parser.dataset
mode = parser.mode

print("Src city", src_city)
print("Tgt city", tgt_city)
print("Data frac", data_frac)

if dataset=="wiki":
    df_city = save_dataframe_wiki(tgt_city,source=False,useSaved=True)
elif dataset=="yelp":
    df_city = save_dataframe_yelp(tgt_city,useSaved=True)
    df_city['category'] = df_city['categories_list']
else:
    df_city = save_dataframe(tgt_city,source=False,useSaved=True)
    
item_df = df_city[['i','category']]  ### To keep all businesses
print(df_city.shape)
len_ = int(df_city.shape[0]*data_frac)
print(len_)
df_city = df_city.head(len_)
print("After filtering", df_city.shape)
min_dst_idx, max_dst_idx = int(min(df_city['i'])), int(max(df_city['i']))
min_src_idx, max_src_idx = int(min(df_city['u'])), int(max(df_city['u']))

print("min_src_idx",min_src_idx,"max_src_idx",max_src_idx,"min_dst_idx",min_dst_idx,"max_dst_idx",max_dst_idx)
b2u_edge_idx, u2b_edge_idx, b2c_edge_idx, c2b_edge_idx, [num_cat, num_poi, num_usr],b2u_edge_wt,u2b_edge_wt,usrStart,busStart = returnEdgeIndexes(df_city,dataset=dataset,item_cat_df=item_df)
print("User start" , usrStart, " Business start", busStart)


print("NUM CAT:", num_cat)
cat_feat = torch.eye(num_cat).to(device)
poi_feat = torch.zeros(num_poi, num_cat).to(device)
usr_feat = torch.zeros(num_usr, num_cat).to(device)
print(cat_feat.shape,poi_feat.shape,usr_feat.shape)

b2c_edge_idx = b2c_edge_idx.to(device)
c2b_edge_idx = c2b_edge_idx.to(device)
b2u_edge_idx = b2u_edge_idx.to(device)
u2b_edge_idx = u2b_edge_idx.to(device)
b2u_edge_wt = b2u_edge_wt.to(device)
u2b_edge_wt = u2b_edge_wt.to(device)


model = GraphAggregator(input_channels=poi_feat.shape[1],hidden_channels=128,device=device,
                        heads=2,dropout=0,num_layers=num_layers_gat).to(device)

if dataset == "wiki":
    model.load_state_dict(torch.load(pwd+'/models/pt_wiki_GAT_layer_1.pth'))
elif dataset == "yelp":
    model.load_state_dict(torch.load(pwd+"/models/pt_yelp_GAT_layer_1.pth"))
elif dataset == "4square":
    model.load_state_dict(torch.load(pwd+'/models/pt_4square_GAT_layer_1.pth')) 

model.eval()
model.device= device
print("Universal gat model loaded")

print("Loading src city embeddings and metadata")
usr_z_src = torch.load(pwd+"/models/"+src_city+"_usr.pth")
poi_z_src = torch.load(pwd+"/models/"+src_city+"_item.pth")
meta_data_src = pickle.load(open(pwd+"/models/"+src_city+"_meta.pkl","rb"))
usr_start_src = meta_data_src['usr_start']
bus_start_src= meta_data_src['bus_start']
df_city_src= meta_data_src['df_city']
print("SRC EMBEDDINGS:", usr_z_src.shape,poi_z_src.shape)
print("Loading src city memory")
src_memory= torch.load(pwd+"/models/"+src_city+"_tgn_memory.pth")
print("Source embedding sizes", usr_z_src.shape,poi_z_src.shape,src_memory.shape)

print("Computing embedding at target node")
usr_z,poi_z,cat_z = model(cat_feat, poi_feat, usr_feat,b2u_edge_idx,
                        u2b_edge_idx, b2c_edge_idx, c2b_edge_idx,
                        b2u_edge_wt,u2b_edge_wt)
print(usr_z.shape,poi_z.shape)

u2b_edge_idx_neg = negative_sampling(u2b_edge_idx,num_nodes=(num_usr,num_poi),method='dense') 
u2b_edge_idx_combined = torch.cat([u2b_edge_idx, u2b_edge_idx_neg],dim=-1,).to(device)
u2b_edge_idx_label_combined = torch.cat([torch.ones(u2b_edge_idx.shape[1],dtype=torch.int),torch.zeros(u2b_edge_idx_neg.shape[1])], dim=0).to(device) 

out = model.decode(usr_z,poi_z, u2b_edge_idx_combined).view(-1)
approx_auc = roc_auc_score(u2b_edge_idx_label_combined.cpu().numpy(), out.sigmoid().cpu().detach().numpy())

print("This is the auc on target dataset", approx_auc)


print("Initializing memory for transfer city")
memory = torch.zeros((max_dst_idx+1,src_memory.shape[1]))
print(memory.shape)


#### First getting user memory #####
tgt_src = sklearn.metrics.pairwise.cosine_similarity(usr_z.cpu().detach().numpy(),usr_z_src.cpu().detach().numpy())
# print("TGT SRC- ", tgt_src.shape)

tgt_src_topk = np.argsort((-1)*tgt_src)[:,0]
# print("TGT SRC- ", tgt_src_topk.shape)


for i in range(0, tgt_src_topk.shape[0]):
    tgt = i+usrStart
    src = tgt_src_topk[i]+usr_start_src    
    memory[tgt] = src_memory[src]
    
####  Business memory #####
tgt_src = sklearn.metrics.pairwise.cosine_similarity(poi_z.cpu().detach().numpy(),poi_z_src.cpu().detach().numpy())
# print("bs SRC- ", tgt_src.shape)

tgt_src_topk = np.argsort((-1)*tgt_src)[:,0]
# print("bs SRC- ", tgt_src_topk.shape)

for i in range(0, tgt_src_topk.shape[0]):
    tgt = i+busStart
    src = tgt_src_topk[i]+bus_start_src
    # print(bus_start_src,tgt_src_topk[i])    
    memory[tgt] = src_memory[src]  
    

if mode == "GAT2":
    print("Saving pre-trained memory for {}".format(tgt_city)," on ",pwd+"/models/"+src_city + "_" + tgt_city+"_tgn_pt_memory.pth" )
    torch.save(memory,pwd+"/models/"+src_city + "_" + tgt_city+"_tgn_pt_memory.pth")