Spaces:
Sleeping
Sleeping
File size: 3,661 Bytes
b8aefd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import pandas as pd
import numpy as np
import pickle
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import ast
DATA_PATH = '../data/'
MAX_ITEMS = 12000
def process_movies():
print("🎬 Processing TMDB Movies...")
movies = pd.read_csv(DATA_PATH + 'tmdb_5000_movies.csv')
movies['release_date'] = pd.to_datetime(movies['release_date'], errors='coerce')
movies = movies.dropna(subset=['release_date'])
movies = movies[
(movies['release_date'].dt.year >= 2000) |
((movies['release_date'].dt.year < 2000) & (movies['vote_count'] > 1500))
].copy()
def parse_genres(x):
try:
return " ".join([i['name'] for i in ast.literal_eval(x)])
except:
return ""
movies['genres_str'] = movies['genres'].apply(parse_genres)
movies['tags'] = movies['overview'].fillna('') + " " + movies['genres_str']
movies['type'] = 'Movie'
movies = movies[['id', 'title', 'tags', 'vote_average', 'vote_count', 'type', 'genres_str']]
movies.rename(columns={'vote_average': 'rating', 'genres_str': 'genre_list'}, inplace=True)
return movies
def process_anime():
print("🍙 Processing Anime...")
anime = pd.read_csv(DATA_PATH + 'anime.csv')
anime = anime[anime['members'] > 40000].copy()
anime['name'] = anime['name'].fillna('')
anime['genre'] = anime['genre'].fillna('')
anime['type'] = anime['type'].fillna('Anime')
anime['tags'] = anime['genre'] + " " + anime['type'] + " " + anime['name']
anime['genre_list'] = "Anime"
anime.rename(columns={'anime_id': 'id', 'name': 'title', 'rating': 'rating', 'members': 'vote_count'}, inplace=True)
anime['type'] = 'Anime'
anime = anime[['id', 'title', 'tags', 'rating', 'vote_count', 'type', 'genre_list']]
return anime
def build_engine():
df_movies = process_movies()
df_anime = process_anime()
combined = pd.concat([df_movies, df_anime], ignore_index=True)
combined = combined.sample(frac=1, random_state=42).reset_index(drop=True)
if len(combined) > MAX_ITEMS:
print(f"⚠️ Trimming dataset from {len(combined)} to {MAX_ITEMS}...")
combined = combined.head(MAX_ITEMS)
print(f"📊 Total Database: {len(combined)} items.")
print("🧠 Training NLP Model...")
cv = CountVectorizer(max_features=5000, stop_words='english')
vectors = cv.fit_transform(combined['tags']).toarray()
print("📐 Calculating Cosine Similarity...")
similarity = cosine_similarity(vectors)
print("📝 Generating Quiz Data...")
all_genres = set()
for g in combined['genre_list'].dropna():
cleaned = g.replace(" ", ",").split(",")
for item in cleaned:
if item and len(item) > 2: all_genres.add(item.strip())
quiz_data = {}
for genre in all_genres:
if genre == "Anime":
mask = combined['type'] == 'Anime'
else:
mask = (combined['genre_list'].str.contains(genre, case=False, na=False)) & (combined['type'] == 'Movie')
top_items = combined[mask].sort_values(by='rating', ascending=False).head(20)
if not top_items.empty:
quiz_data[genre] = top_items[['id', 'title', 'type']].to_dict('records')
print("💾 Saving Artifacts...")
pickle.dump(combined, open('movie_list.pkl', 'wb'))
pickle.dump(similarity, open('similarity.pkl', 'wb'))
pickle.dump(quiz_data, open('quiz_data.pkl', 'wb'))
print("🎉 DONE! Backend ready.")
if __name__ == "__main__":
build_engine() |