app.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. import streamlit as st
  2. import pandas as pd
  3. from main import init_qdrant, load_data, input_ratings, recommend_movies
  4. # Initialize Qdrant and load data
  5. qdrant = init_qdrant()
  6. load_data(qdrant)
  7. # Load movies
  8. movies = pd.read_csv("./data/ml-latest-small/movies.csv")
  9. ratings = pd.read_csv("./data/ml-latest-small/ratings.csv")
  10. # Initialize session state to store user ratings
  11. if "user_ratings" not in st.session_state:
  12. st.session_state["user_ratings"] = {}
  13. # Streamlit app interface
  14. st.title("Movie Recommendation System")
  15. # Movie selection and rating
  16. movie_titles = movies["title"].tolist()
  17. # Movie search and selection using multiselect
  18. selected_movies = st.multiselect("Search and select movies to rate", movie_titles)
  19. if selected_movies:
  20. for movie in selected_movies:
  21. rating = st.slider(f"Rate {movie}", 0.0, 5.0, 0.0, 0.5)
  22. if st.button(f"Add {movie}"):
  23. movie_id = movies[movies.title == movie].movieId.iloc[0]
  24. st.session_state["user_ratings"][movie] = (movie_id, rating)
  25. st.write(f"Added: {movie} with a rating of {rating}")
  26. else:
  27. st.write("Select movies to rate from the dropdown.")
  28. # Clear button to reset all inputs
  29. if st.button("Clear Selections"):
  30. st.session_state["user_ratings"] = {}
  31. st._set_query_params() # Reset the app state
  32. # Display current ratings
  33. if st.session_state["user_ratings"]:
  34. st.write("Current Movie Ratings:")
  35. for movie, (movie_id, rating) in st.session_state["user_ratings"].items():
  36. st.write(f"{movie}: {rating}")
  37. # Get recommendations
  38. if st.button("Get Recommendations"):
  39. if st.session_state["user_ratings"]:
  40. final_ratings = input_ratings(st.session_state["user_ratings"], ratings)
  41. recommendations = recommend_movies(qdrant, movies, final_ratings)
  42. if recommendations:
  43. st.header("Recommended Movies for You")
  44. for movie in recommendations:
  45. st.write(movie)
  46. else:
  47. st.info("No recommendations found based on your ratings.")
  48. else:
  49. st.warning("Please rate at least one movie to get recommendations.")