core.py

import pip

def install(package):
    if hasattr(pip, 'main'):
        pip.main(['install', package])
    else:
        pip._internal.main(['install', package])

print("Everything goes bang.")
install('torch_geometric')
install('torch_scatter')
install('torch_sparse')
print("It's havoc baby!")

import pickle
import numpy as np 
import pandas as pd
import random
from tqdm import tqdm
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import torch
from torch import nn, optim, Tensor
from torch_sparse import SparseTensor, matmul
from torch_geometric.utils import structured_negative_sampling
from torch_geometric.data import download_url, extract_zip
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.typing import Adj
from sklearn.neighbors import BallTree
from thefuzz import fuzz
from thefuzz import process

class LightGCN(MessagePassing):
    def __init__(self, num_users, num_items, embedding_dim=64, diffusion_steps=3, add_self_loops=False):
        super().__init__()
        
        # Number of users and items in the graph
        self.num_users = num_users
        self.num_items = num_items
        
        # Embedding dimension for user and item nodes
        self.embedding_dim = embedding_dim
        
        # Number of diffusion steps (K) for multi-scale diffusion
        self.diffusion_steps = diffusion_steps
        
        # Whether to add self-loops to the adjacency matrix
        self.add_self_loops = add_self_loops
        
        # Initialize embeddings for users and items (E^0)
        self.users_emb = nn.Embedding(num_embeddings=self.num_users, embedding_dim=self.embedding_dim)  # e_u^0
        self.items_emb = nn.Embedding(num_embeddings=self.num_items, embedding_dim=self.embedding_dim)  # e_i^0

        # Initialize embedding weights with a normal distribution (mean=0, std=0.1)
        nn.init.normal_(self.users_emb.weight, std=0.1)
        nn.init.normal_(self.items_emb.weight, std=0.1)

    def forward(self, edge_index: SparseTensor):
        # Compute the symmetrically normalized adjacency matrix (A_hat or \tilde{A})
        edge_index_norm = gcn_norm(edge_index, add_self_loops=self.add_self_loops)

        # Get initial embeddings E^0 for all nodes (users and items)
        emb_0 = torch.cat([self.users_emb.weight, self.items_emb.weight])  # E^0
        
        # List to store embeddings at each diffusion step (E^1, E^2, ..., E^K)
        embs = [emb_0]
        
        # Initialize the current embeddings to E^0
        emb_k = emb_0

        # Perform multi-scale diffusion for K steps
        for _ in range(self.diffusion_steps):
            # Propagate embeddings and update emb_k using the normalized adjacency matrix
            emb_k = self.propagate(edge_index_norm, x=emb_k)
            # Save embeddings at each diffusion step for later use
            embs.append(emb_k)

        # Stack all the embeddings along the second dimension (stack E^0, E^1, ..., E^K)
        embs = torch.stack(embs, dim=1)
        
        # Calculate the final embeddings by taking the mean of all diffusion embeddings (E^K)
        emb_final = torch.mean(embs, dim=1)  # E^K

        # Split the final embeddings into user embeddings (e_u^K) and item embeddings (e_i^K)
        users_emb_final, items_emb_final = torch.split(emb_final, [self.num_users, self.num_items])  # Splits into e_u^K and e_i^K

        # Returns the final embeddings for users (e_u^K), initial embeddings for users (e_u^0),
        # final embeddings for items (e_i^K), and initial embeddings for items (e_i^0)
        return users_emb_final, self.users_emb.weight, items_emb_final, self.items_emb.weight

    def message(self, x_j: Tensor) -> Tensor:
        # The message function is an identity function, i.e., it returns x_j itself
        return x_j

    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
        # Perform message passing and aggregation using the normalized adjacency matrix (A_hat or \tilde{A})
        return matmul(adj_t, x)


model = LightGCN(671, 9125)

def get_movie_recommendations(user_id, num_recomms):
    # Map the user ID to the corresponding index in the model's user embeddings
    user_index = user_mapping[user_id]

    # Retrieve the user embedding for the specified user
    user_embedding = model.users_emb.weight[user_index]

    # Calculate scores for all items using the user embedding
    scores = model.items_emb.weight @ user_embedding

    # Get the indices of the highest scores, including positive items and additional recommendations
    values, indices = torch.topk(scores, k=len(user_pos_items[user_id]) + num_recomms)

    # Retrieve the recommended movies that the user has already rated highly
    rated_movies = [index.cpu().item() for index in indices if index in user_pos_items[user_id]][:num_recomms]
    rated_movie_ids = [list(movie_mapping.keys())[list(movie_mapping.values()).index(movie)] for movie in rated_movies]

    # Retrieve the suggested movies for the user that they have not rated
    suggested_movies = [index.cpu().item() for index in indices if index not in user_pos_items[user_id]][:num_recomms]
    suggested_movie_ids = [list(movie_mapping.keys())[list(movie_mapping.values()).index(movie)] for movie in suggested_movies]

    return rated_movie_ids, suggested_movie_ids

addr = './'

model.load_state_dict(torch.load(addr + 'model.pth'))

final_movies_file = open(addr + 'final_movies.pkl', "rb")
final_movies = pickle.load(final_movies_file)
final_movies_file.close()

movie_embeds_file = open(addr + 'movie_embeds.pkl', "rb")
movie_embeds = pickle.load(movie_embeds_file)
movie_embeds_file.close()

btree_file = open(addr + 'btree.pkl', "rb")
btree = pickle.load(btree_file)
btree_file.close()

user_embeds_file = open(addr + 'user_embeds.pkl', "rb")
user_embeds = pickle.load(user_embeds_file)
user_embeds_file.close()

user_mapping_file = open(addr + 'user_mapping.pkl', "rb")
user_mapping = pickle.load(user_mapping_file)
user_mapping_file.close()

movie_mapping_file = open(addr + 'movie_mapping.pkl', "rb")
movie_mapping = pickle.load(movie_mapping_file)
movie_mapping_file.close()

user_pos_items_file = open(addr + 'user_pos_items.pkl', "rb")
user_pos_items = pickle.load(user_pos_items_file)
user_pos_items_file.close()

def create_user_embedding(movie_ratings, movies_df):
    # Convert the movie_ratings dictionary to a dataframe
    user_ratings_df = pd.DataFrame.from_dict(movie_ratings, orient='index', columns=['rating'])
    user_ratings_df['movieId'] = user_ratings_df.index
    print(user_ratings_df)
    print(user_movie_embeddings)
    # Merge the user_ratings_df with the movies_df to get the movie embeddings
    user_movie_embeddings = user_ratings_df.merge(movies_df, on='movieId', how='left')

    # Multiply the ratings with the movie embeddings
    user_movie_embeddings = user_movie_embeddings.iloc[:, 2:].values * user_movie_embeddings['rating'].values[:, np.newaxis]

    # Calculate the user embedding as the sum of the movie embeddings
    user_embedding = np.sum(user_movie_embeddings, axis=0)
    np.nan_to_num(user_embedding, 0)
    print(user_movie_embeddings.shape)
    return user_embedding

def find_closest_user(user_embedding, tree, user_embeddings):
    # Query the BallTree to find the closest user to the given user_embedding
    _, closest_user_index = tree.query([user_embedding], k=1)

    # Get the embedding of the closest user
    closest_user_embedding = user_embeddings.iloc[closest_user_index[0][0]]

    return closest_user_embedding


def drop_non_numerical_columns(df):
    non_numerical_columns = df.select_dtypes(exclude=[float, int]).columns
    return df.drop(columns=non_numerical_columns, inplace=False)

def output_list(input_dict, movies_df = movie_embeds, tree = btree, user_embeddings = user_embeds, movies = final_movies):
    movie_ratings = {}
    for movie_title in input_dict:
        matching_title = process.extractOne(movie_title, final_movies['title'].values, scorer=fuzz.partial_token_sort_ratio)[0]
        index = movies.index[movies['title'] == matching_title].tolist()[0]
        movie_ratings[index] = input_dict[movie_title]
    user_embed = create_user_embedding(movie_ratings, movie_embeds)
    # Call the find_closest_user function with the pre-built BallTree
    closest_user_embed = find_closest_user(user_embed, tree, user_embeds)
    rated_movie_ids, suggested_movie_ids = get_movie_recommendations(closest_user_embed['userId'], 5)
    out1 = [movie_id for movie_id in set(rated_movie_ids + suggested_movie_ids) if movie_id not in movie_ratings.keys()]
    out2 = [movies['title'][idx] for idx in out1]
    return out2

# output_list({1:1,2:2,3:3,4:4,5:5})