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("""
Задача 1: Спарсить информацию по 5 тыс. фильмов и обработать ее
Задача 2: Разработать систему поиска фильма по пользовательскому запросу
""", unsafe_allow_html=True)
st.markdown(
f"
Мы любезно позаимствовали данные о фильмах с \
Киноафиши, ни один участник команды при этом не пострадал 💟
",
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"{'Сегодня мы подобрали для вас следующие фильмы:'}", 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"{recs['movie_title'].iloc[i]}", 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("""
#### Мы смогли помочь вам с выбором?
""")
feedback = st.text_input('Поделитесь с нами вашим мнением')
feedback_button = st.button("Отправить отзыв", key="feedback_button")
if feedback_button and feedback:
st.success("Спасибо, каждый день мы стараемся быть лучше для вас 💟")
if __name__ == "__main__":
main()