Update stri.py

stri.py

@@ -17,10 +17,10 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModel.from_pretrained(model_name, output_hidden_states=True)
 
 # Загрузка датасета и аннотаций к книгам
-books = pd.read_csv('books_6000.csv')
+books = pd.read_csv('all+.csv')
 books.dropna(inplace=True)
 
-books = books[books['annotation'].apply(lambda x: len(x.split()) >= 10)]
+books = books[books['annotation'].apply(lambda x: len(x.split()) >= 40)]
 books.drop_duplicates(subset='title', keep='first', inplace=True)
 books = books.reset_index(drop=True)
 
@@ -39,7 +39,7 @@ for i in ['author', 'title', 'annotation']:
 annot = books['annotation']
 
 # Получение эмбеддингов аннотаций каждой книги в датасете
-max_len = 128
+max_len = 256
 
 # Определение запроса пользователя
 query = st.text_input("Введите запрос")
@@ -58,9 +58,11 @@ if st.button('Сгенерировать'):
     query_padded = torch.tensor(query_padded, dtype=torch.long)
     query_mask = torch.tensor(query_mask, dtype=torch.long)
 
-    with torch.no_grad():
+    with torch.inference_mode():
         query_embedding = model(query_padded.unsqueeze(0), query_mask.unsqueeze(0))
-        query_embedding = query_embedding[0][:, 0, :]
+        query_embedding = query_embedding[0][:,0,:]
+        query_embedding = torch.nn.functional.normalize(query_embedding)
+      
 
     # Вычисление косинусного расстояния между эмбеддингом запроса и каждой аннотацией
     cosine_similarities = torch.nn.functional.cosine_similarity(
@@ -83,4 +85,5 @@ if st.button('Сгенерировать'):
         response = requests.get(image_url)
         image = Image.open(BytesIO(response.content))
         cols[0].image(image)
+        cols[0].write(cosine_similarities[i]:.2f)
         cols[1].write("---")