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(): 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): sentence_embeddings = model.encode(df_merged["soup"].tolist()) cos_sim = cosine_similarity(sentence_embeddings) return cos_sim def svd(): 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_movie_indices(title: str, cos_sim: np.ndarray) -> list[int]: """ Retrieve the sorted indices 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. :return: A list of sorted movie indices. """ 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] return sorted_movie_indices def get_qualified_movies( df: pd.DataFrame, df_qualified: pd.DataFrame, sorted_movie_indices: list[int] ) -> pd.DataFrame: """ Filter out movies that are not in the qualified movies chart based on IMDB's weighted rating. :param df: The DataFrame containing movie details. :param df_qualified: The DataFrame containing qualified movie details. :param sorted_movie_indices: A list of movie indices sorted by similarity scores. :return: A Pandas DataFrame containing the qualified movies sorted by similarity scores. """ movie_details = [ "id", "title", "genres", "original_language", "production_countries", "release_date", "runtime", ] sorted_movies = df.loc[sorted_movie_indices, movie_details] qualified_movies = sorted_movies[sorted_movies["id"].isin(df_qualified["id"])] 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, 2) ) final_qualified_movies = qualified_movies.sort_values( by=["predicted_user_rating"], ascending=False ) return final_qualified_movies def get_movie_recommendations_hybrid(title: str, userId: int) -> 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 Pandas DataFrame containing the recommended movies """ # Get recommended movie indices based on the given title sorted_movie_indices = get_sorted_movie_indices(title, cos_sim) # Filter out bad movies and select the top 50 qualified movies qualified_movies = get_qualified_movies( df_merged, df_qualified, sorted_movie_indices ).head(50) # Predict user ratings for qualified movies and select the top recommended movies recommended_movies = predict_user_rating( userId, qualified_movies, indices_map ).head(5) recommended_movies.columns = [ "ID", "Title", "Genres", "Language", "Production Countries", "Release Date", "Runtime", "Predicted User Rating", ] return recommended_movies 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) movie_indices = pd.Series( df_merged.index, index=df_merged["title"].apply(lambda title: title.lower()) ).drop_duplicates() svd = svd() 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") recommend_button.click( get_movie_recommendations_hybrid, inputs=[title, user_id], outputs=recommended_movies, ) examples = gr.Examples( examples=[ "Captain America: The First Avenger", "The Conjuring", "Toy Story", "Final Destination 5", ], inputs=[title], ) demo.launch()