12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061 |
- import streamlit as st
- import pandas as pd
- from main import init_qdrant, load_data, input_ratings, recommend_movies
- # Initialize Qdrant and load data
- qdrant = init_qdrant()
- load_data(qdrant)
- # Load movies
- movies = pd.read_csv("./data/ml-latest-small/movies.csv")
- ratings = pd.read_csv("./data/ml-latest-small/ratings.csv")
- # Initialize session state to store user ratings
- if "user_ratings" not in st.session_state:
- st.session_state["user_ratings"] = {}
- # Streamlit app interface
- st.title("Movie Recommendation System")
- # Movie selection and rating
- movie_titles = movies["title"].tolist()
- # Movie search and selection using multiselect
- selected_movies = st.multiselect("Search and select movies to rate", movie_titles)
- if selected_movies:
- for movie in selected_movies:
- rating = st.slider(f"Rate {movie}", 0.0, 5.0, 0.0, 0.5)
- if st.button(f"Add {movie}"):
- movie_id = movies[movies.title == movie].movieId.iloc[0]
- st.session_state["user_ratings"][movie] = (movie_id, rating)
- st.write(f"Added: {movie} with a rating of {rating}")
- else:
- st.write("Select movies to rate from the dropdown.")
- # Clear button to reset all inputs
- if st.button("Clear Selections"):
- st.session_state["user_ratings"] = {}
- st._set_query_params() # Reset the app state
-
- # Display current ratings
- if st.session_state["user_ratings"]:
- st.write("Current Movie Ratings:")
- for movie, (movie_id, rating) in st.session_state["user_ratings"].items():
- st.write(f"{movie}: {rating}")
- # Get recommendations
- if st.button("Get Recommendations"):
- if st.session_state["user_ratings"]:
- final_ratings = input_ratings(st.session_state["user_ratings"], ratings)
- recommendations = recommend_movies(qdrant, movies, final_ratings)
- if recommendations:
- st.header("Recommended Movies for You")
- for movie in recommendations:
- st.write(movie)
- else:
- st.info("No recommendations found based on your ratings.")
- else:
- st.warning("Please rate at least one movie to get recommendations.")
|