Update app.py

app.py

@@ -11,7 +11,6 @@ import numpy as np
 from urllib.parse import urlparse
 import logging
 from sklearn.preprocessing import normalize
-import json
 
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -40,16 +39,27 @@ logging.info("Модель загружена успешно.")
 # Имена таблиц
 embeddings_table = "movie_embeddings"
 query_cache_table = "query_cache"
-movies_table = "Movies"  # Новая таблица Movies
 
 # Максимальный размер таблицы кэша запросов в байтах (50MB)
 MAX_CACHE_SIZE = 50 * 1024 * 1024
 
+# Загружаем данные из файла movies.json
+try:
+    import json
+    with open("movies.json", "r", encoding="utf-8") as f:
+        movies_data = json.load(f)
+    logging.info(f"Загружено {len(movies_data)} фильмов из movies.json")
+except FileNotFoundError:
+    logging.error("Ошибка: Файл movies.json не найден.")
+    movies_data = []
+
 # Очередь для необработанных фильмов
 movies_queue = queue.Queue()
 
-# Флаги
+# Флаг, указывающий, что обработка фильмов завершена
 processing_complete = False
+
+# Флаг, указывающий, что выполняется поиск
 search_in_progress = False
 
 # Блокировка для доступа к базе данных
@@ -72,49 +82,40 @@ def setup_database():
     conn = get_db_connection()
     if conn is None:
         return
+
     try:
         with conn.cursor() as cur:
             # Создаем расширение pgvector если его нет
             cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
-            
+
             # Удаляем существующие таблицы если они есть
             cur.execute(f"DROP TABLE IF EXISTS {embeddings_table}, {query_cache_table};")
 
             # Создаем таблицу для хранения эмбеддингов фильмов
             cur.execute(f"""
-            CREATE TABLE IF NOT EXISTS {embeddings_table} (
-                movie_id INTEGER PRIMARY KEY,
-                embedding_crc32 BIGINT,
-                string_crc32 BIGINT,
-                model_name TEXT,
-                embedding vector(1024)
-            );
-            CREATE INDEX IF NOT EXISTS idx_embeddings_string_crc32 ON {embeddings_table} (string_crc32);
+                CREATE TABLE {embeddings_table} (
+                    movie_id INTEGER PRIMARY KEY,
+                    embedding_crc32 BIGINT,
+                    string_crc32 BIGINT,
+                    model_name TEXT,
+                    embedding vector(1024)
+                );
+                CREATE INDEX ON {embeddings_table} (string_crc32);
             """)
-            
+
             # Создаем таблицу для кэширования запросов
             cur.execute(f"""
-            CREATE TABLE IF NOT EXISTS {query_cache_table} (
-                query_crc32 BIGINT PRIMARY KEY,
-                query TEXT,
-                model_name TEXT,
-                embedding vector(1024),
-                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
-            );
-            CREATE INDEX IF NOT EXISTS idx_cache_query_crc32 ON {query_cache_table} (query_crc32);
-            CREATE INDEX IF NOT EXISTS idx_cache_created_at ON {query_cache_table} (created_at);
-            """)
-            
-            # Проверяем существование таблицы Movies
-            cur.execute(f"""
-            SELECT EXISTS (
-                SELECT FROM information_schema.tables 
-                WHERE table_name = '{movies_table}'
-            );
+                CREATE TABLE {query_cache_table} (
+                    query_crc32 BIGINT PRIMARY KEY,
+                    query TEXT,
+                    model_name TEXT,
+                    embedding vector(1024),
+                    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
+                );
+                CREATE INDEX ON {query_cache_table} (query_crc32);
+                CREATE INDEX ON {query_cache_table} (created_at);
             """)
-            if not cur.fetchone()[0]:
-                logging.error(f"Таблица {movies_table} не существует в базе данных.")
-            
+
         conn.commit()
         logging.info("База данных успешно настроена.")
     except Exception as e:
@@ -140,43 +141,33 @@ def get_movies_without_embeddings():
     conn = get_db_connection()
     if conn is None:
         return []
-    
+
     movies_to_process = []
     try:
         with conn.cursor() as cur:
-            # Получаем список ID фильмов, которые уже есть в таблице эмбеддингов
+            # Получаем список ID фильмов, которые уже есть в базе
             cur.execute(f"SELECT movie_id FROM {embeddings_table}")
             existing_ids = {row[0] for row in cur.fetchall()}
-            
-            # Получаем фильмы из таблицы Movies, которых нет в таблице эмбеддингов
-            cur.execute(f"""
-            SELECT id, data FROM {movies_table}
-            WHERE id NOT IN (SELECT movie_id FROM {embeddings_table})
-            """)
-            
-            for row in cur.fetchall():
-                movie_id, movie_data = row
-                movie_info = json.loads(movie_data)
-                movies_to_process.append({
-                    'id': movie_id,
-                    'name': movie_info.get('name', ''),
-                    'description': movie_info.get('description', ''),
-                    'genres': [genre['name'] for genre in movie_info.get('genres', [])]
-                })
-            
+
+            # Фильтруем только те фильмы, которых нет в базе
+            for movie in movies_data:
+                if movie['id'] not in existing_ids:
+                    movies_to_process.append(movie)
+
         logging.info(f"Найдено {len(movies_to_process)} фильмов для обработки.")
     except Exception as e:
         logging.error(f"Ошибка при получении списка фильмов для обработки: {e}")
     finally:
         conn.close()
-    
+
     return movies_to_process
 
 def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_name):
     """Получает эмбеддинг из базы данных."""
     try:
         with conn.cursor() as cur:
-            cur.execute(f"SELECT embedding FROM {table_name} WHERE {crc32_column} = %s AND model_name = %s", (crc32_value, model_name))
+            cur.execute(f"SELECT embedding FROM {table_name} WHERE {crc32_column} = %s AND model_name = %s", 
+                       (crc32_value, model_name))
             result = cur.fetchone()
             if result and result[0]:
                 # Нормализуем эмбеддинг после извлечения из БД
@@ -192,9 +183,10 @@ def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32,
         normalized_embedding = normalize(embedding.reshape(1, -1))[0]
         with conn.cursor() as cur:
             cur.execute(f"""
-            INSERT INTO {table_name} (movie_id, embedding_crc32, string_crc32, model_name, embedding)
-            VALUES (%s, %s, %s, %s, %s)
-            ON CONFLICT (movie_id) DO NOTHING
+                INSERT INTO {table_name} 
+                (movie_id, embedding_crc32, string_crc32, model_name, embedding)
+                VALUES (%s, %s, %s, %s, %s)
+                ON CONFLICT (movie_id) DO NOTHING
             """, (movie_id, embedding_crc32, string_crc32, model_name, normalized_embedding.tolist()))
         conn.commit()
         return True
@@ -206,10 +198,12 @@ def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32,
 def process_movies():
     """Обрабатывает фильмы, создавая для них эмбеддинги."""
     global processing_complete
+
     logging.info("Начало обработки фильмов.")
-    
+
     # Получаем список фильмов, которые нужно обработать
     movies_to_process = get_movies_without_embeddings()
+
     if not movies_to_process:
         logging.info("Все фильмы уже обработаны.")
         processing_complete = True
@@ -242,22 +236,24 @@ def process_movies():
                 break
 
             logging.info(f"Обработка пакета из {len(batch)} фильмов...")
+
             for movie in batch:
-                embedding_string = f"Название: {movie['name']}\nЖанры: {', '.join(movie['genres'])}\nОписание: {movie['description']}"
+                embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
                 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(str(embedding.tolist()))
+
                     if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
                         logging.info(f"Сохранен эмбеддинг для '{movie['name']}'")
                     else:
                         logging.error(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
                 else:
                     logging.info(f"Эмбеддинг для '{movie['name']}' уже существу��т")
-
     except Exception as e:
         logging.error(f"Ошибка при обработке фильмов: {e}")
     finally:
@@ -270,15 +266,13 @@ def get_movie_embeddings(conn):
     movie_embeddings = {}
     try:
         with conn.cursor() as cur:
-            cur.execute(f"""
-            SELECT m.id, m.data, e.embedding 
-            FROM {movies_table} m
-            JOIN {embeddings_table} e ON m.id = e.movie_id
-            """)
-            for movie_id, movie_data, embedding in cur.fetchall():
-                movie_info = json.loads(movie_data)
-                movie_name = movie_info.get('name', '')
-                movie_embeddings[movie_name] = normalize(np.array(embedding).reshape(1, -1))[0]
+            cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
+            for movie_id, embedding in cur.fetchall():
+                # Находим название фильма по ID
+                for movie in movies_data:
+                    if movie['id'] == movie_id:
+                        movie_embeddings[movie['name']] = normalize(np.array(embedding).reshape(1, -1))[0]
+                        break
         logging.info(f"Загружено {len(movie_embeddings)} эмбеддингов фильмов.")
     except Exception as e:
         logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
@@ -288,72 +282,89 @@ def search_movies(query, top_k=10):
     """Выполняет поиск фильмов по запросу."""
     global search_in_progress
     search_in_progress = True
+    start_time = time.time()
 
     try:
         conn = get_db_connection()
         if conn is None:
-            return []
-
-        # Загружаем эмбеддинги фильмов
-        movie_embeddings = get_movie_embeddings(conn)
-
-        # Получаем эмбеддинг запроса
-        query_embedding = encode_string(query)
-
-        # Выполняем поиск
-        scores = util.dot_score(query_embedding, list(movie_embeddings.values())).cpu().tolist()[0]
-        top_results = sorted(zip(scores, movie_embeddings.keys()), key=lambda x: x[0], reverse=True)[:top_k]
-
-        results = []
-        for score, movie_name in top_results:
-            # Получаем полную информацию о фильме из базы данных
+            return "<p>Ошибка подключения к базе данных</p>"
+
+        query_crc32 = calculate_crc32(query)
+        query_embedding = get_embedding_from_db(conn, query_cache_table, "query_crc32", query_crc32, model_name)
+
+        if query_embedding is None:
+            query_embedding = encode_string(query)
+
+            try:
+                with conn.cursor() as cur:
+                    cur.execute(f"""
+                        INSERT INTO {query_cache_table} (query_crc32, query, model_name, embedding)
+                        VALUES (%s, %s, %s, %s)
+                        ON CONFLICT (query_crc32) DO NOTHING
+                    """, (query_crc32, query, model_name, query_embedding.tolist()))
+                conn.commit()
+                logging.info(f"Сохранен новый эмбеддинг запроса: {query}")
+            except Exception as e:
+                logging.error(f"Ошибка при сохранении эмбеддинга запроса: {e}")
+                conn.rollback()
+
+        # Используем косинусное расстояние для поиска
+        try:
             with conn.cursor() as cur:
-                cur.execute(f"SELECT data FROM {movies_table} WHERE data->>'name' = %s", (movie_name,))
-                movie_data = cur.fetchone()
-                if movie_data:
-                    movie_info = json.loads(movie_data[0])
-                    results.append({
-                        'name': movie_name,
-                        'description': movie_info.get('description', ''),
-                        'genres': [genre['name'] for genre in movie_info.get('genres', [])],
-                        'score': f"{score:.2f}"
-                    })
-
-        return results
+                cur.execute(f"""
+                    WITH query_embedding AS (
+                        SELECT embedding
+                        FROM {query_cache_table}
+                        WHERE query_crc32 = %s
+                    )
+                    SELECT m.movie_id, 1 - (m.embedding <=> (SELECT embedding FROM query_embedding)) as similarity
+                    FROM {embeddings_table} m, query_embedding
+                    ORDER BY similarity DESC
+                    LIMIT %s
+                """, (query_crc32, top_k))
+
+                results = cur.fetchall()
+            logging.info(f"Найдено {len(results)} результатов поиска.")
+        except Exception as e:
+            logging.error(f"Ошибка при выполнении поискового запроса: {e}")
+            results = []
+
+        output = ""
+        for movie_id, similarity in results:
+            # Находим фильм по ID
+            movie = next((m for m in movies_data if m['id'] == movie_id), None)
+            if movie:
+                output += f"<h3>{movie['name']} ({movie['year']})</h3>\n"
+                output += f"<p><strong>Жанры:</strong> {', '.join(movie['genresList'])}</p>\n"
+                output += f"<p><strong>Описание:</strong> {movie['description']}</p>\n"
+                output += f"<p><strong>Релевантность:</strong> {similarity:.4f}</p>\n"
+                output += "<hr>\n"
+
+        search_time = time.time() - start_time
+        logging.info(f"Поиск выполнен за {search_time:.2f} секунд.")
+
+        return f"<p>Время поиска: {search_time:.2f} сек</p>{output}"
+
     except Exception as e:
-        logging.error(f"Ошибка при поиске фильмов: {e}")
-        return []
+        logging.error(f"Ошибка при выполнении поиска: {e}")
+        return "<p>Произошла ошибка при выполнении поиска.</p>"
+
     finally:
         if conn:
             conn.close()
         search_in_progress = False
 
-def start_processing():
-    """Запускает обработку фильмов в отдельном потоке."""
-    thread = threading.Thread(target=process_movies)
-    thread.start()
-
-# Запускаем обработку фильмов при старте приложения
-start_processing()
-
-# Функция для интерфейса Gradio
-def search_interface(query):
-    results = search_movies(query)
-    output = ""
-    for movie in results:
-        output += f"Название: {movie['name']}\n"
-        output += f"Жанры: {', '.join(movie['genres'])}\n"
-        output += f"Описание: {movie['description']}\n"
-        output += f"Оценка: {movie['score']}\n\n"
-    return output
+# Запускаем обработку фильмов в отдельном потоке
+processing_thread = threading.Thread(target=process_movies)
+processing_thread.start()
 
 # Создаем интерфейс Gradio
 iface = gr.Interface(
-    fn=search_interface,
-    inputs="text",
-    outputs="text",
-    title="Поиск фильмов",
-    description="Введите запрос для поиска фильмов"
+    fn=search_movies,
+    inputs=gr.Textbox(lines=2, placeholder="Введите запрос для поиска фильмов..."),
+    outputs=gr.HTML(label="Результаты поиска"),
+    title="Семантический поиск фильмов",
+    description="Введите описание фильма, который вы ищете, и система найдет наиболее похожие фильмы."
 )
 
 # Запускаем интерфейс