Update app.py

app.py

@@ -7,15 +7,19 @@ import threading
 import queue
 import torch
 
-# Загружаем модель
-model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5"
-model = SentenceTransformer(model_name)
-# model.max_seq_length = 8192  # Убираем явное ограничение длины последовательности
+# Загружаем модели
+model_name_kalm = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5"
+model_kalm = SentenceTransformer(model_name_kalm)
 
-# Имя файла для сохранения эмбеддингов
-embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"
-# Имя файла для сохранения эмбеддингов запросов
-query_embeddings_file = f"query_embeddings_{model_name.replace('/', '_')}.json"
+model_name_bge = "BAAI/bge-m3"
+model_bge = SentenceTransformer(model_name_bge)
+
+# Имена файлов для сохранения эмбеддингов
+embeddings_file_kalm = f"movie_embeddings_{model_name_kalm.replace('/', '_')}.json"
+query_embeddings_file_kalm = f"query_embeddings_{model_name_kalm.replace('/', '_')}.json"
+
+embeddings_file_bge = f"movie_embeddings_{model_name_bge.replace('/', '_')}.json"
+query_embeddings_file_bge = f"query_embeddings_{model_name_bge.replace('/', '_')}.json"
 
 # Загружаем данные из файла movies.json
 try:
@@ -25,94 +29,125 @@ except FileNotFoundError:
     print("Ошибка: Файл movies.json не найден.")
     movies_data = []
 
-# Загружаем эмбеддинги фильмов
-if os.path.exists(embeddings_file):
-    with open(embeddings_file, "r", encoding="utf-8") as f:
-        movie_embeddings = json.load(f)
-    print("Загружены эмбеддинги фильмов из файла.")
+# Загружаем эмбеддинги фильмов для KaLM
+if os.path.exists(embeddings_file_kalm):
+    with open(embeddings_file_kalm, "r", encoding="utf-8") as f:
+        movie_embeddings_kalm = json.load(f)
+    print("Загружены эмбеддинги фильмов для KaLM из файла.")
+else:
+    movie_embeddings_kalm = {}
+
+# Загружаем эмбеддинги запросов для KaLM
+if os.path.exists(query_embeddings_file_kalm):
+    with open(query_embeddings_file_kalm, "r", encoding="utf-8") as f:
+        query_embeddings_kalm = json.load(f)
+    print("Загружены эмбеддинги запросов для KaLM из файла.")
 else:
-    movie_embeddings = {}
+    query_embeddings_kalm = {}
 
-# Загружаем эмбеддинги запросов
-if os.path.exists(query_embeddings_file):
-    with open(query_embeddings_file, "r", encoding="utf-8") as f:
-        query_embeddings = json.load(f)
-    print("Загружены эмбеддинги запросов из файла.")
+# Загружаем эмбеддинги фильмов для BGE-M3
+if os.path.exists(embeddings_file_bge):
+    with open(embeddings_file_bge, "r", encoding="utf-8") as f:
+        movie_embeddings_bge = json.load(f)
+    print("Загружены эмбеддинги фильмов для BGE-M3 из файла.")
 else:
-    query_embeddings = {}
+    movie_embeddings_bge = {}
+
+# Загружаем эмбеддинги запросов для BGE-M3
+if os.path.exists(query_embeddings_file_bge):
+    with open(query_embeddings_file_bge, "r", encoding="utf-8") as f:
+        query_embeddings_bge = json.load(f)
+    print("Загружены эмбеддинги запросов для BGE-M3 из файла.")
+else:
+    query_embeddings_bge = {}
+
+# Очереди для необработанных фильмов
+movies_queue_kalm = queue.Queue()
+movies_queue_bge = queue.Queue()
 
-# Очередь для необработанных фильмов
-movies_queue = queue.Queue()
 for movie in movies_data:
-    if movie["name"] not in movie_embeddings:
-        movies_queue.put(movie)
+    if movie["name"] not in movie_embeddings_kalm:
+        movies_queue_kalm.put(movie)
+    if movie["name"] not in movie_embeddings_bge:
+        movies_queue_bge.put(movie)
+
+# Флаги, указывающие, что обработка фильмов завершена
+processing_complete_kalm = False
+processing_complete_bge = False
 
-# Флаг, указывающий, что обработка фильмов завершена
-processing_complete = False
-# Флаг, указывающий, что выполняется поиск
-search_in_progress = False
+# Флаги, указывающие, что выполняется поиск
+search_in_progress_kalm = False
+search_in_progress_bge = False
 
-# Блокировка для доступа к movie_embeddings
-movie_embeddings_lock = threading.Lock()
+# Блокировки для доступа к эмбеддингам
+movie_embeddings_lock_kalm = threading.Lock()
+movie_embeddings_lock_bge = threading.Lock()
 
 # Размер пакета для обработки эмбеддингов
-batch_size = 32  # Увел��чиваем размер пакета в 2 раза
+batch_size = 32
 
-# Инструкция для запроса
-query_prompt = "Инструкция: Найди релевантные фильмы по запросу. \n Запрос: "
+# Инструкция для запроса KaLM
+query_prompt_kalm = "Инструкция: Найди релевантные фильмы по запросу. \n Запрос: "
 
-def encode_string(text, prompt=None):
+def encode_string(text, model, prompt=None):
     """Кодирует строку в эмбеддинг с использованием инструкции, если она задана."""
     if prompt:
-        return model.encode(text, prompt=prompt, convert_to_tensor=True, normalize_embeddings=True)
+        return model.encode(text, prompt=prompt, convert_to_tensor=True, normalize_embeddings=True, batch_size=batch_size)
     else:
-        return model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
+        return model.encode(text, convert_to_tensor=True, normalize_embeddings=True, batch_size=batch_size)
 
-def process_movies():
+def process_movies(model, embeddings_file, movie_embeddings, movies_queue, processing_complete_flag, lock, model_name):
     """
     Обрабатывает фильмы из очереди, создавая для них эмбеддинги.
     """
-    global processing_complete
     while True:
-        if search_in_progress:
-            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("Очередь фильмов пуста.")
-            processing_complete = 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"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
-        embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True).tolist()
-
-        with movie_embeddings_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"Эмбеддинги для фильмов: {', '.join(titles)} созданы и сохранены.")
-
-    print("Обработка фильмов завершена.")
-
-def get_query_embedding(query):
+      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)} созданы и сохранены.")
+
+    print(f"Обработка фильмов для {model_name} завершена.")
+
+def get_query_embedding(query, model, query_embeddings, query_embeddings_file, prompt=None):
     """
-    Возвращает эмбеддинг для запроса с инструкцией. 
+    Возвращает эмбеддинг для запроса с инструкцией.
     Если эмбеддинг уже создан, возвращает его из словаря.
     Иначе создает эмбеддинг, сохраняет его и возвращает.
     """
@@ -121,7 +156,7 @@ def get_query_embedding(query):
         return query_embeddings[query]
     else:
         print(f"Создается эмбеддинг для запроса '{query}'...")
-        embedding = encode_string(query, prompt=query_prompt).tolist()
+        embedding = encode_string(query, model, prompt=prompt).tolist()
         query_embeddings[query] = embedding
         # Сохраняем эмбеддинги запросов в файл
         with open(query_embeddings_file, "w", encoding="utf-8") as f:
@@ -129,32 +164,53 @@ def get_query_embedding(query):
         print(f"Эмбеддинг для запроса '{query}' создан и сохранен.")
         return embedding
 
-def search_movies(query, top_k=10):
+def search_movies(query, model, movie_embeddings, movies_data, top_k=10, search_in_progress_flag=None, query_prompt=None):
     """
     Ищет наиболее похожие фильмы по запросу с использованием инструкции.
 
     Args:
         query: Текстовый запрос.
+        model: Модель для эмбеддингов.
+        movie_embeddings: Словарь с эмбеддингами фильмов.
+        movies_data: Данные о фильмах.
         top_k: Количество возвращаемых результатов.
+        search_in_progress_flag: Флаг, указывающий, что выполняется поиск.
 
     Returns:
         Строку с результатами поиска в формате HTML.
     """
-    global search_in_progress
-    search_in_progress = True
+    if search_in_progress_flag is not None:
+      if model == model_kalm:
+        global search_in_progress_kalm
+        search_in_progress_kalm = True
+      elif model == model_bge:
+        global search_in_progress_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')}")
-    query_embedding_tensor = encode_string(query, prompt=query_prompt)
+    if model == model_kalm:
+        query_embedding_tensor = encode_string(query, model_kalm, prompt=query_prompt)
+    else:
+      query_embedding_tensor = encode_string(query, model)
+    
     print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
-    with movie_embeddings_lock:
-        current_movie_embeddings = movie_embeddings.copy()
+    if model == model_kalm:
+      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()
 
     if not current_movie_embeddings:
-        search_in_progress = False
-        return "<p>Пока что нет обработанных фильмов. Попробуйте позже.</p>"
+      if search_in_progress_flag is not None:
+          if model == model_kalm:
+            search_in_progress_kalm = False
+          elif model == model_bge:
+            search_in_progress_bge = False
+      return "<p>Пока что нет обработанных фильмов. Попробуйте позже.</p>"
 
     # Преобразуем эмбеддинги фильмов в тензор
     movie_titles = list(current_movie_embeddings.keys())
@@ -186,29 +242,38 @@ def search_movies(query, top_k=10):
     end_time = time.time()
     execution_time = end_time - start_time
     print(f"Поиск завершен за {execution_time:.4f} секунд.")
-    search_in_progress = False
+    if search_in_progress_flag is not None:
+      if model == model_kalm:
+        search_in_progress_kalm = False
+      elif model == model_bge:
+        search_in_progress_bge = False
     return results_html
 
-# Поток для обработки фильмов
-processing_thread = threading.Thread(target=process_movies)
-
-# Создаем интерфейс Gradio
-iface = gr.Interface(
-    fn=search_movies,
-    inputs=gr.Textbox(label="Введите запрос:"),
-    outputs=gr.HTML(label="Результаты п��иска:"),
-    title="Поиск фильмов по описанию",
-    description="Введите запрос, и система найдет наиболее похожие фильмы по их описаниям.",
-    examples=[
-        ["Фильм про ограбление"],
-        ["Комедия 2019 года"],
-        ["Фантастика про космос"],
-    ],
-)
-
-# Запускаем поток для обработки фильмов
-processing_thread.start()
-
-# Запускаем приложение
-iface.queue()
-iface.launch()
+# Потоки для обработки фильмов
+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.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)
+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)
+
+with gr.Blocks() as demo:
+    with gr.Tab("KaLM"):
+        text_input_kalm = gr.Textbox(label="Введите запрос для KaLM")
+        text_output_kalm = gr.HTML()
+        text_button_kalm = gr.Button("Поиск с KaLM")
+    with gr.Tab("BGE-M3"):
+        text_input_bge = gr.Textbox(label="Введите запрос для BGE-M3")
+        text_output_bge = gr.HTML()
+        text_button_bge = gr.Button("Поиск с BGE-M3")
+
+    text_button_kalm.click(search_with_kalm, inputs=text_input_kalm, outputs=text_output_kalm)
+    text_button_bge.click(search_with_bge, inputs=text_input_bge, outputs=text_output_bge)
+
+demo.queue()
+demo.launch()