Update app.py

app.py

@@ -3,6 +3,8 @@ from sentence_transformers import SentenceTransformer, util
 import json
 import os
 import time
+import threading
+import queue
 
 # Загружаем модель
 model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1"
@@ -37,25 +39,45 @@ if os.path.exists(query_embeddings_file):
 else:
     query_embeddings = {}
 
-def get_movie_embedding(movie):
+# Очередь для необработанных фильмов
+movies_queue = queue.Queue()
+for movie in movies_data:
+    if movie["name"] not in movie_embeddings:
+        movies_queue.put(movie)
+
+# Флаг, указывающий, что обработка фильмов завершена
+processing_complete = False
+
+# Блокировка для доступа к movie_embeddings
+movie_embeddings_lock = threading.Lock()
+
+def process_movies():
     """
-    Возвращает эмбеддинг для фильма. Если эмбеддинг уже создан, возвращает его из словаря.
-    Иначе создает эмбеддинг, сохраняет его и возвращает.
+    Обрабатывает фильмы из очереди, создавая для них эмбеддинги.
     """
-    title = movie["name"]
-    if title in movie_embeddings:
-        print(f"Эмбеддинг для фильма '{title}' уже существует.")
-        return movie_embeddings[title]
-    else:
+    global processing_complete
+    while True:
+        try:
+            movie = movies_queue.get(timeout=1)  # Тайм-аут 1 секунда
+        except queue.Empty:
+            if processing_complete:
+                break
+            else:
+                continue
+
+        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()
-        movie_embeddings[title] = embedding
-        # Сохраняем эмбеддинги в файл
-        with open(embeddings_file, "w", encoding="utf-8") as f:
-            json.dump(movie_embeddings, f, ensure_ascii=False, indent=4)
+
+        with movie_embeddings_lock:
+            movie_embeddings[title] = embedding
+            # Сохраняем эмбеддинги в файл
+            with open(embeddings_file, "w", encoding="utf-8") as f:
+                json.dump(movie_embeddings, f, ensure_ascii=False, indent=4)
         print(f"Эмбеддинг для фильма '{title}' создан и сохранен.")
-        return embedding
+
+    print("Обработка фильмов завершена.")
 
 def get_query_embedding(query):
     """
@@ -75,10 +97,6 @@ def get_query_embedding(query):
         print(f"Эмбеддинг для запроса '{query}' создан и сохранен.")
         return embedding
 
-# Создаем эмбеддинги для всех фильмов
-for movie in movies_data:
-    get_movie_embedding(movie)
-
 def search_movies(query, top_k=3):
     """
     Ищет наиболее похожие фильмы по запросу.
@@ -90,22 +108,29 @@ def search_movies(query, top_k=3):
     Returns:
         Строку с результатами поиска в формате HTML.
     """
-    start_time = time.time()  # Засекаем время начала выполнения
-
+    start_time = time.time()
     print(f"\n\033[1mПоиск по запросу: '{query}'\033[0m")
 
     query_embedding = get_query_embedding(query)
-    query_embedding_tensor = util.pytorch_cos_sim(
-        model.encode(query, convert_to_tensor=True),
-        model.encode(query, convert_to_tensor=True)
-    )
+
+    with movie_embeddings_lock:
+        # Делаем копию словаря, чтобы избежать ошибок при изменении оригинала во время итерации
+        current_movie_embeddings = movie_embeddings.copy()
+
+    if not current_movie_embeddings:
+        return "<p>Пока что нет обработанных фильмов. Попробуйте позже.</p>"
+
+    # Создаем словарь с описаниями фильмов
+    movie_descriptions = {}
+    for movie in movies_data:
+      movie_descriptions[movie["name"]] = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
 
     # Сортируем фильмы по убыванию сходства с запросом
     sorted_movies = sorted(
-        movie_embeddings.items(),
+        current_movie_embeddings.items(),
         key=lambda item: util.pytorch_cos_sim(
-            query_embedding_tensor,
-            model.encode(list(movie_descriptions.values())[list(movie_descriptions.keys()).index(item[0])], convert_to_tensor=True)
+            model.encode(query, convert_to_tensor=True),
+            model.encode(movie_descriptions[item[0]], convert_to_tensor=True)
         )[0][0],
         reverse=True
     )
@@ -119,8 +144,8 @@ def search_movies(query, top_k=3):
                 year = movie["year"]
                 genres = movie["genresList"]
                 score = util.pytorch_cos_sim(
-                    query_embedding_tensor,
-                    model.encode(list(movie_descriptions.values())[list(movie_descriptions.keys()).index(title)], convert_to_tensor=True)
+                    model.encode(query, convert_to_tensor=True),
+                    model.encode(movie_descriptions[title], convert_to_tensor=True)
                 )[0][0].item()
                 break
 
@@ -130,9 +155,8 @@ def search_movies(query, top_k=3):
         results_html += f"<p><b>Сходство:</b> {score:.4f}</p>"
         results_html += "<hr>"
 
-    end_time = time.time()  # Засекаем время окончания выполнения
-    execution_time = end_time - start_time  # Вычисляем время выполнения
-
+    end_time = time.time()
+    execution_time = end_time - start_time
     print(f"Поиск завершен за {execution_time:.4f} секунд.")
 
     return results_html
@@ -152,4 +176,14 @@ iface = gr.Interface(
 )
 
 # Запускаем приложение
-iface.launch()
+iface.queue()
+iface.launch(share=True)
+
+# Запускаем поток для обработки фильмов
+processing_thread = threading.Thread(target=process_movies)
+processing_thread.start()
+
+# Устанавливаем флаг, что обработка завершена, когда поток завершается
+processing_thread.join()
+processing_complete = True
+print("Все фильмы обработаны.")