Spaces:
Sleeping
Sleeping
import streamlit as st | |
import pandas as pd | |
import numpy as np | |
from sklearn.metrics.pairwise import cosine_similarity | |
from sentence_transformers import SentenceTransformer, util | |
from transformers import AutoTokenizer, AutoModel | |
import faiss | |
from sentence_transformers import SentenceTransformer, InputExample, losses | |
from torch.utils.data import DataLoader | |
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(f"<span style='font-size:{20}px; color:purple'>{'Задача №1: Спарсить информацию по 5 тыс. [фильмов](https://www.kinoafisha.info/) и обработать ее'}</span>", unsafe_allow_html=True) | |
st.markdown(f"<span style='font-size:{20}px; 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.info('🌟Сервис принимает на вход описание фильма от пользователя и возвращает заданное количество подходящих вариантов') | |
df = pd.read_csv('data_with_adds.csv') | |
st.info('🌟Если не знаете, что посмотреть, нажмите кнопку "Сгенерировать" - сервис предложит вам 10 случайных фильмов из библиотеки') | |
if st.button("Сгенерировать 🎲"): | |
# Получение случайных 10 строк | |
random_rows = df[['movie_title', 'description']].sample(n=10).reset_index(drop=True) | |
random_rows.index = random_rows.index + 1 | |
# Вывод на страничке Streamlit | |
st.write("Случайные 10 фильмов") | |
st.write(random_rows) | |
if page == "какая-то еще": | |
# Загрузка предварительно обученной модели ruBERT | |
tokenizer = AutoTokenizer.from_pretrained("DeepPavlov/rubert-base-cased-sentence") | |
model = AutoModel.from_pretrained("DeepPavlov/rubert-base-cased-sentence") | |
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 | |
embeddings = pd.read_pickle('embeddings.pkl') | |
user_input = st.text_area('Введите описание фильма') | |
input_embedding = encode_description(user_input) | |
mbeddings_tensor = torch.stack(df['description_embedding'].tolist()).numpy() | |
# Рассчитайте косинусное сходство | |
similarity_scores = cosine_similarity(input_embedding.view(1, -1).detach().numpy(), embeddings_tensor.reshape(embeddings_tensor.shape[0], -1))[0] | |
# Получение индексов отсортированных значений | |
sorted_indices = similarity_scores.argsort()[::-1] | |
# Используйте индексы для извлечения строк из DataFrame | |
recs = df.iloc[sorted_indices[:10]].reset_index(drop=True) | |
recs.index = recs.index + 1 | |
st.write(recs[['movie_title', 'description']]) |