| """
|
| DataSynthis_ML_JobTask - Movie Recommendation Model
|
| A movie recommendation system using collaborative filtering and matrix factorization.
|
| """
|
|
|
| import pandas as pd
|
| import numpy as np
|
| from sklearn.metrics.pairwise import cosine_similarity
|
| from sklearn.decomposition import TruncatedSVD
|
| import os
|
| import urllib.request
|
| import zipfile
|
| import pickle
|
| from typing import List, Dict, Optional, Union
|
|
|
|
|
| class MovieRecommender:
|
| """
|
| Movie Recommendation Model using collaborative filtering and SVD.
|
| """
|
|
|
| def __init__(self):
|
| self.ratings = None
|
| self.movies = None
|
| self.user_item_matrix = None
|
| self.item_similarity = None
|
| self.item_similarity_df = None
|
| self.svd_model = None
|
| self.pred_svd_df = None
|
| self.is_trained = False
|
|
|
| def load_data(self):
|
| """Load MovieLens 100k dataset."""
|
| dataset_url = "http://files.grouplens.org/datasets/movielens/ml-100k.zip"
|
| dataset_path = "ml-100k"
|
|
|
| if not os.path.exists(dataset_path):
|
| if os.path.exists("ml-100k.zip"):
|
| print("Extracting existing MovieLens 100k dataset...")
|
| with zipfile.ZipFile("ml-100k.zip", "r") as zip_ref:
|
| zip_ref.extractall(".")
|
| print("Extraction complete.")
|
| else:
|
| print("Downloading MovieLens 100k dataset...")
|
| try:
|
| urllib.request.urlretrieve(dataset_url, "ml-100k.zip")
|
| with zipfile.ZipFile("ml-100k.zip", "r") as zip_ref:
|
| zip_ref.extractall(".")
|
| print("Download complete.")
|
| except Exception as e:
|
| print(f"Download failed: {e}")
|
| raise Exception("Could not download dataset")
|
|
|
|
|
| self.ratings = pd.read_csv(
|
| "ml-100k/u.data",
|
| sep="\t",
|
| names=["user_id", "movie_id", "rating", "timestamp"]
|
| )
|
|
|
|
|
| self.movies = pd.read_csv(
|
| "ml-100k/u.item",
|
| sep="|",
|
| encoding="ISO-8859-1",
|
| names=["movie_id", "title", "release_date", "video_release_date", "IMDb_URL",
|
| "unknown", "Action", "Adventure", "Animation", "Children", "Comedy",
|
| "Crime", "Documentary", "Drama", "Fantasy", "Film-Noir", "Horror",
|
| "Musical", "Mystery", "Romance", "Sci-Fi", "Thriller", "War", "Western"]
|
| )
|
|
|
|
|
| self.ratings.drop("timestamp", axis=1, inplace=True)
|
|
|
| print(f"Loaded {len(self.ratings)} ratings from {len(self.ratings['user_id'].unique())} users")
|
| print(f"Loaded {len(self.movies)} movies")
|
|
|
| def train(self):
|
| """Train the recommendation models."""
|
| if self.ratings is None:
|
| self.load_data()
|
|
|
|
|
| self.user_item_matrix = self.ratings.pivot(
|
| index='user_id', columns='movie_id', values='rating'
|
| )
|
|
|
|
|
| self.item_similarity = cosine_similarity(self.user_item_matrix.T.fillna(0))
|
| self.item_similarity_df = pd.DataFrame(
|
| self.item_similarity,
|
| index=self.user_item_matrix.columns,
|
| columns=self.user_item_matrix.columns
|
| )
|
|
|
|
|
| R = self.user_item_matrix.fillna(0)
|
| self.svd_model = TruncatedSVD(n_components=20, random_state=42)
|
| U = self.svd_model.fit_transform(R)
|
| Sigma = np.diag(self.svd_model.singular_values_)
|
| Vt = self.svd_model.components_
|
| pred_svd = np.dot(np.dot(U, Sigma), Vt)
|
| self.pred_svd_df = pd.DataFrame(pred_svd, index=R.index, columns=R.columns)
|
|
|
| self.is_trained = True
|
| print("Model training completed!")
|
|
|
| def predict_ratings_cf(self, user_id: int) -> pd.Series:
|
| """Predict ratings using collaborative filtering."""
|
| if not self.is_trained:
|
| raise ValueError("Model must be trained first")
|
|
|
| if user_id not in self.user_item_matrix.index:
|
| raise ValueError(f"User {user_id} not found in dataset")
|
|
|
| user_ratings = self.user_item_matrix.loc[user_id]
|
| weighted_sum = self.item_similarity_df.dot(user_ratings.fillna(0))
|
| sim_sum = np.abs(self.item_similarity_df).dot(user_ratings.notna().astype(int))
|
| pred = weighted_sum / np.maximum(sim_sum, 1e-9)
|
| return pred
|
|
|
| def recommend_movies(self, user_id: int, n_recommendations: int = 10,
|
| method: str = "svd") -> List[Dict]:
|
| """
|
| Get movie recommendations for a user.
|
|
|
| Args:
|
| user_id: User ID to get recommendations for
|
| n_recommendations: Number of recommendations to return
|
| method: "svd" or "cf" (collaborative filtering)
|
|
|
| Returns:
|
| List of dictionaries with movie recommendations
|
| """
|
| if not self.is_trained:
|
| self.train()
|
|
|
|
|
| if user_id not in self.user_item_matrix.index:
|
| available_users = sorted(self.user_item_matrix.index.tolist())
|
| return [{
|
| "error": f"User {user_id} not found",
|
| "available_users": f"Available user IDs: {available_users[:10]}... (showing first 10)"
|
| }]
|
|
|
|
|
| if method == "svd":
|
| preds = self.pred_svd_df.loc[user_id]
|
| else:
|
| preds = self.predict_ratings_cf(user_id)
|
|
|
|
|
| watched = self.ratings[self.ratings.user_id == user_id].movie_id.values
|
| preds = preds.drop(watched, errors='ignore')
|
|
|
|
|
| top_movies = preds.sort_values(ascending=False).head(n_recommendations).index
|
| recommendations = self.movies[self.movies.movie_id.isin(top_movies)][["movie_id", "title"]]
|
|
|
|
|
| result = []
|
| for _, row in recommendations.iterrows():
|
| result.append({
|
| "movie_id": int(row["movie_id"]),
|
| "title": row["title"],
|
| "predicted_rating": float(preds[row["movie_id"]])
|
| })
|
|
|
| return result
|
|
|
| def get_user_stats(self, user_id: int) -> Dict:
|
| """Get statistics for a user."""
|
| if not self.is_trained:
|
| self.train()
|
|
|
| if user_id not in self.user_item_matrix.index:
|
| return {"error": f"User {user_id} not found"}
|
|
|
| user_ratings = self.ratings[self.ratings.user_id == user_id]
|
|
|
| return {
|
| "user_id": user_id,
|
| "total_ratings": len(user_ratings),
|
| "average_rating": float(user_ratings["rating"].mean()),
|
| "rating_distribution": user_ratings["rating"].value_counts().to_dict()
|
| }
|
|
|
| def get_available_users(self) -> List[int]:
|
| """Get list of available user IDs."""
|
| if not self.is_trained:
|
| self.train()
|
| return sorted(self.user_item_matrix.index.tolist())
|
|
|
| def save_model(self, path: str):
|
| """Save the trained model."""
|
| if not self.is_trained:
|
| raise ValueError("Model must be trained first")
|
|
|
| model_data = {
|
| 'ratings': self.ratings,
|
| 'movies': self.movies,
|
| 'user_item_matrix': self.user_item_matrix,
|
| 'item_similarity_df': self.item_similarity_df,
|
| 'svd_model': self.svd_model,
|
| 'pred_svd_df': self.pred_svd_df,
|
| 'is_trained': self.is_trained
|
| }
|
|
|
| with open(path, 'wb') as f:
|
| pickle.dump(model_data, f)
|
|
|
| print(f"Model saved to {path}")
|
|
|
| def load_model(self, path: str):
|
| """Load a trained model."""
|
| with open(path, 'rb') as f:
|
| model_data = pickle.load(f)
|
|
|
| self.ratings = model_data['ratings']
|
| self.movies = model_data['movies']
|
| self.user_item_matrix = model_data['user_item_matrix']
|
| self.item_similarity_df = model_data['item_similarity_df']
|
| self.svd_model = model_data['svd_model']
|
| self.pred_svd_df = model_data['pred_svd_df']
|
| self.is_trained = model_data['is_trained']
|
|
|
| print(f"Model loaded from {path}")
|
|
|
|
|
|
|
| model = MovieRecommender()
|
|
|
| def predict(user_id: int, n_recommendations: int = 10, method: str = "svd") -> List[Dict]:
|
| """
|
| Inference function for Hugging Face model.
|
|
|
| Args:
|
| user_id: User ID to get recommendations for
|
| n_recommendations: Number of recommendations (default: 10)
|
| method: Recommendation method - "svd" or "cf" (default: "svd")
|
|
|
| Returns:
|
| List of movie recommendations
|
| """
|
| return model.recommend_movies(user_id, n_recommendations, method) |
