Update stri.py

stri.py

@@ -45,7 +45,7 @@ length = 512
 query = st.text_input("Введите запрос")
 
 if st.button('Сгенерировать'):
-    with open("book_embeddings512.pkl", "rb") as f:
+    with open("book_embeddingsN.pkl", "rb") as f:
         book_embeddings = pickle.load(f)
 
     query_tokens = tokenizer.encode_plus(
@@ -58,11 +58,11 @@ if st.button('Сгенерировать'):
 
     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()
-      
+            query_hidden_states = query_outputs.hidden_states[-1][:,0,:]
+            query_hidden_states = torch.nn.functional.normalize(query_hidden_states)query_hidden_states = query_outputs.hidden_states[-1][:,0,:]
+            query_hidden_states = torch.nn.functional.normalize(query_hidden_states)
 
+        
     # Вычисление косинусного расстояния между эмбеддингом запроса и каждой аннотацией
     cosine_similarities = torch.nn.functional.cosine_similarity(
         query_embedding.squeeze(0),