movie_rec / app.py
HaggiVaggi's picture
Update app.py
059bb54
raw
history blame
6.79 kB
import streamlit as st
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModel
import faiss
import numpy as np
st.title('Умный поиск фильмов 🔍🎦')
with st.sidebar:
st.header('Выберите страницу')
page = st.selectbox("Выберите страницу", ["Главная", "Подбор фильмов по описанию", "Итоги"])
if page == "Главная":
st.header('Выполнила команда "FindMyMovie":')
st.subheader('🎥Алексей')
st.subheader('🎬Светлана')
st.subheader('🍿Тата')
st.subheader('Наши задачи:')
st.markdown("""
<span style='font-size:18px; color:purple;'>Задача 1:</span> Спарсить информацию по 5 тыс. фильмов и обработать ее
<span style='font-size:18px; color:purple;'>Задача 2:</span> Разработать систему поиска фильма по пользовательскому запросу
""", unsafe_allow_html=True)
st.markdown(
f"<div style='border: 2px solid #800080; padding: 10px; text-align: center;'><span style='font-size: 20px; color: violet;'>Мы любезно позаимствовали данные о фильмах с \
Киноафиши, ни один участник команды при этом не пострадал 💟</span></div>",
unsafe_allow_html=True
)
st.subheader(' '*10)
st.info('🌟Сервис принимает на вход описание фильма от пользователя и возвращает заданное количество подходящих вариантов')
df = pd.read_csv('fina_data.csv')
st.info('🌟Если не знаете, что посмотреть, нажмите кнопку "Сгенерировать" - сервис предложит вам 10 случайных фильмов из библиотеки')
# Отображаем HTML-разметку в Streamlit
if st.button("Сгенерировать 🎲"):
# Получение случайных 10 строк
random_rows = df[['movie_title', 'description']].sample(n=10).reset_index(drop=True)
random_rows.index = random_rows.index + 1
st.markdown(f"<span style='font-size:{20}px; color:violet'>{'Сегодня мы подобрали для вас следующие фильмы:'}</span>", unsafe_allow_html=True)
st.write(random_rows)
st.image("1.png", use_column_width=True)
if page == "Подбор фильмов по описанию":
# Загрузка предварительно обученной модели ruBERT
tokenizer = AutoTokenizer.from_pretrained("DeepPavlov/rubert-base-cased-sentence")
model = AutoModel.from_pretrained("DeepPavlov/rubert-base-cased-sentence")
df2 = pd.read_csv('fina_data.csv')
embeddings_array = np.load('embeddings_final.npy')
index = faiss.read_index('desc_faiss_index_final.index')
def encode_description(description):
tokens = tokenizer(description, return_tensors="pt")
with torch.no_grad():
outputs = model(**tokens)
embeddings = outputs.last_hidden_state.mean(dim=1)
return embeddings.cpu().numpy().astype('float32')
# embeddings = pd.read_pickle('embeddings.pkl')
def main():
# Пользовательский ввод
user_input = st.text_input("Введите описание фильма:", value="", help="Чем подробнее будет ваше описание, тем точнее мы сможем подобрать для вас фильм 🤗'")
if st.button("Искать🔍🎦"):
if user_input:
# Векторизация введенного запроса
input_embedding = encode_description(user_input)
# Поиск с использованием Faiss
_, sorted_indices = index.search(input_embedding.reshape(1, -1), 5) # Изменил на 5
# Используйте индексы для извлечения строк из DataFrame
recs = df2.iloc[sorted_indices[0]].reset_index(drop=True)
recs.index = recs.index + 1
# Вывод рекомендованных фильмов с изображениями
st.subheader("Рекомендованные фильмы 🎉:")
for i in range(5):
st.markdown(f"<span style='font-size:{20}px; color:purple'>{recs['movie_title'].iloc[i]}</span>", unsafe_allow_html=True)
# Создаем две колонки: одну для текста, другую для изображения
col1, col2 = st.columns([2, 1])
# В колонке отображаем название фильма, описание, роли и ссылку
col1.info(recs['description'].iloc[i])
col1.markdown(f"**В ролях:** {recs['actors'].iloc[i]}")
col1.markdown(f"**Фильм можно посмотреть [здесь]({recs['page_url'].iloc[i]})**")
# В колонке отображаем изображение
col2.image(recs['image_url'].iloc[i], caption=recs['movie_title'].iloc[i], width=200)
with st.sidebar:
st.info("""
#### Понравились ли Вам наши рекомендации?
""")
col1, col2 = st.columns(2)
with col1:
st.button('**Да, очень**🔥', type='primary')
with col2:
st.button('**Нет,можно лучше**👎🏻', type='primary')
# with st.sidebar:
# feedback = st.text_input("Мы смогли помочь вам с выбором?", key="feedback_input")
# feedback_button = st.button("Отправить отзыв", key="feedback_button")
# # Обработка отзыва, например, сохранение в файл или отправка на сервер
# if feedback_button:
# st.sidebar.success("Спасибо, каждый день мы стараемся быть лучше для вас🤗")
if __name__ == "__main__":
main()