Update app.py

app.py

@@ -1,23 +1,17 @@
+import gradio as gr
+from sentence_transformers import SentenceTransformer, util
 import os
 import time
 import threading
 import queue
-from typing import List, Dict, Any, Optional
-import logging
-from urllib.parse import urlparse
-
-import gradio as gr
 import torch
 import psycopg2
 import zlib
 import numpy as np
-from sentence_transformers import SentenceTransformer, util
+from urllib.parse import urlparse
+import logging
 from sklearn.preprocessing import normalize
 
-# Рекомендуется использовать python-dotenv для загрузки переменных окружения
-# from dotenv import load_dotenv
-# load_dotenv()
-
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
@@ -60,10 +54,12 @@ except FileNotFoundError:
     movies_data = []
 
 # Очередь для необработанных фильмов
-movies_queue: queue.Queue = queue.Queue()
+movies_queue = queue.Queue()
 
-# Флаги состояния
+# Флаг, указывающий, что обработка фильмов завершена
 processing_complete = False
+
+# Флаг, указывающий, что выполняется поиск
 search_in_progress = False
 
 # Блокировка для доступа к базе данных
@@ -83,19 +79,20 @@ def get_db_connection():
 
 def setup_database():
     """Настраивает базу данных: создает расширение, таблицы и индексы."""
-    with get_db_connection() as conn:
-        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"""
+    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 {embeddings_table} (
                     movie_id INTEGER PRIMARY KEY,
                     embedding_crc32 BIGINT,
@@ -104,10 +101,10 @@ def setup_database():
                     embedding vector(1024)
                 );
                 CREATE INDEX ON {embeddings_table} (string_crc32);
-                """)
-                
-                # Создаем таблицу для кэширования запросов
-                cur.execute(f"""
+            """)
+
+            # Создаем таблицу для кэширования запросов
+            cur.execute(f"""
                 CREATE TABLE {query_cache_table} (
                     query_crc32 BIGINT PRIMARY KEY,
                     query TEXT,
@@ -117,51 +114,60 @@ def setup_database():
                 );
                 CREATE INDEX ON {query_cache_table} (query_crc32);
                 CREATE INDEX ON {query_cache_table} (created_at);
-                """)
-            
-            conn.commit()
-            logging.info("База данных успешно настроена.")
-        except Exception as e:
-            logging.error(f"Ошибка при настройке базы данных: {e}")
-            conn.rollback()
+            """)
+
+        conn.commit()
+        logging.info("База данных успешно настроена.")
+    except Exception as e:
+        logging.error(f"Ошибка при настройке базы данных: {e}")
+        conn.rollback()
+    finally:
+        conn.close()
 
 # Настраиваем базу данных при запуске
 setup_database()
 
-def calculate_crc32(text: str) -> int:
+def calculate_crc32(text):
     """Вычисляет CRC32 для строки."""
     return zlib.crc32(text.encode('utf-8')) & 0xFFFFFFFF
 
-def encode_string(text: str) -> np.ndarray:
+def encode_string(text):
     """Кодирует строку в эмбеддинг."""
     embedding = model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
     return embedding.cpu().numpy()
 
-def get_movies_without_embeddings() -> List[Dict[str, Any]]:
+def get_movies_without_embeddings():
     """Получает список фильм��в, для которых нужно создать эмбеддинги."""
-    with get_db_connection() as conn:
-        if conn is None:
-            return []
-        try:
-            with conn.cursor() as cur:
-                # Получаем список ID фильмов, которые уже есть в базе
-                cur.execute(f"SELECT movie_id FROM {embeddings_table}")
-                existing_ids = {row[0] for row in cur.fetchall()}
-                
-                # Фильтруем только те фильмы, которых нет в базе
-                movies_to_process = [movie for movie in movies_data if movie['id'] not in existing_ids]
-                
-            logging.info(f"Найдено {len(movies_to_process)} фильмов для обработки.")
-            return movies_to_process
-        except Exception as e:
-            logging.error(f"Ошибка при получении списка фильмов для обработки: {e}")
-            return []
+    conn = get_db_connection()
+    if conn is None:
+        return []
+
+    movies_to_process = []
+    try:
+        with conn.cursor() as cur:
+            # Получаем список 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)
+
+        logging.info(f"Найдено {len(movies_to_process)} фильмов для обработки.")
+    except Exception as e:
+        logging.error(f"Ошибка при получении списка фильмов для обработки: {e}")
+    finally:
+        conn.close()
 
-def get_embedding_from_db(conn, table_name: str, crc32_column: str, crc32_value: int, model_name: str) -> Optional[np.ndarray]:
+    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]:
                 # Нормализуем эмбеддинг после извлечения из БД
@@ -170,16 +176,17 @@ def get_embedding_from_db(conn, table_name: str, crc32_column: str, crc32_value:
         logging.error(f"Ошибка при получении эмбеддинга из БД: {e}")
     return None
 
-def insert_embedding(conn, table_name: str, movie_id: int, embedding_crc32: int, string_crc32: int, embedding: np.ndarray) -> bool:
+def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32, embedding):
     """Вставляет эмбеддинг в базу данных."""
     try:
         # Нормализуем эмбеддинг перед сохранением
         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
@@ -191,10 +198,12 @@ def insert_embedding(conn, table_name: str, movie_id: int, embedding_crc32: int,
 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
@@ -204,51 +213,55 @@ def process_movies():
     for movie in movies_to_process:
         movies_queue.put(movie)
 
-    with get_db_connection() as conn:
-        if conn is None:
-            processing_complete = True
-            return
-        try:
-            while not movies_queue.empty():
-                if search_in_progress:
-                    time.sleep(1)
-                    continue
-
-                batch = []
-                while not movies_queue.empty() and len(batch) < batch_size:
-                    try:
-                        movie = movies_queue.get_nowait()
-                        batch.append(movie)
-                    except queue.Empty:
-                        break
+    conn = get_db_connection()
+    if conn is None:
+        processing_complete = True
+        return
 
-                if not batch:
+    try:
+        while not movies_queue.empty():
+            if search_in_progress:
+                time.sleep(1)
+                continue
+
+            batch = []
+            while not movies_queue.empty() and len(batch) < batch_size:
+                try:
+                    movie = movies_queue.get_nowait()
+                    batch.append(movie)
+                except queue.Empty:
                     break
 
-                logging.info(f"Обработка пакета из {len(batch)} фильмов...")
-                for movie in batch:
-                    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']}' уже существует")
+            if not batch:
+                break
 
-        except Exception as e:
-            logging.error(f"Ошибка при обработке фильмов: {e}")
-        finally:
-            processing_complete = True
-            logging.info("Обработка фильмов завершена")
+            logging.info(f"Обработка пакета из {len(batch)} фильмов...")
+
+            for movie in batch:
+                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:
+        conn.close()
+        processing_complete = True
+        logging.info("Обработка фильмов завершена")
 
-def get_movie_embeddings(conn) -> Dict[str, np.ndarray]:
+def get_movie_embeddings(conn):
     """Загружает все эмбеддинги фильмов из базы данных."""
     movie_embeddings = {}
     try:
@@ -265,97 +278,96 @@ def get_movie_embeddings(conn) -> Dict[str, np.ndarray]:
         logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
     return movie_embeddings
 
-def clean_query_cache(conn):
-    """Очищает устаревшие записи из кэша запросов."""
-    try:
-        with conn.cursor() as cur:
-            # Получаем общий размер кэша
-            cur.execute(f"SELECT pg_total_relation_size('{query_cache_table}')")
-            total_size = cur.fetchone()[0]
-            
-            if total_size > MAX_CACHE_SIZE:
-                # Удаляем старые записи, пока размер не стан��т меньше максимального
-                cur.execute(f"""
-                DELETE FROM {query_cache_table}
-                WHERE ctid IN (
-                    SELECT ctid
-                    FROM {query_cache_table}
-                    ORDER BY created_at ASC
-                    LIMIT (SELECT COUNT(*) / 2 FROM {query_cache_table})
-                )
-                """)
-                conn.commit()
-                logging.info("Кэш запросов очищен.")
-    except Exception as e:
-        logging.error(f"Ошибка при очистке кэша запросов: {e}")
-        conn.rollback()
-
-def search_movies(query: str, top_k: int = 10) -> List[Dict[str, Any]]:
+def search_movies(query, top_k=10):
     """Выполняет поиск фильмов по запросу."""
     global search_in_progress
     search_in_progress = True
-    
+    start_time = time.time()
+
     try:
-        with get_db_connection() as conn:
-            if conn is None:
-                return []
+        conn = get_db_connection()
+        if conn is None:
+            return "<p>Ошибка подключения к базе данных</p>"
 
-            clean_query_cache(conn)
+        query_crc32 = calculate_crc32(query)
+        query_embedding = get_embedding_from_db(conn, query_cache_table, "query_crc32", query_crc32, model_name)
 
-            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)
 
-            if query_embedding is None:
-                query_embedding = encode_string(query)
-                insert_embedding(conn, query_cache_table, -1, -1, query_crc32, query_embedding)
+            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()
 
-            movie_embeddings = get_movie_embeddings(conn)
+        # Используем косинусное расстояние для поиска
+        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 = []
 
-            # Вычисляем косинусное сходство
-            similarities = util.cos_sim(query_embedding, list(movie_embeddings.values()))[0]
+        results_html = "<ol>"
+        for movie_id, similarity in results:
+            # Находим название фильма по ID
+            movie_title = None
+            for movie in movies_data:
+                if movie['id'] == movie_id:
+                    movie_title = movie['name']
+                    break
 
-            # Сортируем результаты
-            top_results = sorted(zip(similarities, movie_embeddings.keys()), key=lambda x: x[0], reverse=True)[:top_k]
+            if movie_title:
+                results_html += f"<li><strong>{movie_title}</strong> (Сходство: {similarity:.4f})</li>"
+        results_html += "</ol>"
+
+        search_time = time.time() - start_time
+        logging.info(f"Поиск выполнен за {search_time:.2f} секунд.")
+
+        return f"<p>Время поиска: {search_time:.2f} сек</p>{results_html}"
 
-            results = []
-            for score, movie_name in top_results:
-                movie = next((m for m in movies_data if m['name'] == movie_name), None)
-                if movie:
-                    results.append({
-                        "name": movie['name'],
-                        "year": movie['year'],
-                        "genres": movie['genresList'],
-                        "description": movie['description'],
-                        "score": float(score)
-                    })
-
-        return results
     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
 
 # Запускаем обработку фильмов в отдельном потоке
-threading.Thread(target=process_movies, daemon=True).start()
+processing_thread = threading.Thread(target=process_movies)
+processing_thread.start()
 
 # Создаем интерфейс Gradio
-def gradio_search(query: str) -> str:
-    results = search_movies(query)
-    output = ""
-    for movie in results:
-        output += f"Название: {movie['name']} ({movie['year']})\n"
-        output += f"Жанры: {', '.join(movie['genres'])}\n"
-        output += f"Описание: {movie['description']}\n"
-        output += f"Релевантность: {movie['score']:.2f}\n\n"
-    return output
-
 iface = gr.Interface(
-    fn=gradio_search,
-    inputs="text",
-    outputs="text",
-    title="Поиск фильмов",
-    description="Введите запрос для поиска фильмов"
+    fn=search_movies,
+    inputs=gr.Textbox(lines=2, placeholder="Введите запрос для поиска фильмов..."),
+    outputs=gr.HTML(label="Результаты поиска"),
+    title="Семантический поиск фильмов",
+    description="Введите описание фильма, который вы ищете, и система найдет наиболее похожие фильмы."
 )
 
+# Запускаем интерфейс
 iface.launch()