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: st.sidebar.success("Спасибо, каждый день мы стараемся быть лучше для вас🤗") if __name__ == "__main__": main()