Spaces:
Runtime error
Runtime error
File size: 5,162 Bytes
34e52df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import streamlit as st
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModel
import faiss
import numpy as np
@st.cache_data
def load_data(url):
df = pd.read_csv(url)
return df
@st.cache_data
def embedding_and_index():
embeddings_array = np.load('data/embeddings_final.npy')
index = faiss.read_index('data/desc_faiss_index_final.index')
return embeddings_array, index
@st.cache_data
def load_model():
model = AutoModel.from_pretrained("DeepPavlov/rubert-base-cased-sentence")
return model
st.header("Подбор фильмов по описанию ✏️🔍")
# Загрузка данных
tokenizer = AutoTokenizer.from_pretrained("DeepPavlov/rubert-base-cased-sentence")
df = load_data('data/final_data.csv')
embeddings_array, index = embedding_and_index()
model = load_model()
# Пользовательский ввод
user_input = st.text_input("Введите описание фильма:", value="", help="Чем подробнее будет ваше описание, тем точнее мы сможем подобрать для вас фильм 🤗'")
genre_list = ['анимация', 'аниме', 'балет', 'биография', 'боевик', 'вестерн', 'военный', 'детектив', 'детский', 'документальный', 'драма', 'исторический', 'катастрофа', 'комедия', 'концерт', 'короткометражный', 'криминал', 'мелодрама', 'мистика', 'музыка', 'мюзикл', 'нуар', 'приключения', 'сборник', 'семейный', 'сказка', 'спорт', 'триллер', 'ужасы', 'фантастика', 'фэнтези', 'эротика']
user_select_genre = st.multiselect('Выберите жанр', genre_list)
if st.button("Искать🔍🎦"):
if user_input:
def encode_description(description, tokenizer, model):
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')
# Векторизация введенного запроса с использованием переданных tokenizer и model
input_embedding = encode_description(user_input, tokenizer, model)
# Поиск с использованием Faiss
_, sorted_indices = index.search(input_embedding.reshape(1, -1), 5)
# Используйте индексы для извлечения строк из DataFrame
recs = df.iloc[sorted_indices[0]].reset_index(drop=True)
recs.index = recs.index + 1
if user_select_genre:
genres_selected = pd.Series(user_select_genre)
genre_mask = df['genre'].str.contains('')
for i in range(len(genres_selected)):
genre_mask_i = df['genre'].str.contains(genres_selected.iloc[i])
genre_mask = genre_mask & genre_mask_i
recs = recs[genre_mask]
if not recs.empty:
# Вывод рекомендованных фильмов с изображениями
st.subheader("Рекомендованные фильмы 🎉:")
for i in range(min(5, len(recs))):
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("""
#### Мы смогли помочь вам с выбором?
""")
feedback = st.text_input('Поделитесь с нами вашим мнением')
feedback_button = st.button("Отправить отзыв", key="feedback_button")
if feedback_button and feedback:
st.success("Спасибо, каждый день мы стараемся быть лучше для вас 💟")
elif feedback_button:
st.warning("Пожалуйста, введите отзыв перед отправкой.")
else:
st.subheader("Подходящих фильмов не найдено, ослабьте фильтры 😔:")
|