Update app file

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModel
 from sklearn.metrics.pairwise import pairwise_distances, cosine_similarity
+import faiss
 
 tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
 model = AutoModel.from_pretrained("cointegrated/rubert-tiny2")
@@ -12,7 +13,6 @@ df = pd.read_csv('data_final.csv')
 
 MAX_LEN = 300
 
-# @st.cache_resource
 def embed_bert_cls(text, model, tokenizer):
     t = tokenizer(text, padding=True, truncation=True, return_tensors='pt', max_length=MAX_LEN)
     with torch.no_grad():
@@ -23,6 +23,9 @@ def embed_bert_cls(text, model, tokenizer):
 
 books_vector = np.loadtxt('vectors.txt')
 
+index = faiss.IndexFlatIP(books_vector.shape[1])
+index.add(books_vector)
+
 st.title('Приложение для рекомендации книг')
 
 text = st.text_input('Введите запрос:')
@@ -32,17 +35,18 @@ recommend_button = st.button('Найти')
 
 if text and recommend_button:
     user_text_pred = embed_bert_cls(text, model, tokenizer)
-    list_ = pairwise_distances(user_text_pred.reshape(1, -1), books_vector).argsort()[0][:num_results]
+    D, I = index.search(user_text_pred.reshape(1, -1), num_results)
 
     st.subheader('Топ рекомендуемых книг:')
 
-    for i in list_:
+    for i, j in zip(I[0], D[0]):
         col_1, col_2 = st.columns([1, 3])
 
         with col_1:
             st.image(df['image_url'][i], use_column_width=True)
+            st.write(round(j* 100, 2))
         with col_2:
             st.write(f'Название книги: {df["title"][i]}')
             st.write(f'Название книги: {df["author"][i]}')
-            st.write(f'Название книги: {df["annotation"][i]}')
-            st.write(f'{df["page_url"][i]}')
+            st.write(f'Ссылка: {df["page_url"][i]}')
+            st.write(f'Название книги: {df["annotation"][i]}')