Upload app.py

app.py

@@ -0,0 +1,79 @@
+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, 1]) 
+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')
+else:
+    model = SentenceTransformer('msmarco-MiniLM-L-12-v3')
+    emb = torch.load('resources/corpus_embeddings_ms.pth')
+
+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:
+        st.write(f'<p style="font-family: Arial, sans-serif; font-size: 24px; color: pink; font-weight: bold;"><strong>\
+            Подходящих вариантов нет. Измените критерии поиска.</strong></p>', unsafe_allow_html=True)