Update app.py

app.py

@@ -7,7 +7,6 @@ import queue
 import torch
 import psycopg2
 import zlib
-import numpy as np
 from urllib.parse import urlparse
 
 # Настройки базы данных PostgreSQL
@@ -76,32 +75,37 @@ def setup_database():
         return
     
     with conn.cursor() as cur:
-        # Создаем расширение pgvector
+        # Создаем расширение 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} (
+            CREATE TABLE {embeddings_table} (
                 movie_id INTEGER PRIMARY KEY,
                 embedding_crc32 BIGINT,
                 string_crc32 BIGINT,
                 model_name TEXT,
-                embedding float8[]
+                embedding vector(1024)
             );
-            CREATE INDEX IF NOT EXISTS idx_movie_embeddings_crc32 ON {embeddings_table} (string_crc32);
+            CREATE INDEX ON {embeddings_table} USING ivfflat (embedding vector_cosine_ops);
+            CREATE INDEX ON {embeddings_table} (string_crc32);
         """)
         
         # Создаем таблицу для кэширования запросов
         cur.execute(f"""
-            CREATE TABLE IF NOT EXISTS {query_cache_table} (
+            CREATE TABLE {query_cache_table} (
                 query_crc32 BIGINT PRIMARY KEY,
                 query TEXT,
                 model_name TEXT,
-                embedding float8[],
+                embedding vector(1024),
                 created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
             );
-            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);
+            CREATE INDEX ON {query_cache_table} USING ivfflat (embedding vector_cosine_ops);
+            CREATE INDEX ON {query_cache_table} (query_crc32);
+            CREATE INDEX ON {query_cache_table} (created_at);
         """)
     
     conn.commit()
@@ -138,14 +142,6 @@ def get_movies_without_embeddings():
     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):
     """Получает эмбеддинг из базы данных."""
     with conn.cursor() as cur:
@@ -153,12 +149,11 @@ def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_nam
                    (crc32_value, model_name))
         result = cur.fetchone()
         if result and result[0]:
-            return list_to_vector(result[0])
+            return torch.tensor(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"""
@@ -166,7 +161,7 @@ def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32,
                 (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))
+            """, (movie_id, embedding_crc32, string_crc32, model_name, embedding.tolist()))
             conn.commit()
             return True
         except Exception as e:
@@ -222,7 +217,7 @@ def process_movies():
             
             if existing_embedding is None:
                 embedding = encode_string(embedding_string)
-                embedding_crc32 = calculate_crc32(str(vector_to_list(embedding)))
+                embedding_crc32 = calculate_crc32(str(embedding.tolist()))
                 
                 if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
                     print(f"Сохранен эмбеддинг для '{movie['name']}'")
@@ -239,15 +234,12 @@ def get_movie_embeddings(conn):
     """Загружает все эмбеддинги фильмов из базы данных."""
     movie_embeddings = {}
     with conn.cursor() as cur:
-        cur.execute(f"""
-            SELECT e.movie_id, e.embedding 
-            FROM {embeddings_table} e
-        """)
+        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']] = list_to_vector(embedding)
+                    movie_embeddings[movie['name']] = torch.tensor(embedding)
                     break
     return movie_embeddings
 
@@ -267,29 +259,38 @@ def search_movies(query, top_k=10):
 
         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))
+                """, (query_crc32, query, model_name, query_embedding.tolist()))
                 conn.commit()
 
-        movie_embeddings = get_movie_embeddings(conn)
-        
-        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]
+        # Используем косинусное расстояние для поиска
+        with conn.cursor() as cur:
+            cur.execute(f"""
+                SELECT m.movie_id, m.embedding <=> %s as distance
+                FROM {embeddings_table} m
+                ORDER BY distance ASC
+                LIMIT %s
+            """, (query_embedding.tolist(), top_k))
+            
+            results = cur.fetchall()
 
         results_html = "<ol>"
-        for title, score in top_results:
-            results_html += f"<li><strong>{title}</strong> (Сходство: {score:.4f})</li>"
+        for movie_id, distance in results:
+            # Находим название фильма по ID
+            movie_title = None
+            for movie in movies_data:
+                if movie['id'] == movie_id:
+                    movie_title = movie['name']
+                    break
+            
+            if movie_title:
+                similarity = 1 - distance  # Конвертируем расстояние в сходство
+                results_html += f"<li><strong>{movie_title}</strong> (Сходство: {similarity:.4f})</li>"
         results_html += "</ol>"
 
         search_time = time.time() - start_time