Update app.py

app.py

@@ -96,52 +96,45 @@ def encode_string(text, model, prompt=None):
     else:
         return model.encode(text, convert_to_tensor=True, normalize_embeddings=True, batch_size=batch_size)
 
-def process_movies(model, embeddings_file, movie_embeddings, movies_queue, processing_complete_flag, lock, model_name):
+def process_movies(model, embeddings_file, movie_embeddings, movies_queue, lock, model_name):
     """
     Обрабатывает фильмы из очереди, создавая для них эмбеддинги.
     """
+    global processing_complete_kalm, processing_complete_bge # Добавлено
+
     while True:
-      if model_name == "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5" and search_in_progress_kalm:
-          time.sleep(1)  # Ждем, пока поиск не завершится
-          continue
-      elif model_name == "BAAI/bge-m3" and search_in_progress_bge:
-          time.sleep(1)  # Ждем, пока поиск не завершится
-          continue
-
-      batch = []
-      while not movies_queue.empty() and len(batch) < batch_size:
-          try:
-              movie = movies_queue.get(timeout=1)
-              batch.append(movie)
-          except queue.Empty:
-              break
-
-      if not batch:
-          print(f"Очередь фильмов для {model_name} пуста.")
-          if model_name == "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5":
-            global processing_complete_kalm
-            processing_complete_kalm = True
-          elif model_name == "BAAI/bge-m3":
-            global processing_complete_bge
-            processing_complete_bge = True
-          break
-
-      titles = [movie["name"] for movie in batch]
-      embedding_strings = [
-          f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
-          for movie in batch
-      ]
-
-      print(f"Создаются эмбеддинги для фильмов ({model_name}): {', '.join(titles)}...")
-      embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True).tolist()
-
-      with lock:
-          for title, embedding in zip(titles, embeddings):
-              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"Эмбеддинги для фильмов ({model_name}): {', '.join(titles)} созданы и сохранены.")
+        batch = []
+        while not movies_queue.empty() and len(batch) < batch_size:
+            try:
+                movie = movies_queue.get(timeout=1)
+                batch.append(movie)
+            except queue.Empty:
+                break
+
+        if not batch:
+            print(f"Очередь фильмов для {model_name} пуста.")
+            if model_name == model_name_kalm:
+                processing_complete_kalm = True
+            elif model_name == model_name_bge:
+                processing_complete_bge = True
+            break
+
+        titles = [movie["name"] for movie in batch]
+        embedding_strings = [
+            f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
+            for movie in batch
+        ]
+
+        print(f"Создаются эмбеддинги для фильмов ({model_name}): {', '.join(titles)}...")
+        embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True).tolist()
+
+        with lock:
+            for title, embedding in zip(titles, embeddings):
+                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"Эмбеддинги для фильмов ({model_name}): {', '.join(titles)} созданы и сохранены.")
 
     print(f"Обработка фильмов для {model_name} завершена.")
 
@@ -164,7 +157,7 @@ def get_query_embedding(query, model, query_embeddings, query_embeddings_file, p
         print(f"Эмбеддинг для запроса '{query}' создан и сохранен.")
         return embedding
 
-def search_movies(query, model, movie_embeddings, movies_data, top_k=10, search_in_progress_flag=None, query_prompt=None):
+def search_movies(query, model, movie_embeddings, movies_data, query_embeddings, query_embeddings_file, top_k=10, query_prompt=None):
     """
     Ищет наиболее похожие фильмы по запросу с использова��ием инструкции.
 
@@ -174,43 +167,38 @@ def search_movies(query, model, movie_embeddings, movies_data, top_k=10, search_
         movie_embeddings: Словарь с эмбеддингами фильмов.
         movies_data: Данные о фильмах.
         top_k: Количество возвращаемых результатов.
-        search_in_progress_flag: Флаг, указывающий, что выполняется поиск.
+        query_prompt: Инструкция для запроса (для KaLM).
 
     Returns:
         Строку с результатами поиска в формате HTML.
     """
-    if search_in_progress_flag is not None:
-      if model == model_kalm:
-        global search_in_progress_kalm
+    global search_in_progress_kalm, search_in_progress_bge # Добавлено
+    
+    if model == model_kalm:
         search_in_progress_kalm = True
-      elif model == model_bge:
-        global search_in_progress_bge
+    elif model == model_bge:
         search_in_progress_bge = True
+    
     start_time = time.time()
     print(f"\n\033[1mПоиск по запросу: '{query}'\033[0m")
 
     print(f"Начало создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
-    if model == model_kalm:
-        query_embedding_tensor = encode_string(query, model_kalm, prompt=query_prompt)
-    else:
-      query_embedding_tensor = encode_string(query, model)
-    
+    query_embedding_tensor = torch.tensor(get_query_embedding(query, model, query_embeddings, query_embeddings_file, prompt=query_prompt))
     print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
     if model == model_kalm:
-      with movie_embeddings_lock_kalm:
-          current_movie_embeddings = movie_embeddings.copy()
+        with movie_embeddings_lock_kalm:
+            current_movie_embeddings = movie_embeddings.copy()
     elif model == model_bge:
-      with movie_embeddings_lock_bge:
-          current_movie_embeddings = movie_embeddings.copy()
+        with movie_embeddings_lock_bge:
+            current_movie_embeddings = movie_embeddings.copy()
 
     if not current_movie_embeddings:
-      if search_in_progress_flag is not None:
-          if model == model_kalm:
+        if model == model_kalm:
             search_in_progress_kalm = False
-          elif model == model_bge:
+        elif model == model_bge:
             search_in_progress_bge = False
-      return "<p>Пока что нет обработанных фильмов. Попробуйте позже.</p>"
+        return "<p>Пока что нет обработанных фильмов. Попробуйте позже.</p>"
 
     # Преобразуем эмбеддинги фильмов в тензор
     movie_titles = list(current_movie_embeddings.keys())
@@ -242,25 +230,27 @@ def search_movies(query, model, movie_embeddings, movies_data, top_k=10, search_
     end_time = time.time()
     execution_time = end_time - start_time
     print(f"Поиск завершен за {execution_time:.4f} секунд.")
-    if search_in_progress_flag is not None:
-      if model == model_kalm:
+
+    if model == model_kalm:
         search_in_progress_kalm = False
-      elif model == model_bge:
+    elif model == model_bge:
         search_in_progress_bge = False
+
     return results_html
 
 # Потоки для обработки фильмов
-processing_thread_kalm = threading.Thread(target=process_movies, args=(model_kalm, embeddings_file_kalm, movie_embeddings_kalm, movies_queue_kalm, processing_complete_kalm, movie_embeddings_lock_kalm, model_name_kalm))
-processing_thread_bge = threading.Thread(target=process_movies, args=(model_bge, embeddings_file_bge, movie_embeddings_bge, movies_queue_bge, processing_complete_bge, movie_embeddings_lock_bge, model_name_bge))
+processing_thread_kalm = threading.Thread(target=process_movies, args=(model_kalm, embeddings_file_kalm, movie_embeddings_kalm, movies_queue_kalm, movie_embeddings_lock_kalm, model_name_kalm))
+processing_thread_bge = threading.Thread(target=process_movies, args=(model_bge, embeddings_file_bge, movie_embeddings_bge, movies_queue_bge, movie_embeddings_lock_bge, model_name_bge))
 
 # Запускаем потоки для обработки фильмов
 processing_thread_kalm.start()
 processing_thread_bge.start()
 
 def search_with_kalm(query):
-  return search_movies(query, model_kalm, movie_embeddings_kalm, movies_data, top_k=10, search_in_progress_flag=search_in_progress_kalm, query_prompt=query_prompt_kalm)
+    return search_movies(query, model_kalm, movie_embeddings_kalm, movies_data, query_embeddings_kalm, query_embeddings_file_kalm, top_k=10, query_prompt=query_prompt_kalm)
+
 def search_with_bge(query):
-  return search_movies(query, model_bge, movie_embeddings_bge, movies_data, top_k=10, search_in_progress_flag=search_in_progress_bge)
+    return search_movies(query, model_bge, movie_embeddings_bge, movies_data, query_embeddings_bge, query_embeddings_file_bge, top_k=10)
 
 with gr.Blocks() as demo:
     with gr.Tab("KaLM"):