Spaces:
Runtime error
Runtime error
| import ast | |
| import pandas as pd | |
| from recommendation_model import Model | |
| """ | |
| The dataset is obtained from TMDB 5000 Movie Dataset | |
| https://www.kaggle.com/datasets/tmdb/tmdb-movie-metadata | |
| """ | |
| def get_name(x): | |
| return ', '.join([i['name'].lower() for i in ast.literal_eval(x)][:5]) | |
| def get_director(x): | |
| return ', '.join(i['name'].lower() for i in ast.literal_eval(x) if i['job'].lower() == 'director') | |
| def get_year(x): | |
| return str(x)[:4] | |
| def normalize_data(x): | |
| return (x - x.min()) / (x.max() - x.min()) | |
| raw1 = pd.read_csv('tmdb_5000_movies.csv') | |
| raw2 = pd.read_csv('tmdb_5000_credits.csv') | |
| raw2 = raw2.rename(columns={'movie_id': 'id'}) | |
| df = pd.merge(raw1, raw2, on='id') | |
| df = df.drop([ | |
| 'budget', | |
| 'homepage', | |
| 'overview', | |
| 'tagline', | |
| 'status', | |
| 'production_companies', | |
| 'production_countries', | |
| 'revenue', | |
| 'spoken_languages', | |
| 'title_x', | |
| 'title_y', | |
| 'vote_count' | |
| ], axis=1) | |
| df['genres'] = df['genres'].map(get_name) | |
| df['keywords'] = df['keywords'].map(get_name) | |
| df['cast'] = df['cast'].map(get_name) | |
| df['crew'] = df['crew'].map(get_director) | |
| df['release_date'] = df['release_date'].map(get_year) | |
| for i in range(len(df)): | |
| df.loc[i, 'id'] = i | |
| df = df.rename(columns={ | |
| 'original_language': 'language', | |
| 'original_title': 'title', | |
| 'release_date': 'year', | |
| 'vote_average': 'rating', | |
| 'crew': 'director' | |
| }) | |
| df = df[[ | |
| 'id', | |
| 'title', | |
| 'genres', | |
| 'keywords', | |
| 'director', | |
| 'cast', | |
| 'year', | |
| 'language', | |
| 'runtime', | |
| 'popularity', | |
| 'rating' | |
| ]] | |
| df['id'] = df['id'].apply(lambda x: str(x)) | |
| df['year'] = df['year'].apply(lambda x: str(x)) | |
| df['runtime'] = normalize_data(df['runtime']) | |
| df['popularity'] = normalize_data(df['popularity']) | |
| df['rating'] = normalize_data(df['rating']) | |
| df_trim = df[['title', 'genres', 'keywords', 'director', 'cast']] | |
| model = Model(df) | |
| model.fit(save=True) | |