Update app.py

app.py

@@ -7,6 +7,7 @@ import queue
 import torch
 import psycopg2
 import zlib
+import numpy as np
 from urllib.parse import urlparse
 
 # Настройки базы данных PostgreSQL
@@ -46,8 +47,6 @@ except FileNotFoundError:
 
 # Очередь для необработанных фильмов
 movies_queue = queue.Queue()
-for movie in movies_data:
-    movies_queue.put(movie)
 
 # Флаг, указывающий, что обработка фильмов завершена
 processing_complete = False
@@ -71,7 +70,7 @@ def get_db_connection():
         return None
 
 def setup_database():
-    """Настраивает базу данных: создает расширение, таблицы и триггер."""
+    """Настраивает базу данных: создает расширение, таблицы и индексы."""
     conn = get_db_connection()
     if conn is None:
         return
@@ -83,57 +82,32 @@ def setup_database():
         # Создаем таблицу для хранения эмбеддингов фильмов
         cur.execute(f"""
             CREATE TABLE IF NOT EXISTS {embeddings_table} (
-                movie_id INTEGER,
-                embedding_crc32 BIGINT PRIMARY KEY,
+                movie_id INTEGER PRIMARY KEY,
+                embedding_crc32 BIGINT,
                 string_crc32 BIGINT,
                 model_name TEXT,
-                embedding vector(1024)
+                embedding float8[]
             );
+            CREATE INDEX IF NOT EXISTS idx_movie_embeddings_crc32 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),
+                embedding float8[],
                 created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
             );
-            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()
-            RETURNS TRIGGER AS $$
-            DECLARE
-                table_size BIGINT;
-                row_to_delete RECORD;
-            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
-                        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};
-                    END LOOP;
-                END IF;
-                RETURN NEW;
-            END;
-            $$ LANGUAGE plpgsql;
-
-            CREATE OR REPLACE TRIGGER trg_manage_query_cache_size
-            AFTER INSERT ON {query_cache_table}
-            FOR EACH ROW
-            EXECUTE PROCEDURE manage_query_cache_size();
+            CREATE INDEX IF NOT EXISTS idx_query_cache_crc32 ON {query_cache_table} (query_crc32);
+            CREATE INDEX IF NOT EXISTS idx_query_cache_created ON {query_cache_table} (created_at);
         """)
     
     conn.commit()
     conn.close()
 
-# Настраиваем базу данных при запуске приложения
+# Настраиваем базу данных при запуске
 setup_database()
 
 def calculate_crc32(text):
@@ -144,44 +118,78 @@ def encode_string(text):
     """Кодирует строку в эмбеддинг."""
     return model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
 
+def get_movies_without_embeddings():
+    """Получает список фильмов, для которых нужно создать эмбеддинги."""
+    conn = get_db_connection()
+    if conn is None:
+        return []
+
+    movies_to_process = []
+    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)
+    
+    conn.close()
+    return movies_to_process
+
+def vector_to_list(vector):
+    """Преобразует вектор PyTorch в список float."""
+    return vector.detach().cpu().numpy().tolist()
+
+def list_to_vector(lst):
+    """Преобразует список float в вектор PyTorch."""
+    return torch.tensor(lst)
+
 def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_name):
-    """
-    Пытается получить эмбеддинг из указанной таблицы по CRC32.
-    Возвращает эмбеддинг, если найден, иначе None.
-    """
+    """Получает эмбеддинг из базы данных."""
     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:
-            return torch.tensor(result[0])
-        else:
-            return None
-
-def insert_embedding(conn, table_name, crc32_column, crc32_value, other_columns, embedding):
-    """Вставляет эмбеддинг в указанную таблицу."""
-    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())
-    
+        if result and result[0]:
+            return list_to_vector(result[0])
+    return None
+
+def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32, embedding):
+    """Вставляет эмбеддинг в базу данных."""
+    embedding_list = vector_to_list(embedding)
     with conn.cursor() as cur:
         try:
             cur.execute(f"""
-                INSERT INTO {table_name} ({columns})
-                VALUES ({placeholders})
-                ON CONFLICT ({crc32_column}) DO NOTHING;
-            """, values)
+                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, embedding_list))
             conn.commit()
             return True
         except Exception as e:
-            print(f"Ошибка при вставке эмбеддинга в таблицу {table_name}: {e}")
+            print(f"Ошибка при вставке эмбеддинга: {e}")
             conn.rollback()
             return False
 
 def process_movies():
-    """
-    Обрабатывает фильмы из очереди, создавая для них эмбеддинги и сохраняя их в базу данных.
-    """
+    """Обрабатывает фильмы, создавая для них эмбеддинги."""
     global processing_complete
+    
+    # Получаем список фильмов, которые нужно обработать
+    movies_to_process = get_movies_without_embeddings()
+    
+    if not movies_to_process:
+        print("Все фильмы уже обработаны.")
+        processing_complete = True
+        return
+
+    # Добавляем фильмы в очередь
+    for movie in movies_to_process:
+        movies_queue.put(movie)
+
     conn = get_db_connection()
     if conn is None:
         processing_complete = True
@@ -189,136 +197,108 @@ def process_movies():
 
     while True:
         if search_in_progress:
-            time.sleep(1)  # Ждем, пока поиск не завершится
+            time.sleep(1)
             continue
 
         batch = []
         while not movies_queue.empty() and len(batch) < batch_size:
             try:
-                movie = movies_queue.get(timeout=1)
+                movie = movies_queue.get_nowait()
                 batch.append(movie)
             except queue.Empty:
                 break
 
         if not batch:
-            print("Очередь фильмов пуста.")
-            processing_complete = True
             break
 
-        titles = [movie["name"] for movie in batch]
-        embedding_strings = [
-            f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
-            for movie in batch
-        ]
-        
-        print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
+        print(f"Обработка пакета из {len(batch)} фильмов...")
         
-        with db_lock:
-            for movie, embedding_string in zip(batch, embedding_strings):
-                movie_id = movie['id']
-                string_crc32 = calculate_crc32(embedding_string)
+        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(vector_to_list(embedding)))
                 
-                # Проверяем, есть ли уже эмбеддинг для этого фильма в базе данных
-                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.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:
-                        print(f"Ошибка сохранения эмбеддинга для фильма '{movie['name']}'.")
+                if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
+                    print(f"Сохранен эмбеддинг для '{movie['name']}'")
                 else:
-                    print(f"Эмбеддинг для фильма '{movie['name']}' уже существует в базе данных.")
+                    print(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
+            else:
+                print(f"Эмбеддинг для '{movie['name']}' уже существует")
 
     conn.close()
-    print("Обработка фильмов завершена.")
+    processing_complete = True
+    print("Обработка фильмов завершена")
 
 def get_movie_embeddings(conn):
     """Загружает все эмбеддинги фильмов из базы данных."""
     movie_embeddings = {}
     with conn.cursor() as cur:
-        cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
-        rows = cur.fetchall()
-        for row in rows:
-            movie_id, embedding = row
-            # Находим название фильма по его ID
+        cur.execute(f"""
+            SELECT e.movie_id, e.embedding 
+            FROM {embeddings_table} e
+        """)
+        for movie_id, embedding in cur.fetchall():
+            # Находим название фильма по ID
             for movie in movies_data:
                 if movie['id'] == movie_id:
-                    title = movie["name"]
-                    movie_embeddings[title] = torch.tensor(embedding)
+                    movie_embeddings[movie['name']] = list_to_vector(embedding)
                     break
     return movie_embeddings
 
 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')}")
-        query_embedding_tensor = encode_string(query)
-        print(f"Окончание создания эмбеддинга запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    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)
+            embedding_list = vector_to_list(query_embedding)
+            
+            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, embedding_list))
+                conn.commit()
+
+        movie_embeddings = get_movie_embeddings(conn)
         
-        # Сохраняем эмбеддинг запроса в кэш
-        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 = []
+        for title, movie_embedding in movie_embeddings.items():
+            similarity = util.pytorch_cos_sim(query_embedding, movie_embedding).item()
+            similarities.append((title, similarity))
+
+        similarities.sort(key=lambda x: x[1], reverse=True)
+        top_results = similarities[:top_k]
+
+        results_html = "<ol>"
+        for title, score in top_results:
+            results_html += f"<li><strong>{title}</strong> (Сходство: {score:.4f})</li>"
+        results_html += "</ol>"
+
+        search_time = time.time() - start_time
+        conn.close()
+        
+        return f"<p>Время поиска: {search_time:.2f} сек</p>{results_html}"
     
-    # Сортируем результаты
-    similarities.sort(key=lambda x: x[1], reverse=True)
-    top_results = similarities[:top_k]
-    print(f"Окончание вычисления косинусного сходства: {time.strftime('%Y-%m-%d %H:%M:%S')}")
-
-    # Формируем HTML-строку с результатами
-    results_html = "<ol>"
-    for title, score in top_results:
-        results_html += f"<li><strong>{title}</strong> (Сходство: {score:.4f})</li>"
-    results_html += "</ol>"
-
-    search_in_progress = False
-    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
+    finally:
+        search_in_progress = False
 
 # Запускаем обработку фильмов в отдельном потоке
 processing_thread = threading.Thread(target=process_movies)