Update stri.py

stri.py

@@ -45,23 +45,22 @@ max_len = 256
 query = st.text_input("Введите запрос")
 
 if st.button('Сгенерировать'):
-    with open("book_embeddings.pkl", "rb") as f:
+    with open("book_embeddings512.pkl", "rb") as f:
         book_embeddings = pickle.load(f)
 
-    query_tokens = tokenizer.encode(query, add_special_tokens=True,
-                                    truncation=True, max_length=max_len)
-
-    query_padded = np.array(query_tokens + [0] * (max_len - len(query_tokens)))
-    query_mask = np.where(query_padded != 0, 1, 0)
-
-    # Переведем numpy массивы в тензоры PyTorch
-    query_padded = torch.tensor(query_padded, dtype=torch.long)
-    query_mask = torch.tensor(query_mask, dtype=torch.long)
-
-    with torch.inference_mode():
-        query_embedding = model(query_padded.unsqueeze(0), query_mask.unsqueeze(0))
-        query_embedding = query_embedding[0][:,0,:]
-        query_embedding = torch.nn.functional.normalize(query_embedding)
+    query_tokens = tokenizer.encode_plus(
+            query,
+            add_special_tokens=True,
+            max_length=length, # Ограничение на максимальную длину входной последовательности
+            pad_to_max_length=True, # Дополним последовательность нулями до максимальной длины
+            return_tensors='pt' # Вернём тензоры PyTorch
+        )
+
+    with torch.no_grad():
+            query_outputs = model(**query_tokens)
+            query_hidden_states = query_outputs.hidden_states
+            query_last_hidden_state = query_hidden_states[-1] # Используем предпоследний слой для эмбеддинга
+            query_embedding = torch.mean(query_last_hidden_state, dim=1).squeeze()
       
 
     # Вычисление косинусного расстояния между эмбеддингом запроса и каждой аннотацией