|
|
"""
|
|
|
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) |
