Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from sentence_transformers import SentenceTransformer, util
 import json
 import os
+import time
 
 # Загружаем модель
 model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1"
@@ -9,6 +10,8 @@ model = SentenceTransformer(model_name)
 
 # Имя файла для сохранения эмбеддингов
 embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"
+# Имя файла для сохранения эмбеддингов запросов
+query_embeddings_file = f"query_embeddings_{model_name.replace('/', '_')}.json"
 
 # Загружаем данные из файла movies.json
 try:
@@ -18,49 +21,63 @@ 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_loaded = json.load(f)
-    print("Загружены эмбеддинги из файла.")
+        movie_embeddings = json.load(f)
+    print("Загружены эмбеддинги фильмов из файла.")
 else:
-    movie_embeddings_loaded = {}
+    movie_embeddings = {}
 
-# Создаем словарь с описаниями фильмов и эмбеддингами
-movie_descriptions = {}
-movie_embeddings = {}
+# Загружаем эмбеддинги запросов
+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 = {}
 
-for movie in movies_data:
+def get_movie_embedding(movie):
+    """
+    Возвращает эмбеддинг для фильма. Если эмбеддинг уже создан, возвращает его из словаря.
+    Иначе создает эмбеддинг, сохраняет его и возвращает.
+    """
     title = movie["name"]
-    # Формируем строку для эмбеддинга из всех данных фильма
-    embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
-    movie_descriptions[title] = embedding_string
-
-    # Проверяем, есть ли эмбеддинг для фильма в загруженных
-    if title in movie_embeddings_loaded:
-        movie_embeddings[title] = movie_embeddings_loaded[title]
+    if title in movie_embeddings:
+        print(f"Эмбеддинг для фильма '{title}' уже существует.")
+        return movie_embeddings[title]
     else:
-        # Создаем эмбеддинг и добавляем в словарь
+        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)
+        print(f"Эмбеддинг для фильма '{title}' создан и сохранен.")
+        return embedding
 
-# Сохраняем эмбеддинги в файл, если были созданы новые
-if len(movie_embeddings_loaded) < len(movie_embeddings):
-    with open(embeddings_file, "w", encoding="utf-8") as f:
-        json.dump(movie_embeddings, f, ensure_ascii=False, indent=4)
-    print("Эмбеддинги сохранены в файл.")
-
-# Конвертируем эмбеддинги в тензоры
-if len(movie_embeddings) > 0:
-    movie_embeddings_tensor = {
-        title: util.pytorch_cos_sim(
-            model.encode(query, convert_to_tensor=True),
-            model.encode(embedding_string, convert_to_tensor=True)
-        )
-        for title, embedding_string in movie_descriptions.items()
-    }
-else:
-    movie_embeddings_tensor = None
+def get_query_embedding(query):
+    """
+    Возвращает эмбеддинг для запроса. Если эмбеддинг уже создан, возвращает его из словаря.
+    Иначе создает эмбеддинг, сохраняет его и возвращает.
+    """
+    if query in query_embeddings:
+        print(f"Эмбеддинг для запроса '{query}' уже существует.")
+        return query_embeddings[query]
+    else:
+        print(f"Создается эмбеддинг для запроса '{query}'...")
+        embedding = model.encode(query, convert_to_tensor=True).tolist()
+        query_embeddings[query] = embedding
+        # Сохраняем эмбеддинги запросов в файл
+        with open(query_embeddings_file, "w", encoding="utf-8") as f:
+            json.dump(query_embeddings, f, ensure_ascii=False, indent=4)
+        print(f"Эмбеддинг для запроса '{query}' создан и сохранен.")
+        return embedding
+
+# Создаем эмбеддинги для всех фильмов
+for movie in movies_data:
+    get_movie_embedding(movie)
 
 def search_movies(query, top_k=3):
     """
@@ -73,15 +90,22 @@ def search_movies(query, top_k=3):
     Returns:
         Строку с результатами поиска в формате HTML.
     """
-    if movie_embeddings_tensor is None:
-        return "<p>Ошибка: Данные фильмов не загружены.</p>"
+    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)
+    )
 
     # Сортируем фильмы по убыванию сходства с запросом
     sorted_movies = sorted(
-        movie_embeddings_tensor.items(),
+        movie_embeddings.items(),
         key=lambda item: util.pytorch_cos_sim(
-            model.encode(query, convert_to_tensor=True),
-            model.encode(movie_descriptions[item[0]], convert_to_tensor=True)
+            query_embedding_tensor,
+            model.encode(list(movie_descriptions.values())[list(movie_descriptions.keys()).index(item[0])], convert_to_tensor=True)
         )[0][0],
         reverse=True
     )
@@ -95,8 +119,8 @@ def search_movies(query, top_k=3):
                 year = movie["year"]
                 genres = movie["genresList"]
                 score = util.pytorch_cos_sim(
-                    model.encode(query, convert_to_tensor=True),
-                    model.encode(movie_descriptions[title], convert_to_tensor=True)
+                    query_embedding_tensor,
+                    model.encode(list(movie_descriptions.values())[list(movie_descriptions.keys()).index(title)], convert_to_tensor=True)
                 )[0][0].item()
                 break
 
@@ -106,6 +130,11 @@ 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  # Вычисляем время выполнения
+
+    print(f"Поиск завершен за {execution_time:.4f} секунд.")
+
     return results_html
 
 # Создаем интерфейс Gradio