Update stri.py

stri.py

@@ -16,10 +16,10 @@ model = AutoModel.from_pretrained(model_name, output_hidden_states=True)
 books = pd.read_csv('book_6000.csv')
 books.dropna(inplace=True)
 
-books = books.reset_index(drop=True)
+
 books = books[books['annotation'].apply(lambda x: len(x.split()) >= 10)]
 books.drop_duplicates(subset='title', keep='first', inplace=True)
-books.reset_index(drop=True)
+books = books.reset_index(drop=True)
 
 
 def data_preprocessing(text: str) -> str:
@@ -50,7 +50,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))
 
 # Определение запроса пользователя