movie_rec / app.py
HaggiVaggi's picture
Update app.py
ac6735c
raw history blame
No virus
4.51 kB
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
from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader
import torch
import ast
import re
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.cpu().numpy().astype('float32')
embeddings_array = np.load('embeddings.npy')
index = faiss.read_index('desc_faiss_index.index')
df2 = pd.read_csv('data_with_embeddings.csv')
# embeddings = pd.read_pickle('embeddings.pkl')
user_input = st.text_area('Введите описание фильма')
input_embedding = encode_description(user_input)
embeddings_tensor = torch.stack([torch.Tensor(ast.literal_eval(embedding_str)) for embedding_str in df2['description_embedding']]).numpy()
if st.button("Искать 🔍"):
if user_input:
# Векторизация введенного запроса
input_embedding = encode_description(user_input)
# Поиск с использованием Faiss
_, sorted_indices = index.search(input_embedding.reshape(1, -1), 10)
# Используйте индексы для извлечения строк из DataFrame
recs = df.iloc[sorted_indices[0]].reset_index(drop=True)
recs.index = recs.index + 1
# Вывод рекомендованных фильмов
st.write("Рекомендованные фильмы:")
st.table(recs[['movie_title', 'description']])