Update app.py

app.py

@@ -1,15 +1,9 @@
 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
+from transformers import AutoTokenizer, AutoModel
 import faiss
+import numpy as np
 
 
 
@@ -58,7 +52,11 @@ if page == "какая-то еще":
     # Загрузка предварительно обученной модели ruBERT
     tokenizer = AutoTokenizer.from_pretrained("DeepPavlov/rubert-base-cased-sentence")
     model = AutoModel.from_pretrained("DeepPavlov/rubert-base-cased-sentence")
-
+        
+    df2 = pd.read_csv('data_with_embeddings.csv')   
+    embeddings_array = np.load('embeddings.npy')
+    index = faiss.read_index('desc_faiss_index.index')
+    
     def encode_description(description):
         tokens = tokenizer(description, return_tensors="pt")
         with torch.no_grad():
@@ -66,29 +64,30 @@ if page == "какая-то еще":
         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()
+    def main():
+    st.title("Система поиска фильмов")
+
+    # Пользовательский ввод
+    user_input = st.text_input("Введите описание фильма:")
 
-    if st.button("Искать 🔍"):
+    if st.button("Искать🔍🎦')"):
         if user_input:
             # Векторизация введенного запроса
             input_embedding = encode_description(user_input)
 
             # Поиск с использованием Faiss
-            _, sorted_indices = index.search(input_embedding.reshape(1, -1), 10)
+            _, sorted_indices = index.search(input_embedding.reshape(1, -1), 5)  # Изменил на 5
 
             # Используйте индексы для извлечения строк из DataFrame
             recs = df.iloc[sorted_indices[0]].reset_index(drop=True)
             recs.index = recs.index + 1
 
-            # Вывод рекомендованных фильмов
+            # Вывод рекомендованных фильмов с изображениями
             st.write("Рекомендованные фильмы:")
-            st.table(recs[['movie_title', 'description']])
-   
+            for i in range(5):  # Изменил на 5
+                st.image(recs['image_url'].iloc[i], caption=recs[['movie_title', 'description']].iloc[i], use_column_width=True)
+
+if __name__ == "__main__":
+    main()