Update stri.py

stri.py

@@ -51,7 +51,7 @@ 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_embedding = book_embedding[0][:, 0, :] #.detach().cpu().numpy()
         book_embeddings.append(np.squeeze(book_embedding))
 
 # Определение запроса пользователя
@@ -68,17 +68,17 @@ query_mask = torch.tensor(query_mask, dtype=torch.long)
 
 with torch.no_grad():
     query_embedding = model(query_padded.unsqueeze(0), query_mask.unsqueeze(0))
-    query_embedding = query_embedding[0][:, 0, :].detach().cpu().numpy()
+    query_embedding = query_embedding[0][:, 0, :] #.detach().cpu().numpy()
 
 # Вычисление косинусного расстояния между эмбеддингом запроса и каждой аннотацией
-#cosine_similarities = torch.nn.functional.cosine_similarity(
-#    query_embedding.squeeze(0),
-#    torch.stack(book_embeddings)
-#)
 cosine_similarities = torch.nn.functional.cosine_similarity(
-    torch.tensor(query_embedding.squeeze(0)),
-    torch.stack([torch.tensor(embedding) for embedding in book_embeddings])
+    query_embedding.squeeze(0),
+    torch.stack(book_embeddings)
 )
+#cosine_similarities = torch.nn.functional.cosine_similarity(
+#    torch.tensor(query_embedding.squeeze(0)),
+#    torch.stack([torch.tensor(embedding) for embedding in book_embeddings])
+#)
 
 cosine_similarities = cosine_similarities.numpy()