jiminbae/coldstarter
0
1# -*- coding: utf-8 -*-2 3from __future__ import annotations4 5import csv6import os7import random8from dataclasses import dataclass, field9from datetime import datetime10from typing import Dict, List11 12import gradio as gr13import pandas as pd14 15# =========================================================16# 1. Lightweight MovieLens-style movie metadata17# =========================================================18# The original base code loads MovieLens 100K through scikit-surprise.19# For Hugging Face Spaces, this version removes scikit-surprise and keeps20# a compact MovieLens-style catalog directly in the app for fast deployment.21 22MOVIES = [23 {"id": "50", "title": "Star Wars (1977)", "genres": ["Action", "Adventure", "Romance", "Sci-Fi", "War"]},24 {"id": "100", "title": "Fargo (1996)", "genres": ["Crime", "Drama", "Thriller"]},25 {"id": "181", "title": "Return of the Jedi (1983)", "genres": ["Action", "Adventure", "Romance", "Sci-Fi", "War"]},26 {"id": "258", "title": "Contact (1997)", "genres": ["Drama", "Sci-Fi"]},27 {"id": "294", "title": "Liar Liar (1997)", "genres": ["Comedy"]},28 {"id": "286", "title": "English Patient, The (1996)", "genres": ["Drama", "Romance", "War"]},29 {"id": "288", "title": "Scream (1996)", "genres": ["Horror", "Thriller"]},30 {"id": "1", "title": "Toy Story (1995)", "genres": ["Animation", "Children", "Comedy"]},31 {"id": "300", "title": "Air Force One (1997)", "genres": ["Action", "Thriller"]},32 {"id": "121", "title": "Independence Day (ID4) (1996)", "genres": ["Action", "Sci-Fi", "War"]},33 {"id": "174", "title": "Raiders of the Lost Ark (1981)", "genres": ["Action", "Adventure"]},34 {"id": "127", "title": "Godfather, The (1972)", "genres": ["Action", "Crime", "Drama"]},35 {"id": "56", "title": "Pulp Fiction (1994)", "genres": ["Crime", "Drama"]},36 {"id": "98", "title": "Silence of the Lambs, The (1991)", "genres": ["Drama", "Thriller"]},37 {"id": "7", "title": "Twelve Monkeys (1995)", "genres": ["Drama", "Sci-Fi"]},38 {"id": "237", "title": "Jerry Maguire (1996)", "genres": ["Drama", "Romance"]},39 {"id": "117", "title": "Rock, The (1996)", "genres": ["Action", "Adventure", "Thriller"]},40 {"id": "172", "title": "Empire Strikes Back, The (1980)", "genres": ["Action", "Adventure", "Drama", "Romance", "Sci-Fi", "War"]},41 {"id": "222", "title": "Star Trek: First Contact (1996)", "genres": ["Action", "Adventure", "Sci-Fi"]},42 {"id": "313", "title": "Titanic (1997)", "genres": ["Action", "Drama", "Romance"]},43 {"id": "204", "title": "Back to the Future (1985)", "genres": ["Comedy", "Sci-Fi"]},44 {"id": "405", "title": "Mission: Impossible (1996)", "genres": ["Action", "Adventure", "Mystery"]},45 {"id": "79", "title": "Fugitive, The (1993)", "genres": ["Action", "Thriller"]},46 {"id": "210", "title": "Indiana Jones and the Last Crusade (1989)", "genres": ["Action", "Adventure"]},47 {"id": "151", "title": "Willy Wonka and the Chocolate Factory (1971)", "genres": ["Adventure", "Children", "Comedy", "Fantasy"]},48 {"id": "173", "title": "Princess Bride, The (1987)", "genres": ["Action", "Adventure", "Comedy", "Romance"]},49 {"id": "69", "title": "Forrest Gump (1994)", "genres": ["Comedy", "Romance", "War"]},50 {"id": "168", "title": "Monty Python and the Holy Grail (1974)", "genres": ["Comedy"]},51 {"id": "269", "title": "Full Monty, The (1997)", "genres": ["Comedy"]},52 {"id": "257", "title": "Men in Black (1997)", "genres": ["Action", "Adventure", "Comedy", "Sci-Fi"]},53 {"id": "318", "title": "Schindler's List (1993)", "genres": ["Drama", "War"]},54 {"id": "302", "title": "L.A. Confidential (1997)", "genres": ["Crime", "Film-Noir", "Mystery", "Thriller"]},55 {"id": "313", "title": "Titanic (1997)", "genres": ["Action", "Drama", "Romance"]},56 {"id": "22", "title": "Braveheart (1995)", "genres": ["Action", "Drama", "War"]},57 {"id": "96", "title": "Terminator 2: Judgment Day (1991)", "genres": ["Action", "Sci-Fi", "Thriller"]},58 {"id": "15", "title": "Mr. Holland's Opus (1995)", "genres": ["Drama"]},59 {"id": "176", "title": "Aliens (1986)", "genres": ["Action", "Sci-Fi", "Thriller", "War"]},60 {"id": "28", "title": "Apollo 13 (1995)", "genres": ["Action", "Drama", "Thriller"]},61 {"id": "195", "title": "Terminator, The (1984)", "genres": ["Action", "Sci-Fi", "Thriller"]},62 {"id": "423", "title": "E.T. the Extra-Terrestrial (1982)", "genres": ["Children", "Drama", "Fantasy", "Sci-Fi"]},63 {"id": "64", "title": "Shawshank Redemption, The (1994)", "genres": ["Drama"]},64 {"id": "12", "title": "Usual Suspects, The (1995)", "genres": ["Crime", "Thriller"]},65 {"id": "483", "title": "Casablanca (1942)", "genres": ["Drama", "Romance", "War"]},66 {"id": "603", "title": "Rear Window (1954)", "genres": ["Mystery", "Thriller"]},67 {"id": "132", "title": "Wizard of Oz, The (1939)", "genres": ["Adventure", "Children", "Drama", "Musical"]},68 {"id": "89", "title": "Blade Runner (1982)", "genres": ["Film-Noir", "Sci-Fi"]},69 {"id": "183", "title": "Alien (1979)", "genres": ["Action", "Horror", "Sci-Fi", "Thriller"]},70 {"id": "202", "title": "Groundhog Day (1993)", "genres": ["Comedy", "Romance"]},71 {"id": "208", "title": "Young Frankenstein (1974)", "genres": ["Comedy", "Horror"]},72 {"id": "216", "title": "When Harry Met Sally... (1989)", "genres": ["Comedy", "Romance"]},73 {"id": "97", "title": "Dances with Wolves (1990)", "genres": ["Adventure", "Drama", "Western"]},74 {"id": "144", "title": "Die Hard (1988)", "genres": ["Action", "Thriller"]},75 {"id": "187", "title": "Godfather: Part II, The (1974)", "genres": ["Action", "Crime", "Drama"]},76 {"id": "194", "title": "Sting, The (1973)", "genres": ["Comedy", "Crime"]},77 {"id": "196", "title": "Dead Poets Society (1989)", "genres": ["Drama"]},78 {"id": "200", "title": "Shining, The (1980)", "genres": ["Horror"]},79 {"id": "203", "title": "Unforgiven (1992)", "genres": ["Western"]},80 {"id": "211", "title": "M*A*S*H (1970)", "genres": ["Comedy", "War"]},81 {"id": "218", "title": "Cape Fear (1991)", "genres": ["Thriller"]},82 {"id": "234", "title": "Jaws (1975)", "genres": ["Action", "Horror"]},83 {"id": "238", "title": "Raising Arizona (1987)", "genres": ["Comedy"]},84 {"id": "276", "title": "Leaving Las Vegas (1995)", "genres": ["Drama", "Romance"]},85 {"id": "357", "title": "One Flew Over the Cuckoo's Nest (1975)", "genres": ["Drama"]},86 {"id": "427", "title": "To Kill a Mockingbird (1962)", "genres": ["Drama"]},87 {"id": "480", "title": "North by Northwest (1959)", "genres": ["Action", "Thriller"]},88 {"id": "496", "title": "It's a Wonderful Life (1946)", "genres": ["Drama"]},89 {"id": "498", "title": "African Queen, The (1951)", "genres": ["Action", "Adventure", "Romance", "War"]},90 {"id": "509", "title": "My Left Foot (1989)", "genres": ["Drama"]},91 {"id": "520", "title": "Great Escape, The (1963)", "genres": ["Adventure", "War"]},92 {"id": "521", "title": "Deer Hunter, The (1978)", "genres": ["Drama", "War"]},93 {"id": "528", "title": "Killing Fields, The (1984)", "genres": ["Drama", "War"]},94 {"id": "528", "title": "Killing Fields, The (1984)", "genres": ["Drama", "War"]},95 {"id": "531", "title": "Shine (1996)", "genres": ["Drama", "Romance"]},96 {"id": "603", "title": "Rear Window (1954)", "genres": ["Mystery", "Thriller"]},97 {"id": "651", "title": "Glory (1989)", "genres": ["Action", "Drama", "War"]},98 {"id": "657", "title": "Manchurian Candidate, The (1962)", "genres": ["Film-Noir", "Thriller"]},99 {"id": "705", "title": "Singin' in the Rain (1952)", "genres": ["Musical", "Romance"]},100 {"id": "742", "title": "Ransom (1996)", "genres": ["Crime", "Thriller"]},101 {"id": "748", "title": "Saint, The (1997)", "genres": ["Action", "Romance", "Thriller"]},102 {"id": "879", "title": "Peacemaker, The (1997)", "genres": ["Action", "Thriller", "War"]},103]104 105# Remove accidental duplicate movie ids while preserving order.106seen = set()107MOVIES = [m for m in MOVIES if not (m["id"] in seen or seen.add(m["id"]))]108MOVIE_DF = pd.DataFrame(MOVIES)109TITLE_TO_MOVIE: Dict[str, dict] = {m["title"]: m for m in MOVIES}110ID_TO_MOVIE: Dict[str, dict] = {m["id"]: m for m in MOVIES}111 112SAVE_PATH = "cold_start_selected_movies.csv"113TARGET_NUM_SELECTED = 5114NUM_OPTIONS = 5115POPULAR_POOL_SIZE = min(40, len(MOVIES))116 117APP_CSS = """118.gradio-container {119 max-width: 980px !important;120 margin: auto !important;121}122#title-card {123 background: linear-gradient(135deg, #fff7ed 0%, #fef2f2 45%, #eef2ff 100%);124 border: 1px solid #fed7aa;125 border-radius: 24px;126 padding: 28px;127 box-shadow: 0 12px 28px rgba(15, 23, 42, 0.08);128}129#title-card h1 {130 font-size: 2.2rem;131 margin-bottom: 0.4rem;132 color: #111827 !important;133}134#title-card h3, #title-card p, #title-card strong {135 color: #334155 !important;136}137#subtitle {138 font-size: 1rem;139 color: #475569;140}141#status-box {142 border-radius: 18px;143 padding: 16px;144 background: #f8fafc;145 border: 1px solid #e2e8f0;146}147.movie-card {148 padding: 12px 16px;149 border-radius: 16px;150 border: 1px solid #e5e7eb;151 background: #ffffff;152}153button.primary-btn {154 border-radius: 999px !important;155}156"""157 158# =========================================================159# 2. State and recommendation logic160# =========================================================161 162@dataclass163class RecommenderState:164 user_name: str = ""165 selected_ids: List[str] = field(default_factory=list)166 shown_ids: List[str] = field(default_factory=list)167 current_options: List[str] = field(default_factory=list)168 finished: bool = False169 170 171def movie_label(movie: dict) -> str:172 return f"{movie['title']} · {', '.join(movie['genres'])}"173 174 175def title_from_label(label: str) -> str:176 return label.split(" · ")[0]177 178 179def get_random_popular(exclude_ids: List[str], n: int = NUM_OPTIONS) -> List[str]:180 """Return n popular movies not shown before."""181 candidates = [m["id"] for m in MOVIES[:POPULAR_POOL_SIZE] if m["id"] not in exclude_ids]182 if len(candidates) < n:183 candidates += [m["id"] for m in MOVIES if m["id"] not in exclude_ids and m["id"] not in candidates]184 return random.sample(candidates, min(n, len(candidates)))185 186 187def genre_overlap_score(source_genres: List[str], target_genres: List[str]) -> int:188 return len(set(source_genres).intersection(set(target_genres)))189 190 191def get_related_movies(selected_movie_ids: List[str], exclude_ids: List[str], n: int = NUM_OPTIONS) -> List[str]:192 """Recommend movies by genre overlap with the most recently selected movie."""193 if not selected_movie_ids:194 return get_random_popular(exclude_ids, n)195 196 last_movie = ID_TO_MOVIE[selected_movie_ids[-1]]197 scored = []198 for movie in MOVIES:199 if movie["id"] in exclude_ids or movie["id"] in selected_movie_ids:200 continue201 score = genre_overlap_score(last_movie["genres"], movie["genres"])202 scored.append((movie["id"], score, random.random()))203 204 scored.sort(key=lambda x: (x[1], x[2]), reverse=True)205 related = [movie_id for movie_id, score, _ in scored if score > 0][:n]206 207 if len(related) < n:208 fallback = get_random_popular(exclude_ids + related + selected_movie_ids, n - len(related))209 related.extend(fallback)210 211 return related[:n]212 213 214def make_initial_state(user_name: str) -> RecommenderState:215 initial_ids = get_random_popular([], NUM_OPTIONS)216 state = RecommenderState(user_name=user_name)217 state.shown_ids.extend(initial_ids)218 state.current_options = [movie_label(ID_TO_MOVIE[m_id]) for m_id in initial_ids]219 return state220 221 222def selected_markdown(state: RecommenderState | None) -> str:223 if state is None or not state.selected_ids:224 return f"### 🍿 Selected movies: 0/{TARGET_NUM_SELECTED}\n아직 선택한 영화가 없습니다."225 226 lines = [f"### 🍿 Selected movies: {len(state.selected_ids)}/{TARGET_NUM_SELECTED}"]227 for idx, movie_id in enumerate(state.selected_ids, start=1):228 movie = ID_TO_MOVIE[movie_id]229 lines.append(f"{idx}. **{movie['title']}** \n <span style='color:#64748b'>Genres: {', '.join(movie['genres'])}</span>")230 return "\n".join(lines)231 232 233def progress_text(state: RecommenderState | None) -> str:234 if state is None:235 return "0%"236 pct = min(100, int(len(state.selected_ids) / TARGET_NUM_SELECTED * 100))237 return f"{pct}% complete"238 239 240def save_selected_movies_to_csv(state: RecommenderState) -> str:241 file_exists = os.path.exists(SAVE_PATH)242 selected_movies = [ID_TO_MOVIE[m_id] for m_id in state.selected_ids[:TARGET_NUM_SELECTED]]243 244 row = {245 "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),246 "user_name": state.user_name,247 "num_selected": len(selected_movies),248 }249 for i, movie in enumerate(selected_movies, start=1):250 row[f"movie_{i}_id"] = movie["id"]251 row[f"movie_{i}_title"] = movie["title"]252 row[f"movie_{i}_genres"] = ", ".join(movie["genres"])253 254 fieldnames = ["timestamp", "user_name", "num_selected"]255 for i in range(1, TARGET_NUM_SELECTED + 1):256 fieldnames.extend([f"movie_{i}_id", f"movie_{i}_title", f"movie_{i}_genres"])257 258 with open(SAVE_PATH, "a", newline="", encoding="utf-8-sig") as f:259 writer = csv.DictWriter(f, fieldnames=fieldnames)260 if not file_exists:261 writer.writeheader()262 writer.writerow(row)263 264 return SAVE_PATH265 266 267def build_result_markdown(state: RecommenderState, csv_path: str) -> str:268 lines = [269 "## ✅ Preference collection completed!",270 f"**User:** {state.user_name}",271 "",272 "### Final selected movies",273 ]274 for idx, movie_id in enumerate(state.selected_ids[:TARGET_NUM_SELECTED], start=1):275 movie = ID_TO_MOVIE[movie_id]276 lines.append(f"{idx}. **{movie['title']}** — {', '.join(movie['genres'])}")277 lines.extend([278 "",279 f"CSV saved as: `{os.path.abspath(csv_path)}`",280 "",281 "아래 파일 영역에서 CSV를 다운로드할 수 있습니다.",282 ])283 return "\n".join(lines)284 285# =========================================================286# 3. Gradio event handlers287# =========================================================288 289def start_survey(user_name: str):290 user_name = (user_name or "").strip()291 if not user_name:292 return (293 gr.update(visible=True),294 gr.update(visible=False),295 None,296 gr.update(value="⚠️ 이름을 먼저 입력해주세요.", visible=True),297 gr.update(choices=[], value=[]),298 gr.update(value=f"### 🍿 Selected movies: 0/{TARGET_NUM_SELECTED}\n아직 선택한 영화가 없습니다."),299 gr.update(value="0%"),300 gr.update(visible=False),301 gr.update(visible=False),302 gr.update(visible=False),303 )304 305 state = make_initial_state(user_name)306 return (307 gr.update(visible=False),308 gr.update(visible=True),309 state,310 gr.update(value=f"### 👋 Welcome, **{user_name}**\n좋아하는 영화를 고르면 다음 목록이 취향에 맞게 조금씩 바뀝니다.", visible=True),311 gr.update(312 choices=state.current_options,313 value=[],314 label=f"Pick movies you like — {len(state.selected_ids)}/{TARGET_NUM_SELECTED} selected",315 visible=True,316 ),317 gr.update(value=selected_markdown(state), visible=True),318 gr.update(value=progress_text(state), visible=True),319 gr.update(value="", visible=False),320 gr.update(visible=True),321 gr.update(visible=False),322 )323 324 325def finish_survey(state: RecommenderState):326 state.finished = True327 csv_path = save_selected_movies_to_csv(state)328 return (329 state,330 gr.update(choices=[], value=[], visible=False),331 gr.update(value=selected_markdown(state), visible=True),332 gr.update(value=progress_text(state), visible=True),333 gr.update(value=build_result_markdown(state, csv_path), visible=True),334 gr.update(visible=False),335 gr.update(value=csv_path, visible=True),336 gr.update(visible=True),337 )338 339 340def next_movies(choices: List[str], state: RecommenderState | None):341 if state is None:342 return (343 None,344 gr.update(),345 gr.update(value=f"### 🍿 Selected movies: 0/{TARGET_NUM_SELECTED}\n먼저 이름을 입력하고 시작해주세요."),346 gr.update(value="0%"),347 gr.update(value="⚠️ 먼저 시작하기 버튼을 눌러주세요.", visible=True),348 gr.update(visible=False),349 gr.update(visible=False),350 gr.update(visible=False),351 )352 353 if state.finished:354 return finish_survey(state)355 356 for label in choices or []:357 title = title_from_label(label)358 movie = TITLE_TO_MOVIE.get(title)359 if movie and movie["id"] not in state.selected_ids:360 state.selected_ids.append(movie["id"])361 362 if len(state.selected_ids) >= TARGET_NUM_SELECTED:363 return finish_survey(state)364 365 next_ids = get_related_movies(state.selected_ids, state.shown_ids, NUM_OPTIONS)366 state.shown_ids.extend(next_ids)367 state.current_options = [movie_label(ID_TO_MOVIE[m_id]) for m_id in next_ids]368 369 hint = ""370 if not choices:371 hint = "마음에 드는 영화가 없어서 새로운 인기 영화/유사 영화를 보여드렸습니다."372 373 return (374 state,375 gr.update(376 choices=state.current_options,377 value=[],378 label=f"Pick movies you like — {len(state.selected_ids)}/{TARGET_NUM_SELECTED} selected",379 visible=True,380 ),381 gr.update(value=selected_markdown(state), visible=True),382 gr.update(value=progress_text(state), visible=True),383 gr.update(value=hint, visible=bool(hint)),384 gr.update(visible=True),385 gr.update(visible=False),386 gr.update(visible=False),387 )388 389 390def reset_app():391 return (392 gr.update(visible=True),393 gr.update(value="", visible=True),394 gr.update(visible=False),395 None,396 gr.update(value="", visible=False),397 gr.update(choices=[], value=[], visible=True),398 gr.update(value=f"### 🍿 Selected movies: 0/{TARGET_NUM_SELECTED}\n아직 선택한 영화가 없습니다.", visible=True),399 gr.update(value="0%", visible=True),400 gr.update(value="", visible=False),401 gr.update(visible=False),402 gr.update(visible=False),403 gr.update(visible=False),404 )405 406 407def download_current_csv():408 if os.path.exists(SAVE_PATH):409 return gr.update(value=SAVE_PATH, visible=True)410 return gr.update(value=None, visible=False)411 412# =========================================================413# 4. Gradio UI414# =========================================================415 416with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange", secondary_hue="rose"), css=APP_CSS) as demo:417 app_state = gr.State(value=None)418 419 with gr.Column(elem_id="title-card"):420 gr.Markdown(421 """422 # 🎬 Cold Start Movie Preference Collector423 ### MovieLens-style onboarding demo for recommender systems424 425 Pick at least **5 movies** you like. The app starts with popular movies,426 then recommends genre-similar movies based on your choices.427 """428 )429 430 with gr.Column(visible=True) as name_col:431 user_name_box = gr.Textbox(432 label="User name",433 placeholder="예: jiminbae",434 info="사용자 이름을 입력하세요.",435 )436 btn_start = gr.Button("🚀 Start preference survey", variant="primary", elem_classes=["primary-btn"])437 438 with gr.Column(visible=False) as survey_col:439 user_display = gr.Markdown(visible=False)440 progress = gr.Textbox(label="Progress", value="0%", interactive=False)441 movie_choices = gr.CheckboxGroup(442 choices=[],443 label=f"Pick movies you like — 0/{TARGET_NUM_SELECTED} selected",444 )445 selected_display = gr.Markdown(446 f"### 🍿 Selected movies: 0/{TARGET_NUM_SELECTED}\n아직 선택한 영화가 없습니다.",447 elem_id="status-box",448 )449 helper_msg = gr.Markdown(visible=False)450 btn_next = gr.Button("✨ Reflect selection & show next movies", variant="primary", elem_classes=["primary-btn"])451 452 result_output = gr.Markdown(visible=False)453 csv_file = gr.File(label="📄 Download saved CSV", visible=False, file_count="single")454 455 with gr.Row():456 btn_reset = gr.Button("🔄 Restart", visible=False)457 btn_download = gr.Button("📥 Show current CSV file")458 459 btn_start.click(460 start_survey,461 inputs=[user_name_box],462 outputs=[463 name_col,464 survey_col,465 app_state,466 user_display,467 movie_choices,468 selected_display,469 progress,470 result_output,471 btn_next,472 csv_file,473 ],474 )475 476 btn_next.click(477 next_movies,478 inputs=[movie_choices, app_state],479 outputs=[480 app_state,481 movie_choices,482 selected_display,483 progress,484 result_output,485 btn_next,486 csv_file,487 btn_reset,488 ],489 )490 491 btn_reset.click(492 reset_app,493 outputs=[494 name_col,495 user_name_box,496 survey_col,497 app_state,498 user_display,499 movie_choices,500 selected_display,501 progress,502 result_output,503 btn_next,504 csv_file,505 btn_reset,506 ],507 )508 509 btn_download.click(510 download_current_csv,511 inputs=[],512 outputs=[csv_file],513 )514 515if __name__ == "__main__":516 demo.launch()517