cinmon/coldstarter_example
0
1# -*- coding: utf-8 -*-2"""coldstarter_part1.ipynb3 4Automatically generated by Colab.5 6Original file is located at7 https://colab.research.google.com/drive/1wWPQCSAFIPCQORc763J-LB-MZTFXmIJT8"""9 10import sys11 12# If the kernel restarts due to numpy version change, print a message.13if 'numpy' in sys.modules:14 if sys.modules['numpy'].__version__.startswith('2.'):15 print("⚠️ NumPy 2.x is still present. Please restart the runtime "16 "('Runtime > Restart runtime') and run all cells again.")17 else:18 print("✅ NumPy 1.x installed successfully. Proceeding with other installations.")19 20# ==========================================21# 1. Imports22# ==========================================23 24import numpy as np25import pandas as pd26import matplotlib.pyplot as plt27import seaborn as sns28from collections import defaultdict29import warnings30warnings.filterwarnings('ignore')31 32import gradio as gr33import random34 35# Surprise library36from surprise import Dataset, Reader, accuracy37from surprise import NormalPredictor, BaselineOnly, KNNBasic, KNNWithMeans38from surprise import SVD, NMF39from surprise.model_selection import cross_validate, KFold, train_test_split40 41# Plot settings42plt.rcParams['figure.figsize'] = (12, 6)43plt.rcParams['font.size'] = 1244sns.set_style('whitegrid')45 46# Reproducibility47RANDOM_SEED = 4248np.random.seed(RANDOM_SEED)49 50print("✅ All libraries loaded successfully.")51 52# ==========================================53# 2. Data loading and preprocessing54# ==========================================55 56# Load MovieLens 100K (built into Surprise)57data = Dataset.load_builtin('ml-100k', prompt=False)58 59# Also load as a Pandas DataFrame (for analysis)60raw_ratings = data.raw_ratings # list of (user, item, rating, timestamp) tuples61df = pd.DataFrame(raw_ratings, columns=['user_id', 'item_id', 'rating', 'timestamp'])62df['rating'] = df['rating'].astype(float)63 64print(f"Total ratings: {len(df):,}")65print(f"Number of users: {df['user_id'].nunique()}")66print(f"Number of items: {df['item_id'].nunique()}")67print(f"Density: {len(df) / (df['user_id'].nunique() * df['item_id'].nunique()) * 100:.2f}%")68print(f"\nRating distribution:")69print(df['rating'].describe())70 71# Load movie metadata (for titles and genres)72import os73item_path = os.path.expanduser('~/.surprise_data/ml-100k/ml-100k/u.item')74movie_cols = ['item_id', 'title', 'release_date', 'video_release_date', 'IMDb_URL',75 'unknown', 'Action', 'Adventure', 'Animation', 'Children', 'Comedy',76 'Crime', 'Documentary', 'Drama', 'Fantasy', 'Film-Noir', 'Horror',77 'Musical', 'Mystery', 'Romance', 'Sci-Fi', 'Thriller', 'War', 'Western']78movies = pd.read_csv(item_path, sep='|', names=movie_cols, encoding='latin-1')79movie_id_to_title = dict(zip(movies['item_id'].astype(str), movies['title']))80title_to_movie_id = dict(zip(movies['title'], movies['item_id'].astype(str)))81 82# --- NEW: Build genre-enriched display labels ----------------------------83# Goal: instead of just "Toy Story (1995)" in the checkbox, show84# "Toy Story (1995) — Animation · Children · Comedy".85# We keep a separate mapping from this richer label back to the movie_id86# so the selection logic does not need to change.87GENRE_COLS = ['Action', 'Adventure', 'Animation', 'Children', 'Comedy',88 'Crime', 'Documentary', 'Drama', 'Fantasy', 'Film-Noir', 'Horror',89 'Musical', 'Mystery', 'Romance', 'Sci-Fi', 'Thriller', 'War', 'Western']90 91def _genres_for_movie_row(row):92 return [g for g in GENRE_COLS if row[g] == 1]93 94def make_enriched_label(m_id):95 """Return a richer label like 'Toy Story (1995) — Animation · Children · Comedy'."""96 m_id = str(m_id)97 title = movie_id_to_title.get(m_id, f"Movie {m_id}")98 row = movies[movies['item_id'].astype(str) == m_id]99 if row.empty:100 return title101 genres = _genres_for_movie_row(row.iloc[0])102 if not genres:103 return title104 return f"{title} — {' · '.join(genres)}"105 106# Maps used by the UI for converting between rich labels and movie IDs107movie_id_to_label = {m_id: make_enriched_label(m_id) for m_id in movie_id_to_title}108label_to_movie_id = {label: m_id for m_id, label in movie_id_to_label.items()}109# -------------------------------------------------------------------------110 111# Popular movies (for cold start initial recommendation)112movie_popularity = df.groupby('item_id').size().sort_values(ascending=False)113popular_movie_ids = movie_popularity.index.tolist()114 115# =====================================================================116# 3. Fix model and split test set117# =====================================================================118from surprise import KNNWithMeans, accuracy119from surprise.model_selection import train_test_split120import pandas as pd121import numpy as np122import matplotlib.pyplot as plt123 124# Reproducible train/test split125trainset, testset = train_test_split(data, test_size=0.2, random_state=42)126 127sim_options = {128 "name": "cosine",129 "user_based": False # Item-based KNN130}131 132algo = KNNWithMeans(133 k=40,134 min_k=1,135 sim_options=sim_options136)137 138algo.fit(trainset)139base_predictions = algo.test(testset)140 141# =====================================================================142# 4. NDCG@K function143# =====================================================================144def ndcg_at_k_from_predictions(predictions, k=10):145 """Compute NDCG@K from a Surprise prediction list."""146 pred_df = pd.DataFrame([147 {148 "user_id": pred.uid,149 "item_id": pred.iid,150 "true_rating": pred.r_ui,151 "pred_rating": pred.est152 }153 for pred in predictions154 ])155 156 user_ndcgs = []157 158 for user_id, group in pred_df.groupby("user_id"):159 if len(group) < 2:160 continue161 162 ranked = group.sort_values("pred_rating", ascending=False).head(k)163 ideal = group.sort_values("true_rating", ascending=False).head(k)164 165 dcg = 0.0166 for rank, true_rating in enumerate(ranked["true_rating"].values, start=1):167 gain = true_rating168 discount = np.log2(rank + 1)169 dcg += gain / discount170 171 idcg = 0.0172 for rank, true_rating in enumerate(ideal["true_rating"].values, start=1):173 gain = true_rating174 discount = np.log2(rank + 1)175 idcg += gain / discount176 177 if idcg > 0:178 user_ndcgs.append(dcg / idcg)179 180 if len(user_ndcgs) == 0:181 return np.nan182 183 return np.mean(user_ndcgs)184 185# ==========================================186# 5. Baseline test set performance187# ==========================================188 189K = 10190 191base_rmse = accuracy.rmse(base_predictions, verbose=False)192base_ndcg = ndcg_at_k_from_predictions(base_predictions, k=K)193 194print("Baseline MovieLens test set performance")195print(f"RMSE : {base_rmse:.4f}")196print(f"NDCG@{K} : {base_ndcg:.4f}")197 198# ==========================================199# 6. Recommendation logic (Cold Start)200# ==========================================201 202def get_random_popular(exclude_ids, n=5):203 """Pick n random movies from the popular pool."""204 candidates = [m for m in popular_movie_ids[:200] if m not in exclude_ids]205 return random.sample(candidates, n)206 207def get_related_movies(selected_movie_ids, exclude_ids, n=5):208 """Pick movies whose genres overlap most with the last selected one."""209 last_movie_id = selected_movie_ids[-1]210 last_movie_genres = movies[movies['item_id'].astype(str) == last_movie_id].iloc[0, 6:].values211 212 sim_scores = []213 for idx, row in movies.iterrows():214 m_id = str(row['item_id'])215 if m_id in exclude_ids:216 continue217 m_genres = row[6:].values218 score = np.dot(last_movie_genres, m_genres)219 sim_scores.append((m_id, score))220 221 sim_scores.sort(key=lambda x: x[1], reverse=True)222 return [x[0] for x in sim_scores[:n]]223 224# =====================================================================225# 7. Gradio web interface: Cold-Start Survey + CSV save226# =====================================================================227 228import os229import csv230import gradio as gr231from datetime import datetime232 233SAVE_PATH = "cold_start_selected_movies.csv"234TARGET_NUM_SELECTED = 5235NUM_OPTIONS = 5236 237 238class RecommenderState:239 def __init__(self):240 self.user_name = ""241 self.selected_ids = []242 self.shown_ids = []243 self.current_options = [] # list of enriched labels244 self.finished = False245 246 247def make_initial_state():248 """Create a fresh state for a new user session."""249 state = RecommenderState()250 initial_ids = get_random_popular([], NUM_OPTIONS)251 state.shown_ids.extend(initial_ids)252 state.current_options = [movie_id_to_label[m_id] for m_id in initial_ids]253 return state254 255 256def selected_text(state):257 """Render the 'selected so far' line as Markdown."""258 selected_titles = [movie_id_to_title[m_id] for m_id in state.selected_ids]259 260 if selected_titles:261 return (262 f"**🍿 Selected so far ({len(state.selected_ids)}/{TARGET_NUM_SELECTED}):** "263 + ", ".join(selected_titles)264 )265 return f"**🍿 Selected so far (0/{TARGET_NUM_SELECTED}):** none"266 267 268def save_selected_movies_to_csv(state):269 """Save the user name and their selected movies to a CSV file."""270 file_exists = os.path.exists(SAVE_PATH)271 272 selected_titles = [movie_id_to_title[m_id] for m_id in state.selected_ids[:TARGET_NUM_SELECTED]]273 selected_ids = state.selected_ids[:TARGET_NUM_SELECTED]274 275 row = {276 "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),277 "user_name": state.user_name,278 }279 280 for i, (m_id, title) in enumerate(zip(selected_ids, selected_titles), start=1):281 row[f"movie_{i}_id"] = m_id282 row[f"movie_{i}_title"] = title283 284 fieldnames = ["timestamp", "user_name"]285 for i in range(1, TARGET_NUM_SELECTED + 1):286 fieldnames.extend([f"movie_{i}_id", f"movie_{i}_title"])287 288 with open(SAVE_PATH, "a", newline="", encoding="utf-8-sig") as f:289 writer = csv.DictWriter(f, fieldnames=fieldnames)290 if not file_exists:291 writer.writeheader()292 writer.writerow(row)293 294 return SAVE_PATH295 296 297def start_survey(user_name):298 """After the user enters their name, start the survey."""299 user_name = user_name.strip()300 301 if not user_name:302 return (303 gr.update(visible=True),304 gr.update(visible=False),305 None,306 gr.update(value="⚠️ Please enter your name first.", visible=True),307 gr.update(),308 gr.update(),309 gr.update(),310 gr.update(visible=False),311 gr.update(visible=False),312 gr.update(visible=False),313 )314 315 state = make_initial_state()316 state.user_name = user_name317 318 return (319 gr.update(visible=False),320 gr.update(visible=True),321 state,322 gr.update(value=f"User: **{user_name}**", visible=True),323 gr.update(324 choices=state.current_options,325 value=[],326 label=f"Pick the movies you like ({len(state.selected_ids)}/{TARGET_NUM_SELECTED} selected)",327 visible=True,328 ),329 gr.update(value=selected_text(state), visible=True),330 gr.update(value="", visible=False),331 gr.update(visible=True),332 gr.update(visible=True),333 gr.update(visible=False),334 )335 336 337def finish_survey(state):338 """Once 5 selections are made: save CSV and show finish screen."""339 state.finished = True340 csv_path = save_selected_movies_to_csv(state)341 342 chosen_titles = [movie_id_to_title[m_id] for m_id in state.selected_ids[:TARGET_NUM_SELECTED]]343 344 result_lines = []345 result_lines.append("### ✅ Selection complete and saved")346 result_lines.append("")347 result_lines.append(f"User name: **{state.user_name}**")348 result_lines.append("")349 result_lines.append("Selected movies:")350 351 for i, title in enumerate(chosen_titles, start=1):352 result_lines.append(f"{i}. {title}")353 354 result_lines.append("")355 result_lines.append(f"Saved to CSV file: `{os.path.abspath(csv_path)}`")356 result_lines.append("")357 result_lines.append("You can download the CSV file from the download area below.")358 359 result_md = "\n".join(result_lines)360 361 return (362 state,363 gr.update(visible=False),364 gr.update(value=selected_text(state), visible=True),365 gr.update(value=result_md, visible=True),366 gr.update(visible=False),367 gr.update(visible=True),368 gr.update(value=csv_path, visible=True),369 )370 371 372def next_movies(choices, state):373 """Record current selections and show the next 5 movies."""374 if state is None:375 state = make_initial_state()376 377 if state.finished:378 return (379 state,380 gr.update(),381 gr.update(value=selected_text(state)),382 gr.update(value="Selection is already complete.", visible=True),383 gr.update(visible=False),384 gr.update(visible=True),385 gr.update(visible=False),386 )387 388 # 1. Record what the user selected on this screen389 # `choices` are enriched labels; map them back to movie IDs.390 for label in choices:391 m_id = label_to_movie_id.get(label)392 if m_id is None:393 # Safety net: maybe a plain title slipped through394 m_id = title_to_movie_id.get(label)395 if m_id is not None and m_id not in state.selected_ids:396 state.selected_ids.append(m_id)397 398 # 2. If we have enough selections, finish the survey399 if len(state.selected_ids) >= TARGET_NUM_SELECTED:400 return finish_survey(state)401 402 # 3. Pick the next batch of movies403 if len(state.selected_ids) > 0:404 next_ids = get_related_movies(state.selected_ids, state.shown_ids, NUM_OPTIONS)405 else:406 next_ids = get_random_popular(state.shown_ids, NUM_OPTIONS)407 408 # If we still don't have enough, top up from the popular pool409 if len(next_ids) < NUM_OPTIONS:410 fallback_ids = get_random_popular(state.shown_ids + next_ids, NUM_OPTIONS - len(next_ids))411 next_ids.extend(fallback_ids)412 413 state.shown_ids.extend(next_ids)414 state.current_options = [movie_id_to_label[m_id] for m_id in next_ids]415 416 return (417 state,418 gr.update(419 choices=state.current_options,420 value=[],421 label=f"Pick the movies you like ({len(state.selected_ids)}/{TARGET_NUM_SELECTED} selected)",422 visible=True,423 ),424 gr.update(value=selected_text(state), visible=True),425 gr.update(value="", visible=False),426 gr.update(visible=True),427 gr.update(visible=False),428 gr.update(visible=False),429 )430 431 432def reset_app():433 """Return to the initial screen."""434 return (435 gr.update(visible=True),436 gr.update(value="", visible=True),437 gr.update(visible=False),438 None,439 gr.update(value="", visible=False),440 gr.update(choices=[], value=[], visible=True),441 gr.update(value=f"**🍿 Selected so far (0/{TARGET_NUM_SELECTED}):** none", visible=True),442 gr.update(value="", visible=False),443 gr.update(visible=False),444 gr.update(visible=False),445 gr.update(visible=False),446 )447 448 449with gr.Blocks(theme=gr.themes.Soft()) as demo:450 gr.Markdown("# 🎬 MovieLens Cold-Start Preference Collector")451 gr.Markdown(452 "You'll first be shown 5 random movies. Tick the ones you like, "453 "or click **Next** to see new movies if none appeal to you. "454 f"After selecting **{TARGET_NUM_SELECTED}** movies in total, your choices will be saved to a CSV file."455 )456 457 app_state = gr.State(value=None)458 459 with gr.Column(visible=True) as name_col:460 user_name_box = gr.Textbox(label="Your name", placeholder="e.g., KimSuyeon")461 btn_start = gr.Button("Start")462 463 with gr.Column(visible=False) as survey_col:464 user_display = gr.Markdown(visible=False)465 movie_choices = gr.CheckboxGroup(466 choices=[],467 label=f"Pick the movies you like (0/{TARGET_NUM_SELECTED} selected)",468 )469 selected_display = gr.Markdown(f"**🍿 Selected so far (0/{TARGET_NUM_SELECTED}):** none")470 btn_next = gr.Button("Confirm selection & show next movies (Next)")471 472 result_output = gr.Markdown(visible=False)473 csv_file = gr.File(label="Download saved CSV file", visible=False, file_count="single")474 475 with gr.Row():476 btn_reset = gr.Button("Start over", visible=False)477 478 btn_start.click(479 start_survey,480 inputs=[user_name_box],481 outputs=[482 name_col,483 survey_col,484 app_state,485 user_display,486 movie_choices,487 selected_display,488 result_output,489 btn_next,490 btn_reset,491 csv_file,492 ],493 )494 495 btn_next.click(496 next_movies,497 inputs=[movie_choices, app_state],498 outputs=[499 app_state,500 movie_choices,501 selected_display,502 result_output,503 btn_next,504 btn_reset,505 csv_file,506 ],507 )508 509 btn_reset.click(510 reset_app,511 outputs=[512 name_col,513 user_name_box,514 survey_col,515 app_state,516 user_display,517 movie_choices,518 selected_display,519 result_output,520 btn_next,521 csv_file,522 btn_reset,523 ],524 )525 526demo.launch(share=True, debug=True)