Spaces:
Sleeping
Sleeping
File size: 3,336 Bytes
06d9388 e4f5e4c 06d9388 e4f5e4c 06d9388 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import umap.umap_ as umap
import plotly.express as px
import pandas as pd
import random
import viz_utils
import torch
import torch
import torch.nn.functional as F
from torch.nn import Linear
import torch_geometric.transforms as T
from torch_geometric.nn import SAGEConv, to_hetero
from torch_geometric.transforms import RandomLinkSplit, ToUndirected
from sentence_transformers import SentenceTransformer
from torch_geometric.data import HeteroData
import yaml
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data = torch.load("./PyGdata.pt", map_location=device('cpu'))
movies_df = pd.read_csv("./sampled_movie_dataset/movies_metadata.csv")
class GNNEncoder(torch.nn.Module):
def __init__(self, hidden_channels, out_channels):
super().__init__()
# these convolutions have been replicated to match the number of edge types
self.conv1 = SAGEConv((-1, -1), hidden_channels)
self.conv2 = SAGEConv((-1, -1), out_channels)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index).relu()
x = self.conv2(x, edge_index)
return x
class EdgeDecoder(torch.nn.Module):
def __init__(self, hidden_channels):
super().__init__()
self.lin1 = Linear(2 * hidden_channels, hidden_channels)
self.lin2 = Linear(hidden_channels, 1)
def forward(self, z_dict, edge_label_index):
row, col = edge_label_index
# concat user and movie embeddings
z = torch.cat([z_dict['user'][row], z_dict['movie'][col]], dim=-1)
# concatenated embeddings passed to linear layer
z = self.lin1(z).relu()
z = self.lin2(z)
return z.view(-1)
class Model(torch.nn.Module):
def __init__(self, hidden_channels):
super().__init__()
self.encoder = GNNEncoder(hidden_channels, hidden_channels)
self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
self.decoder = EdgeDecoder(hidden_channels)
def forward(self, x_dict, edge_index_dict, edge_label_index):
# z_dict contains dictionary of movie and user embeddings returned from GraphSage
z_dict = self.encoder(x_dict, edge_index_dict)
return self.decoder(z_dict, edge_label_index)
model = Model(hidden_channels=32).to(device)
model2 = Model(hidden_channels=32).to(device)
model.load_state_dict(torch.load("PyGTrainedModelState.pt"), map_location=device('cpu'))
model.eval()
total_users = data['user'].num_nodes
total_movies = data['movie'].num_nodes
print("total users =", total_users)
print("total movies =", total_movies)
with torch.no_grad():
a = model.encoder(data.x_dict,data.edge_index_dict)
user = pd.DataFrame(a['user'].detach().cpu())
movie = pd.DataFrame(a['movie'].detach().cpu())
embedding_df = pd.concat([user, movie], axis=0)
movie_index = 20
title = movies_df.iloc[movie_index]['title']
print(title)
fig_umap = viz_utils.visualize_embeddings_umap(embedding_df)
viz_utils.save_visualization(fig_umap, "./Visualizations/umap_visualization")
fig_tsne = viz_utils.visualize_embeddings_tsne(embedding_df)
viz_utils.save_visualization(fig_tsne, "./Visualizations/tsne_visualization")
fig_pca = viz_utils.visualize_embeddings_pca(embedding_df)
viz_utils.save_visualization(fig_pca, "./Visualizations/pca_visualization")
|