Update app.py

app.py

@@ -5,6 +5,7 @@ import os
 import time
 import threading
 import queue
+import torch
 
 # Загружаем модель
 model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1"
@@ -53,6 +54,9 @@ search_in_progress = False
 # Блокировка для доступа к movie_embeddings
 movie_embeddings_lock = threading.Lock()
 
+# Размер пакета для обработки эмбеддингов
+batch_size = 16  # Подберите оптимальный размер
+
 def encode_string(text):
     """Кодирует строку в эмбеддинг."""
     return model.encode(text, convert_to_tensor=True)
@@ -67,28 +71,35 @@ def process_movies():
             time.sleep(1)  # Ждем, пока поиск не завершится
             continue
 
-        try:
-            movie = movies_queue.get(timeout=1)
-        except queue.Empty:
+        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
 
-        title = movie["name"]
-        if title in movie_embeddings:
-            print(f"Эмбеддинг для фильма '{title}' уже существует. Пропускаем.")
-            continue
+        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"Создается эмбеддинг для фильма '{title}'...")
-        embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
-        embedding = encode_string(embedding_string).tolist()
+        print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
+        embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size).tolist()
 
         with movie_embeddings_lock:
-            movie_embeddings[title] = embedding
-            # Сохраняем эмбеддинги в файл после обработки каждого фильма
+            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"Эмбеддинг для фильма '{title}' создан и сохранен.")
+            print(f"Эмбеддинги для фильмов: {', '.join(titles)} созданы и сохранены.")
 
     print("Обработка фильмов завершена.")
 
@@ -137,21 +148,20 @@ def search_movies(query, top_k=3):
         search_in_progress = False
         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']}"
-
-    print(f"Начало вычисления косинусного сходства: {time.strftime('%Y-%m-%d %H:%M:%S')}")
-    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))
-    print(f"Окончание вычисления косинусного сходства: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    # Преобразуем эмбеддинги фильмов в тензор
+    movie_titles = list(current_movie_embeddings.keys())
+    movie_embeddings_tensor = torch.tensor(list(current_movie_embeddings.values()))
 
-    sorted_movies = sorted(similarity_scores, key=lambda item: item[1], reverse=True)
+    print(f"Начало поиска похожих фильмов: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    # Используем util.semantic_search для поиска похожих фильмов
+    hits = util.semantic_search(query_embedding_tensor, movie_embeddings_tensor, top_k=top_k)[0]
+    print(f"Окончание поиска похожих фильмов: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
     results_html = ""
-    for title, score in sorted_movies[:top_k]:
+    for hit in hits:
+        title = movie_titles[hit['corpus_id']]
+        score = hit['score']
+        # Ищем полное описание фильма в исходных данных
         for movie in movies_data:
             if movie["name"] == title:
                 description = movie["description"]