# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Implementing Probabilistic Matrix Factorization in PyTorch
#
# From this paper: https://papers.nips.cc/paper/3208-probabilistic-matrix-factorization.pdf.
#
# Data from here: https://grouplens.org/datasets/movielens/
#
# Our goal here is to find feature vectors for our users and movies that allow us to compare them. We start with a big matrix $R$ of movie ratings, with users as rows and movies as columns. Since users don't rate all movies, this is a pretty sparse matrix with a lot of missing values. 
#
# What we can do is factor $R$ into two matrices $U$ and $V$ representing latent feature vectors for users and movies. With these latent features, we can compare users with each other, movies with each other, and predict ratings for movies users haven't seen.
#
# ![matrix factorization](matrix_factorization.png)
#
# The big idea behind this paper is that we're going to treat the latent vectors as parameters in a Bayesian model. As a reminder, Bayes theorem:
#
# $$
# \large P(\theta \mid D) \propto P(D \mid \theta) P(\theta)
# $$
#
# In our model we try to predict the rating matrix $R$ with $U$ and $V$
#
# $$ \large \hat{R} = U^T V$$
#
# In our model we assume the ratings are drawn from a normal distribution with mean $\hat{R}$. What's really cool is that we can place priors on our latent features. 
#
# $$
# \begin{aligned}
# R &\sim \mathrm{Normal}(U^T V, \sigma^2) \\
# U &\sim \mathrm{Normal}(0, \sigma_U^2) \\
# V &\sim \mathrm{Normal}(0, \sigma_V^2)
# \end{aligned}
# $$
#
# The authors of the paper go further and build a hierachical structure for the user vectors
#
# $$
# \begin{aligned}
# U_i &= Y_i + \frac{\sum_k^M I_{ik}W_k}{\sum_k^M I_{ik}} \\
# Y &\sim \mathrm{Normal}(0, \sigma_U^2) \\
# W &\sim \mathrm{Normal}(0, \sigma_W^2)
# \end{aligned}
# $$
#
# With this model, we can maximize the posterior probability with respect to $U$ and $V$, or $V$, $Y$, and $W$ for the hierarchical model. In effect this is just a linear model with fancy regularization.
#
# The authors also do things like converting the ratings to be between 0 and 1, then taking the sigmoid of $U^T V$. For empirical reasons.

import pandas as pd
import torch

ratings = pd.read_csv('ml-latest-small/ratings.csv')

ratings.describe()

rating_matrix = ratings.pivot(index='userId', columns='movieId', values='rating')
n_users, n_movies = rating_matrix.shape
# Scaling ratings to between 0 and 1, this helps our model by constraining predictions
min_rating, max_rating = ratings['rating'].min(), ratings['rating'].max()
rating_matrix = (rating_matrix - min_rating) / (max_rating - min_rating)

# Replacing missing ratings with -1 so we can filter them out later
rating_matrix[rating_matrix.isnull()] = -1
rating_matrix = torch.FloatTensor(rating_matrix.values)

# This is how we can define our feature matrices
# We're going to be training these, so we'll need gradients
latent_vectors = 5
user_features = torch.randn(n_users, latent_vectors, requires_grad=True)
user_features.data.mul_(0.01)
movie_features = torch.randn(n_movies, latent_vectors, requires_grad=True)
movie_features.data.mul_(0.01)


class PMFLoss(torch.nn.Module):
    def __init__(self, lam_u=0.3, lam_v=0.3):
        super().__init__()
        self.lam_u = lam_u
        self.lam_v = lam_v
    
    def forward(self, matrix, u_features, v_features):
        non_zero_mask = (matrix != -1).type(torch.FloatTensor)
        predicted = torch.sigmoid(torch.mm(u_features, v_features.t()))
        
        diff = (matrix - predicted)**2
        prediction_error = torch.sum(diff*non_zero_mask)

        u_regularization = self.lam_u * torch.sum(u_features.norm(dim=1))
        v_regularization = self.lam_v * torch.sum(v_features.norm(dim=1))
        
        return prediction_error + u_regularization + v_regularization


criterion = PMFLoss()
loss = criterion(rating_matrix, user_features, movie_features)

# +
# Actual training loop now

latent_vectors = 30
user_features = torch.randn(n_users, latent_vectors, requires_grad=True)
user_features.data.mul_(0.01)
movie_features = torch.randn(n_movies, latent_vectors, requires_grad=True)
movie_features.data.mul_(0.01)

pmferror = PMFLoss(lam_u=0.05, lam_v=0.05)
optimizer = torch.optim.Adam([user_features, movie_features], lr=0.01)
for step, epoch in enumerate(range(1000)):
    optimizer.zero_grad()
    loss = pmferror(rating_matrix, user_features, movie_features)
    loss.backward()
    optimizer.step()
    if step % 50 == 0:
        print(f"Step {step}, {loss:.3f}")

# +
# Checking if our model can reproduce the true user ratings
user_idx = 7
user_ratings = rating_matrix[user_idx, :]
true_ratings = user_ratings != -1
predictions = torch.sigmoid(torch.mm(user_features[user_idx, :].view(1, -1), movie_features.t()))
predicted_ratings = (predictions.squeeze()[true_ratings]*(max_rating - min_rating) + min_rating).round()
actual_ratings = (user_ratings[true_ratings]*(max_rating - min_rating) + min_rating).round()

print("Predictions: \n", predicted_ratings)
print("Truth: \n", actual_ratings)

StarcoderdataJupyterScriptsDedupFiltered