Update stri.py

stri.py

@@ -4,7 +4,7 @@ import numpy as np
 import pandas as pd
 from transformers import AutoTokenizer, AutoModel
 import re
-import torch.nn.functional as F
+import pickle
 
 st.title("Книжные рекомендации")
 
@@ -38,27 +38,30 @@ annot = books['annotation']
 
 # Получение эмбеддингов аннотаций каждой книги в датасете
 max_len = 128
-token_annot = annot.apply(lambda x: tokenizer.encode(x, add_special_tokens=True,
-                                                     truncation=True, max_length=max_len))
+# token_annot = annot.apply(lambda x: tokenizer.encode(x, add_special_tokens=True,
+#                                                      truncation=True, max_length=max_len))
 
-padded = np.array([i + [0] * (max_len - len(i)) for i in token_annot.values])  # заполним недостающую длину нулями
-attention_mask = np.where(padded != 0, 1, 0)  # создадим маску, отметим где есть значения а где пустота
-# Переведем numpy массивы в тензоры PyTorch
-input_ids = torch.tensor(padded, dtype=torch.long)
-attention_mask = torch.tensor(attention_mask, dtype=torch.long)
+# padded = np.array([i + [0] * (max_len - len(i)) for i in token_annot.values])  # заполним недостающую длину нулями
+# attention_mask = np.where(padded != 0, 1, 0)  # создадим маску, отметим где есть значения а где пустота
+# # Переведем numpy массивы в тензоры PyTorch
+# input_ids = torch.tensor(padded, dtype=torch.long)
+# attention_mask = torch.tensor(attention_mask, dtype=torch.long)
+
+# book_embeddings = []
+# for inputs, attention_masks in zip(input_ids, attention_mask):
+#     with torch.no_grad():
+#         book_embedding = model(inputs.unsqueeze(0), attention_mask=attention_masks.unsqueeze(0))
+#         book_embedding = book_embedding[0][:, 0, :] #.detach().cpu().numpy()
+#         book_embeddings.append(np.squeeze(book_embedding))
 
-book_embeddings = []
-for inputs, attention_masks in zip(input_ids, attention_mask):
-    with torch.no_grad():
-        book_embedding = model(inputs.unsqueeze(0), attention_mask=attention_masks.unsqueeze(0))
-        book_embedding = book_embedding[0][:, 0, :] #.detach().cpu().numpy()
-        book_embeddings.append(np.squeeze(book_embedding))
-        
 
 # Определение запроса пользователя
 query = st.text_input("Введите запрос")
 
 if st.button('**Generating recommendations**'):
+    with open("book_embeddings.pkl", "rb") as f:
+        book_embeddings = pickle.load(f)
+        
     query_tokens = tokenizer.encode(query, add_special_tokens=True,
                                     truncation=True, max_length=max_len)