logic.py

import pandas as pd
from arango import ArangoClient
from tqdm import tqdm
import numpy as np
import itertools
import requests
import sys
import oasis
from arango import ArangoClient

import torch
import torch.nn.functional as F
from torch.nn import Linear
from arango import ArangoClient
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

#-------------------------------------------------------------------------------------------
# Functions

# performs user and movie mappings
def node_mappings(path, index_col):
    df = pd.read_csv(path, index_col=index_col)
    mapping = {index: i for i, index in enumerate(df.index.unique())}

    return mapping


def convert_int(x):
    try:
        return int(x)
    except:
        return np.nan


def remove_movies():
    '''
    # Remove ids which dont have meta data information
    '''
    no_metadata = []
    for idx in range(len(m_id)):
        tmdb_id = id_map.loc[id_map['movieId'] == m_id[idx]]
  
        if tmdb_id.size == 0:
            no_metadata.append(m_id[idx])
            #print('No Meta data information at:', m_id[idx])
    return no_metadata

def populate_user_collection(total_users):
    batch = []
    BATCH_SIZE = 50
    batch_idx = 1
    index = 0
    user_ids = list(user_mapping.keys())
    user_collection = movie_rec_db["Users"]
    for idx in tqdm(range(total_users)):
        insert_doc = {}

        insert_doc["_id"] = "Users/" + str(user_mapping[user_ids[idx]])
        insert_doc["original_id"] = str(user_ids[idx])
        
        batch.append(insert_doc)
        index +=1
        last_record = (idx == (total_users - 1))
        if index % BATCH_SIZE == 0:
            #print("Inserting batch %d" % (batch_idx))
            batch_idx += 1
            user_collection.import_bulk(batch)
            batch = []   
        if last_record and len(batch) > 0:
            print("Inserting batch the last batch!")
            user_collection.import_bulk(batch)

def create_ratings_graph(user_id, movie_id, ratings):
    batch = []
    BATCH_SIZE = 100
    batch_idx = 1
    index = 0
    edge_collection = movie_rec_db["Ratings"]
    for idx in tqdm(range(ratings.shape[0])):
        
        # removing edges (movies) with no metatdata
        if movie_id[idx] in no_metadata:
            print('Removing edges with no metadata', movie_id[idx])
            
        else:
            insert_doc = {}
            insert_doc = {"_from":  ("Users" + "/" + str(user_mapping[user_id[idx]])),
                          "_to":    ("Movie" + "/" + str(movie_mappings[movie_id[idx]])),
                          "_rating": float(ratings[idx])}

            batch.append(insert_doc)
            index += 1
            last_record = (idx == (ratings.shape[0] - 1))

            if index % BATCH_SIZE == 0:
                #print("Inserting batch %d" % (batch_idx))
                batch_idx += 1
                edge_collection.import_bulk(batch)
                batch = []
            if last_record and len(batch) > 0:
                print("Inserting batch the last batch!")
                edge_collection.import_bulk(batch)

def create_pyg_edges(rating_docs):
    src = []
    dst = []
    ratings = []
    for doc in rating_docs:
        _from = int(doc['_from'].split('/')[1])
        _to   = int(doc['_to'].split('/')[1])
         
        src.append(_from)
        dst.append(_to)
        ratings.append(int(doc['_rating']))
        
    edge_index = torch.tensor([src, dst])
    edge_attr = torch.tensor(ratings)

    return edge_index, edge_attr 

def SequenceEncoder(movie_docs , model_name=None):
    movie_titles = [doc['movie_title'] for doc in movie_docs]
    model = SentenceTransformer(model_name, device=device)
    title_embeddings = model.encode(movie_titles, show_progress_bar=True,
                              convert_to_tensor=True, device=device)
    
    return title_embeddings

def GenresEncoder(movie_docs):
    gen = []
    #sep = '|'
    for doc in movie_docs:
        gen.append(doc['genres'])
        #genre = doc['movie_genres']
        #gen.append(genre.split(sep))
    
    # getting unique genres
    unique_gen = set(list(itertools.chain(*gen)))
    print("Number of unqiue genres we have:", unique_gen)
    
    mapping = {g: i for i, g in enumerate(unique_gen)}
    x = torch.zeros(len(gen), len(mapping))
    for i, m_gen in enumerate(gen):
        for genre in m_gen:
            x[i, mapping[genre]] = 1
    return x.to(device)

def weighted_mse_loss(pred, target, weight=None):
    weight = 1. if weight is None else weight[target].to(pred.dtype)
    return (weight * (pred - target.to(pred.dtype)).pow(2)).mean()

@torch.no_grad()
def test(data):
    model.eval()
    pred = model(data.x_dict, data.edge_index_dict,
                 data['user', 'movie'].edge_label_index)
    pred = pred.clamp(min=0, max=5)
    target = data['user', 'movie'].edge_label.float()
    rmse = F.mse_loss(pred, target).sqrt()
    return float(rmse)

def train():
    model.train()
    optimizer.zero_grad()
    pred = model(train_data.x_dict, train_data.edge_index_dict,
                 train_data['user', 'movie'].edge_label_index)
    target = train_data['user', 'movie'].edge_label
    loss = weighted_mse_loss(pred, target, weight)
    loss.backward()
    optimizer.step()
    return float(loss)


#-------------------------------------------------------------------------------------------
# SAGE model
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)
    
#-------------------------------------------------------------------------------------------
def make_graph():
    global movie_mappings, user_mapping, ratings_df, m_id, id_map, sampled_md
    metadata_path = './sampled_movie_dataset/movies_metadata.csv'
    df = pd.read_csv(metadata_path)
    df = df.drop([19730, 29503, 35587])
    df['id'] = df['id'].astype('int')

    links_small = pd.read_csv('./sampled_movie_dataset/links_small.csv')
    links_small = links_small[links_small['tmdbId'].notnull()]['tmdbId'].astype('int') # selecting tmdbId coloumn from links_small file

    sampled_md = df[df['id'].isin(links_small)]
    sampled_md['tagline'] = sampled_md['tagline'].fillna('')
    sampled_md['description'] = sampled_md['overview'] + sampled_md['tagline']
    sampled_md['description'] = sampled_md['description'].fillna('')
    sampled_md = sampled_md.reset_index()

    indices = pd.Series(sampled_md.index, index=sampled_md['title'])
    ind_gen = pd.Series(sampled_md.index, index=sampled_md['genres'])

    ratings_path = './sampled_movie_dataset/ratings_small.csv'
    ratings_df = pd.read_csv(ratings_path)
    m_id = ratings_df['movieId'].tolist()
    m_id = list(dict.fromkeys(m_id))

    user_mapping = node_mappings(ratings_path, index_col='userId')
    movie_mapping = node_mappings(ratings_path, index_col='movieId')

    id_map = pd.read_csv('./sampled_movie_dataset/links_small.csv')[['movieId', 'tmdbId']]
    id_map['tmdbId'] = id_map['tmdbId'].apply(convert_int)
    id_map.columns = ['movieId', 'id']
    id_map = id_map.merge(sampled_md[['title', 'id']], on='id').set_index('title') # tmbdid is same (of links_small) as of id in sampled_md
    indices_map = id_map.set_index('id')

    global no_metadata
    no_metadata = remove_movies()

    ## remove ids which dont have meta data information
    for element in no_metadata:
        if element in m_id:
            print("ids with no metadata information:",element)
            m_id.remove(element)

    # create new movie_mapping dict with only m_ids having metadata information
    movie_mappings = {}
    for idx, m in enumerate(m_id):
        movie_mappings[m] = idx

    

    return movie_mappings, user_mapping, ratings_df, m_id, id_map, sampled_md

    
def login_ArangoDB():
    # get temporary credentials for ArangoDB on cloud
    login = oasis.getTempCredentials(tutorialName="MovieRecommendations", credentialProvider="https://tutorials.arangodb.cloud:8529/_db/_system/tutorialDB/tutorialDB")
    # url to access the ArangoDB Web UI
    url = "https://"+login["hostname"]+":"+str(login["port"])
    username = "Username: " + login["username"]
    password = "Password: " + login["password"]
    dbname = "Database: " + login["dbName"]
    return login,url,username,password,dbname

def create_smart_graph():
  # defining graph schema

  # create a new graph called movie_rating_graph in the temp database if it does not already exist.
  if not movie_rec_db.has_graph("movie_rating_graph"):
      movie_rec_db.create_graph('movie_rating_graph', smart=True)

  # This returns and API wrapper for the above created graphs
  movie_rating_graph = movie_rec_db.graph("movie_rating_graph")

  # Create a new vertex collection named "Users" if it does not exist.
  if not movie_rating_graph.has_vertex_collection("Users"):
      movie_rating_graph.vertex_collection("Users")

  # Create a new vertex collection named "Movie" if it does not exist.
  if not movie_rating_graph.has_vertex_collection("Movie"):
      movie_rating_graph.vertex_collection("Movie")

  # creating edge definitions named "Ratings. This creates any missing
  # collections and returns an API wrapper for "Ratings" edge collection.
  if not movie_rating_graph.has_edge_definition("Ratings"):
      Ratings = movie_rating_graph.create_edge_definition(
          edge_collection='Ratings',
          from_vertex_collections=['Users'],
          to_vertex_collections=['Movie']
      )

  return movie_rating_graph

def load_data_to_ArangoDB(login):
    global movie_rec_db
    movie_rec_db = oasis.connect_python_arango(login)
    
    movie_rating_graph = create_smart_graph()

    if not movie_rec_db.has_collection("Movie"):
        movie_rec_db.create_collection("Movie", replication_factor=3)
        
    batch = []
    BATCH_SIZE = 128
    batch_idx = 1
    index = 0
    movie_collection = movie_rec_db["Movie"]
    
    # loading movies metadata information into ArangoDB's Movie collection
    for idx in tqdm(range(len(m_id))):
        insert_doc = {}
        tmdb_id = id_map.loc[id_map['movieId'] == m_id[idx]]
      
        if tmdb_id.size == 0:
            print('No Meta data information at:', m_id[idx])
            
    
        else:
            tmdb_id = int(tmdb_id.iloc[:,1][0])
            emb_id = "Movie/" + str(movie_mappings[m_id[idx]])
            insert_doc["_id"] = emb_id
            m_meta = sampled_md.loc[sampled_md['id'] == tmdb_id]
            # adding movie metadata information 
            m_title = m_meta.iloc[0]['title']
            m_poster = m_meta.iloc[0]['poster_path']
            m_description = m_meta.iloc[0]['description']
            m_language = m_meta.iloc[0]['original_language']
            m_genre = m_meta.iloc[0]['genres']
            m_genre = yaml.load(m_genre, Loader=yaml.BaseLoader)
            genres = [g['name'] for g in m_genre]
             
            insert_doc["movieId"] = m_id[idx]
            insert_doc["mapped_movieId"] = movie_mappings[m_id[idx]]
            insert_doc["tmdbId"] = tmdb_id
            insert_doc['movie_title'] = m_title
         
            insert_doc['description'] = m_description
            insert_doc['genres'] = genres
            insert_doc['language'] = m_language
            
            if str(m_poster) == "nan":
                insert_doc['poster_path'] = "No poster path available"
            else:
                insert_doc['poster_path'] = m_poster
            
            batch.append(insert_doc)
            index +=1
            last_record = (idx == (len(m_id) - 1))
            if index % BATCH_SIZE == 0:
                #print("Inserting batch %d" % (batch_idx))
                batch_idx += 1
                movie_collection.import_bulk(batch)
                batch = []   
            if last_record and len(batch) > 0:
                print("Inserting batch the last batch!")
                movie_collection.import_bulk(batch)

    if not movie_rec_db.has_collection("Users"):
        movie_rec_db.create_collection("Users", replication_factor=3)

    total_users = np.unique(ratings_df[['userId']].values.flatten()).shape[0]
    print("Total number of Users:", total_users)
    populate_user_collection(total_users)

    # This returns an API wrapper for "Ratings" collection.
    if not movie_rec_db.has_collection("Ratings"):
        movie_rec_db.create_collection("Ratings", edge=True, replication_factor=3)

    user_id, movie_id, ratings = ratings_df[['userId']].values.flatten(), ratings_df[['movieId']].values.flatten() , ratings_df[['rating']].values.flatten()
    create_ratings_graph(user_id, movie_id, ratings)


    return movie_rec_db


def make_pyg_graph(movie_rec_db):
    global device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    users = movie_rec_db.collection('Users')
    movies = movie_rec_db.collection('Movie')
    ratings_graph = movie_rec_db.collection('Ratings')

    edge_index, edge_label = create_pyg_edges(movie_rec_db.aql.execute('FOR doc IN Ratings RETURN doc'))
    
    title_emb = SequenceEncoder(movie_rec_db.aql.execute('FOR doc IN Movie RETURN doc'), model_name='all-MiniLM-L6-v2')
    encoded_genres = GenresEncoder(movie_rec_db.aql.execute('FOR doc IN Movie RETURN doc'))
    movie_x = torch.cat((title_emb, encoded_genres), dim=-1)

    global data
    data = HeteroData()
    data['user'].num_nodes = len(users)  # Users do not have any features.
    data['movie'].x = movie_x
    data['user', 'rates', 'movie'].edge_index = edge_index
    data['user', 'rates', 'movie'].edge_label = edge_label
    
    # Add user node features for message passing:
    data['user'].x = torch.eye(data['user'].num_nodes, device=device)
    del data['user'].num_nodes
    data = ToUndirected()(data)
    del data['movie', 'rev_rates', 'user'].edge_label  # Remove "reverse" label.

    data = data.to(device)

    train_data, val_data, test_data = T.RandomLinkSplit(
        num_val=0.1,
        num_test=0.1,
        neg_sampling_ratio=0.0,
        edge_types=[('user', 'rates', 'movie')],
        rev_edge_types=[('movie', 'rev_rates', 'user')],
    )(data)

    return data,train_data, val_data, test_data


def load_model(train_data, val_data, test_data):
    model = Model(hidden_channels=32)
    with torch.no_grad():
        model.encoder(train_data.x_dict, train_data.edge_index_dict)
    model.load_state_dict(torch.load('model.pt',map_location=torch.device('cpu')))
    model.eval()
    return model

def get_recommendation(model,data,user_id):

    movies = movie_rec_db.collection('Movie')
    total_movies = len(movies)
    
    user_row = torch.tensor([user_id] * total_movies)
    all_movie_ids = torch.arange(total_movies)
    edge_label_index = torch.stack([user_row, all_movie_ids], dim=0)
    pred = model(data.x_dict, data.edge_index_dict,edge_label_index)
    pred = pred.clamp(min=0, max=5)
    # we will only select movies for the user where the predicting rating is =5
    rec_movie_ids = (pred == 5).nonzero(as_tuple=True)
    top_ten_recs = [rec_movies for rec_movies in rec_movie_ids[0].tolist()[:10]] 
    return {'user': user_id, 'rec_movies': top_ten_recs}
    


def train(train_data, val_data, test_data):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    #make weight
    weight = torch.bincount(train_data['user', 'movie'].edge_label)
    weight = weight.max() / weight
    model = Model(hidden_channels=32).to(device)
    with torch.no_grad():
        model.encoder(train_data.x_dict, train_data.edge_index_dict)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    # Train loop
    for epoch in range(1, 300):
        loss = train()
        train_rmse = test(train_data)
        val_rmse = test(val_data)
        test_rmse = test(test_data)
        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_rmse:.4f}, '
              f'Val: {val_rmse:.4f}, Test: {test_rmse:.4f}')