Spaces:
Sleeping
Sleeping
File size: 2,998 Bytes
d87c6b1 |
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 |
import streamlit as st
import pandas as pd
import torch
import yaml
from embeddings import load_model, compute_embeddings
# Load configuration from YAML file
with open("configs.yaml", "r") as file:
configs = yaml.safe_load(file)
# Load the processed movie dataset
movie_data = pd.read_csv(configs["processed_dataset"])
# Streamlit app
st.title("🎬 EasyRec Movie Recommender")
# Dropdown for model selection
model_names = configs['hf_models'] # Assuming this is a list of model names in your configs.yaml
selected_model_name = st.selectbox("Select a model:", model_names)
# Load the model based on user selection
model, tokenizer = load_model(selected_model_name)
# User input for movie description
user_description = st.text_input("Enter a description of the type of movie you're interested in:",
placeholder="e.g. A romantic comedy with a twist...")
if user_description:
# Load the precomputed movie embeddings from a .pt file
embedding_dir_path = f"{configs['movie_embeddings']}/{selected_model_name}"
embedding_file_path = f"{embedding_dir_path}/{configs['movie_embeddings']}.pt"
movie_embeddings = torch.load(embedding_file_path) # Load the .pt file
# Compute the embedding for the user input by passing it as a list
user_embedding = compute_embeddings([user_description], model, tokenizer)
similarity_scores = torch.matmul(movie_embeddings, user_embedding.T).flatten()
# Set the number of top recommendations to display
K = 5
top_k_indices = torch.argsort(similarity_scores, descending=True)[:K].tolist() # Get indices of top K
# Display recommendations
st.write("## 🎉 Top Recommendations:")
for rank, movie_id in enumerate(top_k_indices, start=1):
movie = movie_data.iloc[movie_id]
# Convert runtime from minutes to hours and minutes
hours = movie.runtime // 60
minutes = movie.runtime % 60
# Construct an HTML card for displaying the movie information
st.markdown(f"### {rank}. {movie.title}")
st.markdown(f"**Release Date:** {movie.release_date} **Runtime:** {f'{hours}h {minutes}m' if hours > 0 else f'{minutes}m'}")
st.markdown(f"⭐ {movie.vote_average} ({movie.vote_count} votes)")
st.markdown(f"**Overview:** {movie.overview}")
st.markdown(f"**Genres:** {movie.genres}")
st.markdown(f"**Production Companies:** {movie.production_companies}")
st.markdown(f"**Production Countries:** {movie.production_countries}")
st.markdown("---")
# Additional styling (optional)
st.markdown(
"""
<style>
.stTextInput > div > input {
background-color: #f0f0f5; /* Light gray background */
border: 1px solid #ccc; /* Gray border */
border-radius: 5px; /* Rounded corners */
}
</style>
""",
unsafe_allow_html=True
)
|