Update app.py

app.py

@@ -11,6 +11,7 @@ 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')
@@ -39,27 +40,16 @@ 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
 
 # Блокировка для доступа к базе данных
@@ -82,40 +72,49 @@ 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"DROP TABLE IF EXISTS {embeddings_table}, {query_cache_table};")
 
             # Создаем таблицу для хранения эмбеддингов фильмов
             cur.execute(f"""
-                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);
+            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);
             """)
-
+            
             # Создаем таблицу для кэширования запросов
             cur.execute(f"""
-                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);
+            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}'
+            );
+            """)
+            if not cur.fetchone()[0]:
+                logging.error(f"Таблица {movies_table} не существует в базе данных.")
+            
         conn.commit()
         logging.info("База данных успешно настроена.")
     except Exception as e:
@@ -141,33 +140,43 @@ 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()}
-
-            # Фильтруем только те фильмы, которых нет в базе
-            for movie in movies_data:
-                if movie['id'] not in existing_ids:
-                    movies_to_process.append(movie)
-
+            
+            # Получаем фильмы из таблицы 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', [])]
+                })
+            
         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]:
                 # Нормализуем эмбеддинг после извлечения из БД
@@ -183,10 +192,9 @@ 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
@@ -198,12 +206,10 @@ 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
@@ -236,24 +242,22 @@ def process_movies():
                 break
 
             logging.info(f"Обработка пакета из {len(batch)} фильмов...")
-
             for movie in batch:
-                embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
+                embedding_string = f"Название: {movie['name']}\nЖанры: {', '.join(movie['genres'])}\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:
@@ -266,13 +270,15 @@ def get_movie_embeddings(conn):
     movie_embeddings = {}
     try:
         with conn.cursor() as cur:
-            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
+            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]
         logging.info(f"Загружено {len(movie_embeddings)} эмбеддингов фильмов.")
     except Exception as e:
         logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
@@ -282,89 +288,72 @@ 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 "<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"""
-                    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> {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}"
+            return []
 
-    except Exception as e:
-        logging.error(f"Ошибка при выполнении поиска: {e}")
-        return "<p>Произошла ошибка при выполнении поиска.</p>"
+        # Загружаем эмбеддинги фильмов
+        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:
+            # Получаем полную информацию о фильме из базы данных
+            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
+    except Exception as e:
+        logging.error(f"Ошибка при поиске фильмов: {e}")
+        return []
     finally:
         if conn:
             conn.close()
         search_in_progress = False
-        
-# Запускаем обработку фильмов в отдельном потоке
-processing_thread = threading.Thread(target=process_movies)
-processing_thread.start()
+
+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
 
 # Создаем интерфейс Gradio
 iface = gr.Interface(
-    fn=search_movies,
-    inputs=gr.Textbox(lines=2, placeholder="Введите запрос для поиска фильмов..."),
-    outputs=gr.HTML(label="Результаты поиска"),
-    title="Семантический поиск фильмов",
-    description="Введите описание фильма, который вы ищете, и система найдет наиболее похожие фильмы."
+    fn=search_interface,
+    inputs="text",
+    outputs="text",
+    title="Поиск фильмов",
+    description="Введите запрос для поиска фильмов"
 )
 
 # Запускаем интерфейс