Update app.py

app.py

@@ -1,78 +1,231 @@
 import gradio as gr
 from sentence_transformers import SentenceTransformer, util
-import json
 import os
 import time
 import threading
 import queue
 import torch
+import psycopg2
+import zlib
+from urllib.parse import urlparse
+
+# Настройки базы данных PostgreSQL
+DATABASE_URL = os.environ.get("postgres://avnadmin:[email protected]:22054/Kinopoisk")
+if DATABASE_URL is None:
+    raise ValueError("DATABASE_URL environment variable not set.")
+
+parsed_url = urlparse(DATABASE_URL)
+db_params = {
+    "host": parsed_url.hostname,
+    "port": parsed_url.port,
+    "database": parsed_url.path.lstrip("/"),
+    "user": parsed_url.username,
+    "password": parsed_url.password,
+    "sslmode": "require"
+}
 
 # Загружаем модель
-model_name = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5"
+model_name = "BAAI/bge-m3"
 model = SentenceTransformer(model_name)
-# model.max_seq_length = 8192  # Убираем явное ограничение длины последовательности
 
-# Имя файла для сохранения эмбеддингов
-embeddings_file = f"movie_embeddings_{model_name.replace('/', '_')}.json"
-# Имя файла для сохранения эмбеддингов запросов
-query_embeddings_file = f"query_embeddings_{model_name.replace('/', '_')}.json"
+# Имена таблиц
+embeddings_table = "movie_embeddings"
+query_cache_table = "query_cache"
+
+# Максимальный размер таблицы кэша запросов в байтах (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)
 except FileNotFoundError:
     print("Ошибка: Файл movies.json не найден.")
     movies_data = []
 
-# Загружаем эмбеддинги фильмов
-if os.path.exists(embeddings_file):
-    with open(embeddings_file, "r", encoding="utf-8") as f:
-        movie_embeddings = json.load(f)
-    print("Загружены эмбеддинги фильмов из файла.")
-else:
-    movie_embeddings = {}
-
-# Загружаем эмбеддинги запросов
-if os.path.exists(query_embeddings_file):
-    with open(query_embeddings_file, "r", encoding="utf-8") as f:
-        query_embeddings = json.load(f)
-    print("Загружены эмбеддинги запросов из файла.")
-else:
-    query_embeddings = {}
-
 # Очередь для необработанных фильмов
 movies_queue = queue.Queue()
 for movie in movies_data:
-    if movie["name"] not in movie_embeddings:
-        movies_queue.put(movie)
+    movies_queue.put(movie)
 
 # Флаг, указывающий, что обработка фильмов завершена
 processing_complete = False
 # Флаг, указывающий, что выполняется поиск
 search_in_progress = False
 
-# Блокировка для доступа к movie_embeddings
-movie_embeddings_lock = threading.Lock()
+# Блокировка для доступа к базе данных
+db_lock = threading.Lock()
 
 # Размер пакета для обработки эмбеддингов
-batch_size = 32  # Увеличиваем размер пакета в 2 раза
-
-# Инструкция для запроса
-query_prompt = "Инструкция: Найди релевантные фильмы по запросу. \n Запрос: "
-
-def encode_string(text, prompt=None):
-    """Кодирует строку в эмбеддинг с использованием инструкции, если она задана."""
-    if prompt:
-        return model.encode(text, prompt=prompt, convert_to_tensor=True, normalize_embeddings=True)
-    else:
-        return model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
+batch_size = 32
+
+def get_db_connection():
+    """Устанавливает соединение с базой данных."""
+    try:
+        conn = psycopg2.connect(**db_params)
+        return conn
+    except Exception as e:
+        print(f"Ошибка подключения к базе данных: {e}")
+        return None
+
+def create_embeddings_table():
+    """Создает таблицу для хранения эмбеддингов фильмов, если она не существует."""
+    conn = get_db_connection()
+    if conn is None:
+        return
+
+    with conn.cursor() as cur:
+        cur.execute(f"""
+            CREATE TABLE IF NOT EXISTS {embeddings_table} (
+                movie_id INTEGER,
+                embedding_crc32 BIGINT PRIMARY KEY,
+                string_crc32 BIGINT,
+                model_name TEXT,
+                embedding vector(1024)
+            );
+        """)
+        conn.commit()
+    conn.close()
+
+def create_query_cache_table():
+    """Создает таблицу для кэширования эмбеддингов запросов, если она не существует."""
+    conn = get_db_connection()
+    if conn is None:
+        return
+
+    with conn.cursor() as cur:
+        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_query_crc32 ON {query_cache_table} (query_crc32);
+            CREATE INDEX IF NOT EXISTS idx_created_at ON {query_cache_table} (created_at);
+        """)
+        conn.commit()
+    conn.close()
+
+def create_trigger_function():
+    """Создает функцию и триггер для автоматического удаления старых записей из таблицы кэша запросов"""
+    conn = get_db_connection()
+    if conn:
+        with conn.cursor() as cur:
+            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();
+            """)
+            conn.commit()
+        conn.close()
+
+# Создаем таблицы, индексы и триггер при запуске приложения
+create_embeddings_table()
+create_query_cache_table()
+create_trigger_function()
+
+def calculate_crc32(text):
+    """Вычисляет CRC32 для строки."""
+    return zlib.crc32(text.encode('utf-8')) & 0xFFFFFFFF
+
+def encode_string(text):
+    """Кодирует строку в эмбеддинг."""
+    return model.encode(text, convert_to_tensor=True, normalize_embeddings=True)
+
+def insert_embedding(conn, movie_id, embedding_string, model_name, embedding):
+    """Вставляет эмбеддинг фильма в базу данных."""
+    embedding_crc32 = calculate_crc32(str(embedding.tolist()))
+    string_crc32 = calculate_crc32(embedding_string)
+    with conn.cursor() as cur:
+        try:
+            cur.execute(
+                f"""
+                INSERT INTO {embeddings_table} (movie_id, embedding_crc32, string_crc32, model_name, embedding)
+                VALUES (%s, %s, %s, %s, %s)
+                ON CONFLICT (embedding_crc32) DO NOTHING;
+                """,
+                (movie_id, embedding_crc32, string_crc32, model_name, embedding.tolist())
+            )
+            conn.commit()
+            return True
+        except Exception as e:
+            print(f"Ошибка при вставке эмбеддинга фильма: {e}")
+            conn.rollback()
+            return False
+
+def insert_query_embedding(conn, query, model_name, embedding):
+    """Вставляет эмбеддинг запроса в таблицу кэша."""
+    query_crc32 = calculate_crc32(query)
+    with conn.cursor() as cur:
+        try:
+            cur.execute(
+                f"""
+                INSERT INTO {query_cache_table} (query_crc32, query, model_name, embedding)
+                VALUES (%s, %s, %s, %s)
+                ON CONFLICT (query_crc32) DO UPDATE SET created_at = DEFAULT; 
+                """,
+                (query_crc32, query, model_name, embedding.tolist())
+            )
+            conn.commit()
+            print(f"Эмбеддинг для запроса '{query}' сохранен в кэше.")
+            return True
+        except Exception as e:
+            print(f"Ошибка при вставке эмбеддинга запроса: {e}")
+            conn.rollback()
+            return False
+
+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
+            for movie in movies_data:
+                if movie['id'] == movie_id:
+                    title = movie["name"]
+                    movie_embeddings[title] = torch.tensor(embedding)
+                    break
+    return movie_embeddings
 
 def process_movies():
     """
-    Обрабатывает фильмы из очереди, создавая для них эмбеддинги.
+    Обрабатывает фильмы из очереди, создавая для них эмбеддинги и сохраняя их в базу данных.
     """
     global processing_complete
+    conn = get_db_connection()
+    if conn is None:
+        processing_complete = True
+        return
+
     while True:
         if search_in_progress:
             time.sleep(1)  # Ждем, пока поиск не завершится
@@ -98,40 +251,36 @@ def process_movies():
         ]
 
         print(f"Создаются эмбеддинги для фильмов: {', '.join(titles)}...")
-        embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True).tolist()
+        embeddings = model.encode(embedding_strings, convert_to_tensor=True, batch_size=batch_size, normalize_embeddings=True)
 
-        with movie_embeddings_lock:
-            for title, embedding in zip(titles, embeddings):
-                movie_embeddings[title] = embedding
-            # Сохраняем эмбеддинги в файл после обработки каждого пакета
-            with open(embeddings_file, "w", encoding="utf-8") as f:
-                json.dump(movie_embeddings, f, ensure_ascii=False, indent=4)
-            print(f"Эмбеддинги для фильмов: {', '.join(titles)} созданы и сохранены.")
+        with db_lock:
+            for movie, embedding, embedding_string in zip(batch, embeddings, embedding_strings):
+                if insert_embedding(conn, movie['id'], embedding_string, model_name, embedding):
+                    print(f"Эмбеддинг для фильма '{movie['name']}' сохранен в базе данных.")
+                else:
+                    print(f"Ошибка сохранения эмбеддинга для фильма '{movie['name']}'.")
 
+    conn.close()
     print("Обработка фильмов завершена.")
 
-def get_query_embedding(query):
+def get_query_embedding_from_db(conn, query):
     """
-    Возвращает эмбеддинг для запроса с инструкцией. 
-    Если эмбеддинг уже создан, возвращает его из словаря.
-    Иначе создает эмбеддинг, сохраняет его и возвращает.
+    Пытается получить эмбеддинг запроса из базы данных по CRC32.
+    Возвращает эмбеддинг, если найден, иначе None.
     """
-    if query in query_embeddings:
-        print(f"Эмбеддинг для запроса '{query}' уже существует.")
-        return query_embeddings[query]
-    else:
-        print(f"Создается эмбеддинг для запроса '{query}'...")
-        embedding = encode_string(query, prompt=query_prompt).tolist()
-        query_embeddings[query] = embedding
-        # Сохраняем эмбеддинги запросов в файл
-        with open(query_embeddings_file, "w", encoding="utf-8") as f:
-            json.dump(query_embeddings, f, ensure_ascii=False, indent=4)
-        print(f"Эмбеддинг для запроса '{query}' создан и сохранен.")
-        return embedding
+    query_crc32 = calculate_crc32(query)
+    with conn.cursor() as cur:
+        cur.execute(f"SELECT embedding FROM {query_cache_table} WHERE query_crc32 = %s AND model_name = %s", (query_crc32, model_name))
+        result = cur.fetchone()
+        if result:
+            print(f"Эмбеддинг для запроса '{query}' найден в кэше.")
+            return torch.tensor(result[0])
+        else:
+            return None
 
 def search_movies(query, top_k=10):
     """
-    Ищет наиболее похожие фильмы по запросу с использованием инструкции.
+    Ищет наиболее похожие фильмы по запросу.
 
     Args:
         query: Текстовый запрос.
@@ -145,12 +294,24 @@ def search_movies(query, top_k=10):
     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>"
+
     print(f"Начало соз��ания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
-    query_embedding_tensor = encode_string(query, prompt=query_prompt)
+    query_embedding_tensor = get_query_embedding_from_db(conn, query)
+
+    if query_embedding_tensor is None:
+        query_embedding_tensor = encode_string(query)
+        # Вставляем эмбеддинг запроса в базу данных
+        insert_query_embedding(conn, query, model_name, query_embedding_tensor)
     print(f"Окончание создания эмбеддинга для запроса: {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
-    with movie_embeddings_lock:
-        current_movie_embeddings = movie_embeddings.copy()
+    with db_lock:
+        current_movie_embeddings = get_movie_embeddings(conn)
+
+    conn.close()
 
     if not current_movie_embeddings:
         search_in_progress = False
@@ -158,7 +319,7 @@ def search_movies(query, top_k=10):
 
     # Преобразуем эмбеддинги фильмов в тензор
     movie_titles = list(current_movie_embeddings.keys())
-    movie_embeddings_tensor = torch.tensor(list(current_movie_embeddings.values()))
+    movie_embeddings_tensor = torch.stack(list(current_movie_embeddings.values()))
 
     print(f"Начало поиска похожих фильмов: {time.strftime('%Y-%m-%d %H:%M:%S')}")
     # Используем util.semantic_search для поиска похожих фильмов