Update app.py

app.py

@@ -7,19 +7,15 @@ import threading
 import queue
 import torch
 
-# Загружаем модели
-model_name_kalm = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5"
-model_kalm = SentenceTransformer(model_name_kalm)
+# Загружаем модель
+model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5"
+model = SentenceTransformer(model_name)
+# model.max_seq_length = 8192  # Убираем явное ограничение длины последовательности
 
-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"
+# Имя файла для сохранения эмбеддингов
+embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"
+# Имя файла для сохранения эмбеддингов запросов
+query_embeddings_file = f"query_embeddings_{model_name.replace('/', '_')}.json"
 
 # Загружаем данные из файла movies.json
 try:
@@ -29,80 +25,59 @@ except FileNotFoundError:
     print("Ошибка: Файл movies.json не найден.")
     movies_data = []
 
-# Загружаем эмбеддинги фильмов для 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 из файла.")
+# Загружаем эмбеддинги фильмов
+if os.path.exists(embeddings_file):
+    with open(embeddings_file, "r", encoding="utf-8") as f:
+        movie_embeddings = json.load(f)
+    print("Загружены эмбеддинги фильмов из файла.")
 else:
-    movie_embeddings_kalm = {}
+    movie_embeddings = {}
 
-# Загружаем эмбеддинги запросов для 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 из файла.")
+# Загружаем эмбеддинги запросов
+if os.path.exists(query_embeddings_file):
+    with open(query_embeddings_file, "r", encoding="utf-8") as f:
+        query_embeddings = json.load(f)
+    print("Загружены эмбеддинги запросов из файла.")
 else:
-    query_embeddings_kalm = {}
-
-# Загружаем эмбеддинги фильмов для 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:
-    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()
+    query_embeddings = {}
 
+# Очередь для необработанных фильмов
+movies_queue = queue.Queue()
 for movie in movies_data:
-    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)
+    if movie["name"] not in movie_embeddings:
+        movies_queue.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_lock_kalm = threading.Lock()
-movie_embeddings_lock_bge = threading.Lock()
+# Блокировка для доступа к movie_embeddings
+movie_embeddings_lock = threading.Lock()
 
 # Размер пакета для обработки эмбеддингов
-batch_size = 32
+batch_size = 32  # Увеличиваем размер пакета в 2 раза
 
-# Инструкция для запроса KaLM
-query_prompt_kalm = "Инструкция: Найди релевантные фильмы по запросу. \n Запрос: "
+# Инструкция для запроса
+query_prompt = "Инструкция: Найди релевантные фильмы по запросу. \n Запрос: "
 
-def encode_string(text, model, prompt=None):
+def encode_string(text, prompt=None):
     """Кодирует строку в эмбеддинг с использованием инструкции, если она задана."""
     if prompt:
-        return model.encode(text, prompt=prompt, convert_to_tensor=True, normalize_embeddings=True, batch_size=batch_size)
+        return model.encode(text, prompt=prompt, convert_to_tensor=True, normalize_embeddings=True)
     else:
-        return model.encode(text, convert_to_tensor=True, normalize_embeddings=True, batch_size=batch_size)
+        return model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
 
-def process_movies(model, embeddings_file, movie_embeddings, movies_queue, lock, model_name):
+def process_movies():
     """
     Обрабатывает фильмы из очереди, создавая для них эмбеддинги.
     """
-    global processing_complete_kalm, processing_complete_bge # Добавлено
-
+    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:
@@ -112,11 +87,8 @@ def process_movies(model, embeddings_file, movie_embeddings, movies_queue, lock,
                 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
+            print("Очередь фильмов пуста.")
+            processing_complete = True
             break
 
         titles = [movie["name"] for movie in batch]
@@ -125,22 +97,22 @@ def process_movies(model, embeddings_file, movie_embeddings, movies_queue, lock,
             for movie in batch
         ]
 
-        print(f"Создаются эмбеддинги для фильмов ({model_name}): {', '.join(titles)}...")
+        print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
         embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True).tolist()
 
-        with lock:
+        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"Эмбеддинги для фильмов ({model_name}): {', '.join(titles)} созданы и сохранены.")
+            print(f"Эмбеддинги для фильмов: {', '.join(titles)} созданы и сохранены.")
 
-    print(f"Обработка фильмов для {model_name} завершена.")
+    print("Обработка фильмов завершена.")
 
-def get_query_embedding(query, model, query_embeddings, query_embeddings_file, prompt=None):
+def get_query_embedding(query):
     """
-    Возвращает эмбеддинг для запроса с инструкцией.
+    Возвращает эмбеддинг для запроса с инструкцией. 
     Если эмбеддинг уже создан, возвращает его из словаря.
     Иначе создает эмбеддинг, сохраняет его и возвращает.
     """
@@ -149,7 +121,7 @@ def get_query_embedding(query, model, query_embeddings, query_embeddings_file, p
         return query_embeddings[query]
     else:
         print(f"Создается эмбеддинг для запроса '{query}'...")
-        embedding = encode_string(query, model, prompt=prompt).tolist()
+        embedding = encode_string(query, prompt=query_prompt).tolist()
         query_embeddings[query] = embedding
         # Сохраняем эмбеддинги запросов в файл
         with open(query_embeddings_file, "w", encoding="utf-8") as f:
@@ -157,47 +129,31 @@ 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, query_embeddings, query_embeddings_file, top_k=10, query_prompt=None):
+def search_movies(query, top_k=10):
     """
     Ищет наиболее похожие фильмы по запросу с использованием инструкции.
 
     Args:
         query: Текстовый запрос.
-        model: Модель для эмбеддингов.
-        movie_embeddings: Словарь с эмбеддингами фильмов.
-        movies_data: Данные о фильмах.
         top_k: Количество возвращаемых результатов.
-        query_prompt: Инструкция для запроса (для KaLM).
 
     Returns:
         Строку с результатами поиска в формате HTML.
     """
-    global search_in_progress_kalm, search_in_progress_bge # Добавлено
-    
-    if model == model_kalm:
-        search_in_progress_kalm = True
-    elif model == model_bge:
-        search_in_progress_bge = True
-    
+    global search_in_progress
+    search_in_progress = 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 = torch.tensor(get_query_embedding(query, model, query_embeddings, query_embeddings_file, prompt=query_prompt))
+    query_embedding_tensor = encode_string(query, 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()
-    elif model == model_bge:
-        with movie_embeddings_lock_bge:
-            current_movie_embeddings = movie_embeddings.copy()
+    with movie_embeddings_lock:
+        current_movie_embeddings = movie_embeddings.copy()
 
     if not current_movie_embeddings:
-        if model == model_kalm:
-            search_in_progress_kalm = False
-        elif model == model_bge:
-            search_in_progress_bge = False
+        search_in_progress = False
         return "<p>Пока что нет обработанных фильмов. Попробуйте позже.</p>"
 
     # Преобразуем эмбеддинги фильмов в тензор
@@ -230,40 +186,29 @@ def search_movies(query, model, movie_embeddings, movies_data, query_embeddings,
     end_time = time.time()
     execution_time = end_time - start_time
     print(f"Поиск завершен за {execution_time:.4f} секунд.")
-
-    if model == model_kalm:
-        search_in_progress_kalm = False
-    elif model == model_bge:
-        search_in_progress_bge = False
-
+    search_in_progress = False
     return results_html
 
-# Потоки для обработки фильмов
-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, 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, query_embeddings_bge, query_embeddings_file_bge, top_k=10)
-
-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()
+# Поток для обработки фильмов
+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()