elozano's picture
remove cached
3078098
raw
history blame contribute delete
No virus
2.15 kB
import numpy as np
import pandas as pd
import streamlit as st
from utils import get_index, get_movie_info
MAX_OVERVIEW_LENGTH = 280
IMAGE_WIDTH = 170
MIN_RECOMMENDATIONS = 1
MAX_RECOMMENDATIONS = 100
DEFAULT_RECOMMENDATIONS = 10
def init() -> None:
movies = pd.read_csv("movies.csv")
st.session_state["movies"] = movies
st.session_state["titles"] = movies["original_title"].values
st.session_state["sim_matrix"] = np.load("sim.npy")
def run() -> None:
st.title("Movie Recommender")
with st.form("search"):
title = st.selectbox("Recommend movies based on:", st.session_state["titles"])
num_recommendations = st.number_input(
"Number of recommendations:",
min_value=MIN_RECOMMENDATIONS,
max_value=MAX_RECOMMENDATIONS,
value=DEFAULT_RECOMMENDATIONS,
)
search = st.form_submit_button("Search")
if search:
_recommend(title, num_recommendations)
def _recommend(based_on: str, num_recommendations: int) -> None:
movie_index = get_index(based_on)
similarities = list(enumerate(st.session_state["sim_matrix"][movie_index]))
similarities.pop(movie_index)
sorted_similarities = sorted(similarities, key=lambda x: x[1], reverse=True)
recommendations = sorted_similarities[:num_recommendations]
st.markdown("---")
for idx, _ in recommendations:
_show_movie_info(idx)
def _show_movie_info(movie_index: str) -> None:
info = get_movie_info(movie_index)
im_col, info_col = st.columns([2, 4])
im_col.image(
f"https://a.ltrbxd.com/resized/{info['image_url']}.jpg", width=IMAGE_WIDTH
)
info_col.markdown(
f"### {info['original_title']} ({info['release_date'].split('-')[0]})"
)
info_col.markdown(f"*{int(info['runtime'])} min*")
short_overview = (
f"{info['overview'][:MAX_OVERVIEW_LENGTH]} ..."
if len(info["overview"]) > MAX_OVERVIEW_LENGTH
else info["overview"]
)
info_col.markdown(short_overview)
info_col.caption(", ".join(eval(info["genres"])))
st.markdown("---")
if __name__ == "__main__":
init()
run()