Update app.py

app.py

@@ -51,6 +51,7 @@ for movie in movies_data:
 
 # Флаг, указывающий, что обработка фильмов завершена
 processing_complete = False
+
 # Флаг, указывающий, что выполняется поиск
 search_in_progress = False
 
@@ -74,11 +75,11 @@ def setup_database():
     conn = get_db_connection()
     if conn is None:
         return
-
+    
     with conn.cursor() as cur:
         # Создаем расширение pgvector
         cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
-
+        
         # Создаем таблицу для хранения эмбеддингов фильмов
         cur.execute(f"""
             CREATE TABLE IF NOT EXISTS {embeddings_table} (
@@ -89,7 +90,7 @@ def setup_database():
                 embedding vector(1024)
             );
         """)
-
+        
         # Создаем таблицу для кэширования эмбеддингов запросов
         cur.execute(f"""
             CREATE TABLE IF NOT EXISTS {query_cache_table} (
@@ -102,7 +103,7 @@ def setup_database():
             CREATE INDEX IF NOT EXISTS idx_query_crc32 ON {query_cache_table} (query_crc32);
             CREATE INDEX IF NOT EXISTS idx_created_at ON {query_cache_table} (created_at);
         """)
-
+        
         # Создаем функцию и триггер для автоматического удаления старых записей из таблицы кэша запросов
         cur.execute(f"""
             CREATE OR REPLACE FUNCTION manage_query_cache_size()
@@ -113,11 +114,7 @@ def setup_database():
             BEGIN
                 SELECT pg_total_relation_size('{query_cache_table}') INTO table_size;
                 IF table_size > {MAX_CACHE_SIZE} THEN
-                    FOR row_to_delete IN
-                        SELECT query_crc32
-                        FROM {query_cache_table}
-                        ORDER BY created_at ASC
-                    LOOP
+                    FOR row_to_delete IN SELECT query_crc32 FROM {query_cache_table} ORDER BY created_at ASC LOOP
                         DELETE FROM {query_cache_table} WHERE query_crc32 = row_to_delete.query_crc32;
                         SELECT pg_total_relation_size('{query_cache_table}') INTO table_size;
                         EXIT WHEN table_size <= {MAX_CACHE_SIZE};
@@ -132,7 +129,8 @@ def setup_database():
             FOR EACH ROW
             EXECUTE PROCEDURE manage_query_cache_size();
         """)
-        conn.commit()
+    
+    conn.commit()
     conn.close()
 
 # Настраиваем базу данных при запуске приложения
@@ -164,14 +162,14 @@ def insert_embedding(conn, table_name, crc32_column, crc32_value, other_columns,
     columns = ', '.join([crc32_column] + list(other_columns.keys()) + ['model_name', 'embedding'])
     placeholders = ', '.join(['%s'] * (len(other_columns) + 3))
     values = (crc32_value,) + tuple(other_columns.values()) + (model_name, embedding.tolist())
-
+    
     with conn.cursor() as cur:
         try:
             cur.execute(f"""
                 INSERT INTO {table_name} ({columns})
                 VALUES ({placeholders})
                 ON CONFLICT ({crc32_column}) DO NOTHING;
-                """, values)
+            """, values)
             conn.commit()
             return True
         except Exception as e:
@@ -212,22 +210,22 @@ def process_movies():
             f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
             for movie in batch
         ]
-
+        
         print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
-
+        
         with db_lock:
             for movie, embedding_string in zip(batch, embedding_strings):
                 movie_id = movie['id']
                 string_crc32 = calculate_crc32(embedding_string)
-
+                
                 # Проверяем, есть ли уже эмбеддинг для этого фильма в базе данных
                 existing_embedding = get_embedding_from_db(conn, embeddings_table, "string_crc32", string_crc32, model_name)
-
+                
                 if existing_embedding is None:
                     # Создаем эмбеддинг, только если его нет в базе данных
                     embedding = encode_string(embedding_string)
-                    embedding_crc32 = calculate_crc32(embedding.numpy().tobytes())  # Исправлено
-
+                    embedding_crc32 = calculate_crc32(embedding.cpu().numpy().tobytes())
+                    
                     if insert_embedding(conn, embeddings_table, "embedding_crc32", embedding_crc32, {"movie_id": movie_id, "string_crc32": string_crc32}, embedding):
                         print(f"Эмбеддинг для фильма '{movie['name']}' сохранен в базе данных.")
                     else:
@@ -237,7 +235,7 @@ def process_movies():
 
     conn.close()
     print("Обработка фильмов завершена.")
-    
+
 def get_movie_embeddings(conn):
     """Загружает все эмбеддинги фильмов из базы данных."""
     movie_embeddings = {}
@@ -257,101 +255,83 @@ def get_movie_embeddings(conn):
 def search_movies(query, top_k=10):
     """
     Ищет наиболее похожие фильмы по запросу.
-
     Args:
         query: Текстовый запрос.
         top_k: Количество возвращаемых результатов.
-
     Returns:
         Строку с результатами поиска в формате HTML.
     """
     global search_in_progress
     search_in_progress = True
     start_time = time.time()
+    
     print(f"\n\033[1mПоиск по запросу: '{query}'\033[0m")
-
+    
     conn = get_db_connection()
     if conn is None:
         search_in_progress = False
         return "<p>Ошибка подключения к базе данных.</p>"
 
     query_crc32 = calculate_crc32(query)
-
+    
     # Проверяем, есть ли уже эмбеддинг для этого запроса в кэше
     print(f"Начало поиска эмбеддинга запроса в кэше: {time.strftime('%Y-%m-%d %H:%M:%S')}")
     query_embedding_tensor = get_embedding_from_db(conn, query_cache_table, "query_crc32", query_crc32, model_name)
     print(f"Окончание поиска эмбеддинга запроса в кэше: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
     if query_embedding_tensor is None:
-        print(f"Начало создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+        # Если эмбеддинга нет в кэше, создаем новый
+        print(f"Начало создания эмбеддинга запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
         query_embedding_tensor = encode_string(query)
-        print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
-
-        # Вставляем эмбеддинг запроса в базу данных
+        print(f"Окончание создания эмбеддинга запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+        
+        # Сохраняем эмбеддинг запроса в кэш
         insert_embedding(conn, query_cache_table, "query_crc32", query_crc32, {"query": query}, query_embedding_tensor)
+    else:
+        print("Эмбеддинг запроса найден в кэше.")
+
+    # Загружаем эмбеддинги фильмов
+    print(f"Начало загрузки эмбеддингов фильмов: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    movie_embeddings = get_movie_embeddings(conn)
+    print(f"Окончание загрузки эмбеддингов фильмов: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+
+    # Вычисляем косинусное сходство
+    print(f"Начало вычисления косинусного сходства: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    similarities = []
+    for title, movie_embedding in movie_embeddings.items():
+        similarity = util.pytorch_cos_sim(query_embedding_tensor, movie_embedding).item()
+        similarities.append((title, similarity))
+    
+    # Сортируем результаты
+    similarities.sort(key=lambda x: x[1], reverse=True)
+    top_results = similarities[:top_k]
+    print(f"Окончание вычисления косинусного сходства: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
-    with db_lock:
-        current_movie_embeddings = get_movie_embeddings(conn)
-
-    conn.close()
-
-    if not current_movie_embeddings:
-        search_in_progress = False
-        return "<p>Пока что нет обработанных фильмов. Попробуйте позже.</p>"
-
-    # Преобразуем эмбеддинги фильмов в тензор
-    movie_titles = list(current_movie_embeddings.keys())
-    movie_embeddings_tensor = torch.stack(list(current_movie_embeddings.values()))
-
-    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 hit in hits:
-        title = movie_titles[hit['corpus_id']]
-        score = hit['score']
-        # Ищем полное описание фильма в исходных данных
-        for movie in movies_data:
-            if movie["name"] == title:
-                description = movie["description"]
-                year = movie["year"]
-                genres = movie["genresList"]
-                break
-
-        results_html += f"<h3><b>{title} ({year})</b></h3>"
-        results_html += f"<p><b>Жанры:</b> {genres}</p>"
-        results_html += f"<p><b>Описание:</b> {description}</p>"
-        results_html += f"<p><b>Сходство:</b> {score:.4f}</p>"
-        results_html += "<hr>"
+    # Формируем HTML-строку с результатами
+    results_html = "<ol>"
+    for title, score in top_results:
+        results_html += f"<li><strong>{title}</strong> (Сходство: {score:.4f})</li>"
+    results_html += "</ol>"
 
-    end_time = time.time()
-    execution_time = end_time - start_time
-    print(f"Поиск завершен за {execution_time:.4f} секунд.")
     search_in_progress = False
-    return results_html
+    end_time = time.time()
+    search_time = end_time - start_time
+    print(f"\033[1mПоиск завершен за {search_time:.2f} секунд.\033[0m")
 
-# Поток для обработки фильмов
+    return f"<p>Время поиска: {search_time:.2f} секунд</p>" + results_html
+
+# Запускаем обработку фильмов в отдельном потоке
 processing_thread = threading.Thread(target=process_movies)
+processing_thread.start()
 
 # Создаем интерфейс Gradio
 iface = gr.Interface(
     fn=search_movies,
-    inputs=gr.Textbox(label="Введите запрос:"),
-    outputs=gr.HTML(label="Результаты поиска:"),
-    title="Поиск фильмов по описанию",
-    description="Введите запрос, и система найдет наиболее похожие фильмы по их описаниям.",
-    examples=[
-        ["Фильм про ограбление"],
-        ["Комедия 2019 года"],
-        ["Фантастика про космос"],
-    ],
+    inputs=gr.Textbox(lines=2, placeholder="Введите запрос для поиска фильмов..."),
+    outputs=gr.HTML(label="Результаты поиска"),
+    title="Семантический поиск фильмов",
+    description="Введите описание фильма, который вы ищете, и система найдет наиболее похожие фильмы."
 )
 
-# Запускаем поток для обработки фильмов
-processing_thread.start()
-
-# Запускаем приложение
-iface.queue()
-iface.launch()
+# Запускаем интерфейс
+iface.launch()