File size: 4,461 Bytes
4835c0a 4ec6fa8 4835c0a 4ec6fa8 4835c0a 461b6f7 4835c0a 8928848 4835c0a 8928848 4835c0a f8fb48f |
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 |
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:
# emb = emb[mask]
# df = df[mask]
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)
|