Update app.py

app.py

@@ -51,6 +51,10 @@ processing_complete = False
 # Блокировка для доступа к movie_embeddings
 movie_embeddings_lock = threading.Lock()
 
+def encode_string(text):
+    """Кодирует строку в эмбеддинг."""
+    return model.encode(text, convert_to_tensor=True)
+
 def process_movies():
     """
     Обрабатывает фильмы из очереди, создавая для них эмбеддинги.
@@ -67,7 +71,7 @@ def process_movies():
         title = movie["name"]
         print(f"Создается эмбеддинг для фильма '{title}'...")
         embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
-        embedding = model.encode(embedding_string, convert_to_tensor=True).tolist()
+        embedding = encode_string(embedding_string).tolist()
 
         with movie_embeddings_lock:
             movie_embeddings[title] = embedding
@@ -88,7 +92,7 @@ def get_query_embedding(query):
         return query_embeddings[query]
     else:
         print(f"Создается эмбеддинг для запроса '{query}'...")
-        embedding = model.encode(query, convert_to_tensor=True).tolist()
+        embedding = encode_string(query).tolist()
         query_embeddings[query] = embedding
         # Сохраняем эмбеддинги запросов в файл
         with open(query_embeddings_file, "w", encoding="utf-8") as f:
@@ -111,6 +115,7 @@ def search_movies(query, top_k=3):
     print(f"\n\033[1mПоиск по запросу: '{query}'\033[0m")
 
     query_embedding = get_query_embedding(query)
+    query_embedding_tensor = encode_string(query)
 
     with movie_embeddings_lock:
         # Делаем копию словаря, чтобы избежать ошибок при изменении оригинала во время итерации
@@ -124,28 +129,23 @@ def search_movies(query, top_k=3):
     for movie in movies_data:
       movie_descriptions[movie["name"]] = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
 
-    # Сортируем фильмы по убыванию сходства с запросом
-    sorted_movies = sorted(
-        current_movie_embeddings.items(),
-        key=lambda item: util.pytorch_cos_sim(
-            model.encode(query, convert_to_tensor=True),
-            model.encode(movie_descriptions[item[0]], convert_to_tensor=True)
-        )[0][0],
-        reverse=True
-    )
+    # Вычисляем косинусное сходство и сохраняем в список кортежей (название фильма, сходство)
+    similarity_scores = []
+    for title, embedding in current_movie_embeddings.items():
+        similarity = util.pytorch_cos_sim(query_embedding_tensor, encode_string(movie_descriptions[title]))[0][0].item()
+        similarity_scores.append((title, similarity))
+
+    # Сортируем фильмы по убыванию сходства
+    sorted_movies = sorted(similarity_scores, key=lambda item: item[1], reverse=True)
 
     results_html = ""
-    for title, _ in sorted_movies[:top_k]:
+    for title, score in sorted_movies[:top_k]:
         # Ищем полное описание фильма в исходных данных
         for movie in movies_data:
             if movie["name"] == title:
                 description = movie["description"]
                 year = movie["year"]
                 genres = movie["genresList"]
-                score = util.pytorch_cos_sim(
-                    model.encode(query, convert_to_tensor=True),
-                    model.encode(movie_descriptions[title], convert_to_tensor=True)
-                )[0][0].item()
                 break
 
         results_html += f"<h3><b>{title} ({year})</b></h3>"