app.py

import gradio as gr
import numpy as np
import pandas as pd
import torch

from surprise import Reader, Dataset, SVD
from surprise.model_selection import cross_validate
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity


def load_model() -> SentenceTransformer:
    """
    Loads a pre-trained SentenceTransformer model.

    :return: The loaded SentenceTransformer model.
    """
    if torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"

    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2").to(device)
    return model


def encode_and_calculate_similarity(
    model: SentenceTransformer, df_merged: pd.DataFrame
) -> np.ndarray:
    """
    Encodes sentences using the provided SentenceTransformer model and calculates cosine similarity.

    :param model: The SentenceTransformer model to use for encoding.
    :param df_merged: The DataFrame containing the sentences to encode.
    :return: The cosine similarity matrix.
    """
    sentence_embeddings = model.encode(df_merged["soup"].tolist())

    cos_sim = cosine_similarity(sentence_embeddings)

    return cos_sim


def svd(df_ratings: pd.DataFrame) -> SVD:
    """
    Performs Singular Value Decomposition (SVD) on the provided DataFrame.

    :param df_ratings: The DataFrame containing user ratings.
    :return: The trained SVD model.
    """
    reader = Reader()
    data = Dataset.load_from_df(df_ratings[["userId", "movieId", "rating"]], reader)
    svd = SVD()
    cross_validate(svd, data, measures=["RMSE", "MAE"], cv=5, verbose=True)

    trainset = data.build_full_trainset()
    svd.fit(trainset)
    return svd


def get_sorted_similar_movies(
    title: str, cos_sim: np.ndarray, df_merged: pd.DataFrame
) -> pd.DataFrame:
    """
    Get a sorted DataFrame of movies based on their similarity scores to a given movie.

    :param title: The title of the movie to find similar movies for.
    :param cos_sim: The cosine similarity matrix of movies.
    :param df_merged: The DataFrame containing movie details.
    :return: A sorted DataFrame of similar movies.
    """
    try:
        # Get the index of the movie that matches the title
        movie_index = movie_indices[title.lower()]

        # If there are multiple movies with the same title, pick the first one.
        if isinstance(movie_index, pd.Series):
            movie_index = movie_index[0]

    except KeyError:
        print(f"Movie '{title}' not found. Please enter a valid movie title.")
        return None

    # Get the pairwise similarity scores of all movies with that movie
    sim_scores = list(enumerate(cos_sim[movie_index]))

    # Sort the movies based on the similarity scores
    sim_scores = sorted(sim_scores, key=lambda x: x[1], reverse=True)[1:]

    # Get the movie indices
    sorted_movie_indices = [sim_score[0] for sim_score in sim_scores]

    # Get the similarity scores
    sorted_similarity_scores = [format(sim_score[1], ".1f") for sim_score in sim_scores]

    movie_details = [
        "id",
        "title",
        "genres",
        "original_language",
        "production_countries",
        "release_date",
        "runtime",
        "weighted_rating",
    ]

    sorted_similar_movies = df_merged.loc[sorted_movie_indices, movie_details]

    sorted_similar_movies["similarity_scores"] = sorted_similarity_scores

    return sorted_similar_movies


def get_qualified_movies(
    df_qualified: pd.DataFrame, sorted_similar_movies: pd.DataFrame
) -> pd.DataFrame:
    """
    Filter out movies that are not in the qualified movies chart and sort the movies based on similarity scores and IMDB's weighted rating.

    :param df_qualified: The DataFrame containing qualified movie details.
    :param sorted_similar_movies: The DataFrame containing movie details sorted by similarity scores.
    :return: A Pandas DataFrame containing the qualified movies sorted by similarity scores and IMDB's weighted rating..
    """

    qualified_movies = sorted_similar_movies[
        sorted_similar_movies["id"].isin(df_qualified["id"])
    ]
    qualified_movies = qualified_movies.sort_values(
        by=["similarity_scores", "weighted_rating"], ascending=False
    )
    return qualified_movies


def predict_user_rating(
    userId: int, qualified_movies: pd.DataFrame, indices_map: pd.DataFrame
) -> pd.DataFrame:
    """
    Predict the user rating for qualified movies using SVD and return the sorted DataFrame.

    :param userId: The ID of the user.
    :param qualified_movies:  A Pandas DataFrame containing qualified movies data.
    :return: A Pandas DataFrame containing the final qualified movies sorted by estimated user ratings.
    """
    # Calculate estimated user ratings for qualified movies using SVD
    qualified_movies["predicted_user_rating"] = qualified_movies["id"].apply(
        lambda x: round(svd.predict(userId, indices_map.loc[x]["movieId"]).est, 1)
    )
    final_qualified_movies = qualified_movies.sort_values(
        by=["predicted_user_rating", "similarity_scores", "weighted_rating"],
        ascending=False,
    )
    return final_qualified_movies


def get_movie_recommendations_hybrid(
    title: str, user_id: int
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Get movie recommendations based on a given title and user ID.

    :param title: The title of the movie to find similar movies for.
    :param userId: The ID of the user.
    :return: A tuple of two Pandas DataFrames.
        The first DataFrame contains the recommended movies.
        The second DataFrame contains the recommendation criteria (ID, Title, Predicted User Rating, Similarity Score, Weighted Rating).
    """
    # Get recommended movie indices based on the given title
    sorted_similar_movies = get_sorted_similar_movies(title, cos_sim, df_merged)

    # Filter out bad movies and select the top 50 qualified movies
    qualified_movies = get_qualified_movies(df_qualified, sorted_similar_movies).head(
        50
    )

    # Predict user ratings for qualified movies and select the top recommended movies
    recommended_movies = predict_user_rating(
        user_id, qualified_movies, indices_map
    ).head(5)

    recommended_movies.columns = [
        "ID",
        "Title",
        "Genres",
        "Language",
        "Production Countries",
        "Release Date",
        "Runtime",
        "Weighted Rating",
        "Similarity Score",
        "Predicted User Rating",
    ]

    recommendation_criteria = recommended_movies[
        ["ID", "Title", "Predicted User Rating", "Similarity Score", "Weighted Rating"]
    ]
    recommended_movies.drop(
        ["Predicted User Rating", "Similarity Score", "Weighted Rating"],
        axis=1,
        inplace=True,
    )
    return recommended_movies, recommendation_criteria


if __name__ == "__main__":
    df_qualified = pd.read_csv("data/qualified_movies.csv")
    df_ratings = pd.read_csv("data/ratings_small.csv")
    df_merged = pd.read_csv("data/df_merged.csv")

    model = load_model()
    cos_sim = encode_and_calculate_similarity(model, df_merged)
    movie_indices = pd.Series(
        df_merged.index, index=df_merged["title"].apply(lambda title: title.lower())
    ).drop_duplicates()

    svd = svd(df_ratings)
    indices_map = df_merged.set_index("id")

    with gr.Blocks(theme=gr.themes.Soft(text_size="lg")) as demo:
        gr.Markdown(
            """
        # Movie Recommendation System
        """
        )
        title = gr.Dropdown(
            choices=df_merged["title"].unique().tolist(),
            label="Movie Title",
            value="Iron Man",
        )
        user_id = gr.Number(
            value=1, label="User ID", info="Please enter a number between 1 and 671!"
        )
        recommend_button = gr.Button("Get Movie Recommendations")
        recommended_movies = gr.DataFrame(label="Movie Recommendations")
        recommendation_criteria = gr.DataFrame(label="Recommendation Criteria")
        recommend_button.click(
            get_movie_recommendations_hybrid,
            inputs=[title, user_id],
            outputs=[recommended_movies, recommendation_criteria],
        )
        examples = gr.Examples(
            examples=[
                "Captain America: The First Avenger",
                "The Conjuring",
                "Toy Story",
                "Final Destination 5",
            ],
            inputs=[title],
        )

    demo.launch()