HarshaX's picture
Upload 4 files
4e3c433 verified
Raw
History Blame Contribute Delete
3.19 kB
import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression
from sklearn.svm import SVR
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error, explained_variance_score
from sklearn.metrics.pairwise import cosine_similarity
def preprocess_data(df):
df = df.dropna(subset=['rating', 'plot']).copy()
df['votes'] = df['votes'].astype(str).str.replace(',', '', regex=True)
df['votes'] = pd.to_numeric(df['votes'], errors='coerce').fillna(0)
df['year'] = pd.to_numeric(df['year'], errors='coerce').fillna(df['year'].median())
df['certificate'] = df['certificate'].fillna('Unknown')
le = LabelEncoder()
df['certificate'] = le.fit_transform(df['certificate'])
tfidf = TfidfVectorizer(stop_words='english', max_features=200)
tfidf_matrix = tfidf.fit_transform(df['plot'])
tfidf_feature_names = [f'tfidf_{name}' for name in tfidf.get_feature_names_out()]
tfidf_df = pd.DataFrame(tfidf_matrix.toarray(), columns=tfidf_feature_names, index=df.index)
genre_cols = ['Action', 'Adventure', 'Comedy', 'Crime', 'Family', 'Fantasy', 'Mystery', 'Sci-Fi', 'Thriller']
genre_df = df[genre_cols].astype(int)
features_df = pd.concat([df[['votes', 'certificate', 'year']], genre_df, tfidf_df], axis=1)
return features_df, df['rating'], tfidf, le, df
def train_models(X, y):
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
models = {
'RandomForest': RandomForestRegressor(random_state=42),
'LinearRegression': LinearRegression(),
'SVR': SVR()
}
trained_models = {}
metrics = {}
for name, model in models.items():
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
rmse = np.sqrt(mse)
r2 = r2_score(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
evs = explained_variance_score(y_test, y_pred)
metrics[name] = {
'RMSE': rmse,
'R2': r2,
'MAE': mae,
'Explained Variance': evs
}
trained_models[name] = model
return trained_models, metrics
def recommend_similar(df, tfidf, user_plot, user_genres, top_n=50):
user_vec = tfidf.transform([user_plot])
genre_cols = ['Action', 'Adventure', 'Comedy', 'Crime', 'Family', 'Fantasy', 'Mystery', 'Sci-Fi', 'Thriller']
user_genre_vec = np.array([1 if genre in user_genres else 0 for genre in genre_cols]).reshape(1, -1)
plot_sim = cosine_similarity(user_vec, tfidf.transform(df['plot']))[0]
genre_sim = cosine_similarity(user_genre_vec, df[genre_cols])[0]
combined_sim = 0.7 * plot_sim + 0.3 * genre_sim
top_indices = combined_sim.argsort()[::-1][:top_n]
return df.iloc[top_indices][['name', 'year', 'rating', 'votes', 'certificate']]