|
import streamlit as st |
|
import pandas as pd |
|
import numpy as np |
|
import torch |
|
from sentence_transformers import SentenceTransformer |
|
from resources.functions import recommend, find_rows_with_genres, get_mask_in_range |
|
|
|
|
|
st.markdown(f"<h1 style='text-align: center;'>Семантический поиск фильмов", unsafe_allow_html=True) |
|
|
|
df = pd.read_csv('resources/DF_FINAL.csv') |
|
genre_lists = df['ganres'].apply(lambda x: x.split(', ') if isinstance(x, str) else []) |
|
all_genres = sorted(list(set([genre for sublist in genre_lists for genre in sublist]))) |
|
|
|
st.write(f'<p style="text-align: center; font-family: Arial, sans-serif; font-size: 20px; color: white;">Количество фильмов \ |
|
для поиска {len(df)}</p>', unsafe_allow_html=True) |
|
|
|
st.header(':wrench: Панель инструментов') |
|
col1, col2, col3 = st.columns([1, 2, 2]) |
|
with col1: |
|
top_k = st.selectbox("Сколько фильмов?", options=[5, 10, 15, 20]) |
|
with col2: |
|
model_type = st.selectbox("Какой моделью пользуемся?\n ", options=['rubert-tiny2', 'msmarco-MiniLM-L-12-v3']) |
|
with col3: |
|
genres_list = st.multiselect("Какого жанра фильм?\n ", options=all_genres) |
|
|
|
if model_type == 'rubert-tiny2': |
|
model = SentenceTransformer('cointegrated/rubert-tiny2') |
|
emb = torch.load('resources/corpus_embeddings_rub.pth', map_location=torch.device('cpu')) |
|
else: |
|
model = SentenceTransformer('msmarco-MiniLM-L-12-v3') |
|
emb = torch.load('resources/corpus_embeddings_ms.pth', map_location=torch.device('cpu')) |
|
|
|
range_years = st.slider("В каком году вышел фильм?", min_value=df['year'].unique().min(), |
|
max_value=df['year'].unique().max(), |
|
value=(df['year'].unique().min(), df['year'].unique().max())) |
|
|
|
text = st.text_input('Что будем искать?') |
|
button = st.button('Начать поиск', type="primary") |
|
|
|
if text and button: |
|
if len(genres_list) == 0: |
|
mask = get_mask_in_range(df=df, range_values=range_years) |
|
else: |
|
mask1 = find_rows_with_genres(df=df, genres_list=genres_list) |
|
mask2 = get_mask_in_range(df=df, range_values=range_years) |
|
mask = mask1 & mask2 |
|
try: |
|
|
|
|
|
hits = recommend(model, text, emb, len(df)) |
|
st.write(f'<p style="font-family: Arial, sans-serif; font-size: 24px; color: pink; font-weight: bold;"><strong>\ |
|
{top_k} лучших рекомендаций</strong></p>', unsafe_allow_html=True) |
|
st.write('\n') |
|
mask_ind = df[mask].index.tolist() |
|
fil_hits = [hits[0][i] for i in range(len(hits[0])) if hits[0][i]['corpus_id'] in mask_ind] |
|
for i in range(top_k): |
|
col4, col5 = st.columns([3, 4]) |
|
with col4: |
|
try: |
|
st.image(df['poster'][fil_hits[i]['corpus_id']], width=300) |
|
except: |
|
st.image('https://cdnn11.img.sputnik.by/img/104126/36/1041263627_235:441:1472:1802_1920x0_80_0_0_fc2acc893b618b7c650d661fafe178b8.jpg', width=300) |
|
with col5: |
|
st.write(f"***Название:*** {df['title'][fil_hits[i]['corpus_id']]}") |
|
st.write(f"***Жанр:*** {(df['ganres'][fil_hits[i]['corpus_id']])}") |
|
st.write(f"***Описание:*** {df['description'][fil_hits[i]['corpus_id']]}") |
|
st.write(f"***Год:*** {df['year'][fil_hits[i]['corpus_id']]}") |
|
st.write(f"***Актерский состав:*** {df['cast'][fil_hits[i]['corpus_id']]}") |
|
st.write(f"***Косинусное сходство:*** {round(fil_hits[i]['score'], 2)}") |
|
st.write(f"***Ссылка на фильм : {df['url'][fil_hits[i]['corpus_id']]}***") |
|
|
|
st.markdown( |
|
"<hr style='border: 2px solid #000; margin-top: 10px; margin-bottom: 10px;'>", |
|
unsafe_allow_html=True |
|
) |
|
except: |
|
message = '<p style="font-family: Arial, sans-serif; font-size: 24px; color: pink; font-weight: bold;"><strong>\ |
|
Подходящих вариантов нет. Измените критерии поиска или отключайте интернет и читайте\ |
|
<a href="https://huggingface.co/spaces/sakoser/rec_sys_books">книги</a>.</strong></p>' |
|
st.write(message, unsafe_allow_html=True) |
|
|
