Update app.py
Browse files
app.py
CHANGED
@@ -8,6 +8,10 @@ import torch
|
|
8 |
import psycopg2
|
9 |
import zlib
|
10 |
from urllib.parse import urlparse
|
|
|
|
|
|
|
|
|
11 |
|
12 |
# Настройки базы данных PostgreSQL
|
13 |
DATABASE_URL = os.environ.get("DATABASE_URL")
|
@@ -26,7 +30,9 @@ db_params = {
|
|
26 |
|
27 |
# Загружаем модель
|
28 |
model_name = "BAAI/bge-m3"
|
|
|
29 |
model = SentenceTransformer(model_name)
|
|
|
30 |
|
31 |
# Имена таблиц
|
32 |
embeddings_table = "movie_embeddings"
|
@@ -40,8 +46,9 @@ try:
|
|
40 |
import json
|
41 |
with open("movies.json", "r", encoding="utf-8") as f:
|
42 |
movies_data = json.load(f)
|
|
|
43 |
except FileNotFoundError:
|
44 |
-
|
45 |
movies_data = []
|
46 |
|
47 |
# Очередь для необработанных фильмов
|
@@ -65,7 +72,7 @@ def get_db_connection():
|
|
65 |
conn = psycopg2.connect(**db_params)
|
66 |
return conn
|
67 |
except Exception as e:
|
68 |
-
|
69 |
return None
|
70 |
|
71 |
def setup_database():
|
@@ -73,43 +80,49 @@ def setup_database():
|
|
73 |
conn = get_db_connection()
|
74 |
if conn is None:
|
75 |
return
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
|
80 |
-
|
81 |
-
|
82 |
-
|
83 |
-
|
84 |
-
|
85 |
-
|
86 |
-
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
113 |
|
114 |
# Настраиваем базу данных при запуске
|
115 |
setup_database()
|
@@ -129,55 +142,66 @@ def get_movies_without_embeddings():
|
|
129 |
return []
|
130 |
|
131 |
movies_to_process = []
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
143 |
return movies_to_process
|
144 |
|
145 |
def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_name):
|
146 |
"""Получает эмбеддинг из базы данных."""
|
147 |
-
|
148 |
-
|
149 |
-
|
150 |
-
|
151 |
-
|
152 |
-
|
|
|
|
|
|
|
153 |
return None
|
154 |
|
155 |
def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32, embedding):
|
156 |
"""Вставляет эмбеддинг в базу данных."""
|
157 |
-
|
158 |
-
|
159 |
cur.execute(f"""
|
160 |
INSERT INTO {table_name}
|
161 |
(movie_id, embedding_crc32, string_crc32, model_name, embedding)
|
162 |
VALUES (%s, %s, %s, %s, %s)
|
163 |
ON CONFLICT (movie_id) DO NOTHING
|
164 |
""", (movie_id, embedding_crc32, string_crc32, model_name, embedding.tolist()))
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
|
172 |
def process_movies():
|
173 |
"""Обрабатывает фильмы, создавая для них эмбеддинги."""
|
174 |
global processing_complete
|
175 |
-
|
|
|
|
|
176 |
# Получаем список фильмов, которые нужно обработать
|
177 |
movies_to_process = get_movies_without_embeddings()
|
178 |
-
|
179 |
if not movies_to_process:
|
180 |
-
|
181 |
processing_complete = True
|
182 |
return
|
183 |
|
@@ -190,57 +214,64 @@ def process_movies():
|
|
190 |
processing_complete = True
|
191 |
return
|
192 |
|
193 |
-
|
194 |
-
|
195 |
-
|
196 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
197 |
|
198 |
-
|
199 |
-
while not movies_queue.empty() and len(batch) < batch_size:
|
200 |
-
try:
|
201 |
-
movie = movies_queue.get_nowait()
|
202 |
-
batch.append(movie)
|
203 |
-
except queue.Empty:
|
204 |
break
|
205 |
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
# Проверяем существующий эмбеддинг
|
216 |
-
existing_embedding = get_embedding_from_db(conn, embeddings_table, "string_crc32", string_crc32, model_name)
|
217 |
-
|
218 |
-
if existing_embedding is None:
|
219 |
-
embedding = encode_string(embedding_string)
|
220 |
-
embedding_crc32 = calculate_crc32(str(embedding.tolist()))
|
221 |
-
|
222 |
-
if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
|
223 |
-
print(f"Сохранен эмбеддинг для '{movie['name']}'")
|
224 |
-
else:
|
225 |
-
print(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
|
226 |
-
else:
|
227 |
-
print(f"Эмбеддинг для '{movie['name']}' уже существует")
|
228 |
|
229 |
-
|
230 |
-
|
231 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
232 |
|
233 |
def get_movie_embeddings(conn):
|
234 |
"""Загружает все эмбеддинги фильмов из базы данных."""
|
235 |
movie_embeddings = {}
|
236 |
-
|
237 |
-
|
238 |
-
|
239 |
-
|
240 |
-
|
241 |
-
|
242 |
-
|
243 |
-
|
|
|
|
|
|
|
|
|
244 |
return movie_embeddings
|
245 |
|
246 |
def search_movies(query, top_k=10):
|
@@ -259,25 +290,33 @@ def search_movies(query, top_k=10):
|
|
259 |
|
260 |
if query_embedding is None:
|
261 |
query_embedding = encode_string(query)
|
262 |
-
|
263 |
-
|
264 |
-
|
265 |
-
|
266 |
-
|
267 |
-
|
268 |
-
|
|
|
269 |
conn.commit()
|
|
|
|
|
|
|
270 |
|
271 |
# Используем косинусное расстояние для поиска
|
272 |
-
|
273 |
-
|
274 |
-
|
275 |
-
|
276 |
-
|
277 |
-
|
278 |
-
|
279 |
-
|
280 |
-
|
|
|
|
|
|
|
|
|
281 |
|
282 |
results_html = "<ol>"
|
283 |
for movie_id, distance in results:
|
@@ -287,18 +326,24 @@ def search_movies(query, top_k=10):
|
|
287 |
if movie['id'] == movie_id:
|
288 |
movie_title = movie['name']
|
289 |
break
|
290 |
-
|
291 |
if movie_title:
|
292 |
similarity = 1 - distance # Конвертируем расстояние в сходство
|
293 |
results_html += f"<li><strong>{movie_title}</strong> (Сходство: {similarity:.4f})</li>"
|
294 |
results_html += "</ol>"
|
295 |
|
296 |
search_time = time.time() - start_time
|
297 |
-
|
298 |
-
|
299 |
return f"<p>Время поиска: {search_time:.2f} сек</p>{results_html}"
|
300 |
-
|
|
|
|
|
|
|
|
|
301 |
finally:
|
|
|
|
|
302 |
search_in_progress = False
|
303 |
|
304 |
# Запускаем обработку фильмов в отдельном потоке
|
|
|
8 |
import psycopg2
|
9 |
import zlib
|
10 |
from urllib.parse import urlparse
|
11 |
+
import logging
|
12 |
+
|
13 |
+
# Настройка логирования
|
14 |
+
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
15 |
|
16 |
# Настройки базы данных PostgreSQL
|
17 |
DATABASE_URL = os.environ.get("DATABASE_URL")
|
|
|
30 |
|
31 |
# Загружаем модель
|
32 |
model_name = "BAAI/bge-m3"
|
33 |
+
logging.info(f"Загрузка модели {model_name}...")
|
34 |
model = SentenceTransformer(model_name)
|
35 |
+
logging.info("Модель загружена успешно.")
|
36 |
|
37 |
# Имена таблиц
|
38 |
embeddings_table = "movie_embeddings"
|
|
|
46 |
import json
|
47 |
with open("movies.json", "r", encoding="utf-8") as f:
|
48 |
movies_data = json.load(f)
|
49 |
+
logging.info(f"Загружено {len(movies_data)} фильмов из movies.json")
|
50 |
except FileNotFoundError:
|
51 |
+
logging.error("Ошибка: Файл movies.json не найден.")
|
52 |
movies_data = []
|
53 |
|
54 |
# Очередь для необработанных фильмов
|
|
|
72 |
conn = psycopg2.connect(**db_params)
|
73 |
return conn
|
74 |
except Exception as e:
|
75 |
+
logging.error(f"Ошибка подключения к базе данных: {e}")
|
76 |
return None
|
77 |
|
78 |
def setup_database():
|
|
|
80 |
conn = get_db_connection()
|
81 |
if conn is None:
|
82 |
return
|
83 |
+
|
84 |
+
try:
|
85 |
+
with conn.cursor() as cur:
|
86 |
+
# Создаем расширение pgvector если его нет
|
87 |
+
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
88 |
+
|
89 |
+
# Удаляем существующие таблицы если они есть
|
90 |
+
cur.execute(f"DROP TABLE IF EXISTS {embeddings_table}, {query_cache_table};")
|
91 |
+
|
92 |
+
# Создаем таблицу для хранения эмбеддингов фильмов
|
93 |
+
cur.execute(f"""
|
94 |
+
CREATE TABLE {embeddings_table} (
|
95 |
+
movie_id INTEGER PRIMARY KEY,
|
96 |
+
embedding_crc32 BIGINT,
|
97 |
+
string_crc32 BIGINT,
|
98 |
+
model_name TEXT,
|
99 |
+
embedding vector(1024)
|
100 |
+
);
|
101 |
+
CREATE INDEX ON {embeddings_table} USING ivfflat (embedding vector_cosine_ops);
|
102 |
+
CREATE INDEX ON {embeddings_table} (string_crc32);
|
103 |
+
""")
|
104 |
+
|
105 |
+
# Создаем таблицу для кэширования запросов
|
106 |
+
cur.execute(f"""
|
107 |
+
CREATE TABLE {query_cache_table} (
|
108 |
+
query_crc32 BIGINT PRIMARY KEY,
|
109 |
+
query TEXT,
|
110 |
+
model_name TEXT,
|
111 |
+
embedding vector(1024),
|
112 |
+
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
113 |
+
);
|
114 |
+
CREATE INDEX ON {query_cache_table} USING ivfflat (embedding vector_cosine_ops);
|
115 |
+
CREATE INDEX ON {query_cache_table} (query_crc32);
|
116 |
+
CREATE INDEX ON {query_cache_table} (created_at);
|
117 |
+
""")
|
118 |
+
|
119 |
+
conn.commit()
|
120 |
+
logging.info("База данных успешно настроена.")
|
121 |
+
except Exception as e:
|
122 |
+
logging.error(f"Ошибка при настройке базы данных: {e}")
|
123 |
+
conn.rollback()
|
124 |
+
finally:
|
125 |
+
conn.close()
|
126 |
|
127 |
# Настраиваем базу данных при запуске
|
128 |
setup_database()
|
|
|
142 |
return []
|
143 |
|
144 |
movies_to_process = []
|
145 |
+
try:
|
146 |
+
with conn.cursor() as cur:
|
147 |
+
# Получаем список ID фильмов, которые уже есть в базе
|
148 |
+
cur.execute(f"SELECT movie_id FROM {embeddings_table}")
|
149 |
+
existing_ids = {row[0] for row in cur.fetchall()}
|
150 |
+
|
151 |
+
# Фильтруем только те фильмы, которых нет в базе
|
152 |
+
for movie in movies_data:
|
153 |
+
if movie['id'] not in existing_ids:
|
154 |
+
movies_to_process.append(movie)
|
155 |
+
|
156 |
+
logging.info(f"Найдено {len(movies_to_process)} фильмов для обработки.")
|
157 |
+
except Exception as e:
|
158 |
+
logging.error(f"Ошибка при получении списка фильмов для обработки: {e}")
|
159 |
+
finally:
|
160 |
+
conn.close()
|
161 |
+
|
162 |
return movies_to_process
|
163 |
|
164 |
def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_name):
|
165 |
"""Получает эмбеддинг из базы данных."""
|
166 |
+
try:
|
167 |
+
with conn.cursor() as cur:
|
168 |
+
cur.execute(f"SELECT embedding FROM {table_name} WHERE {crc32_column} = %s AND model_name = %s",
|
169 |
+
(crc32_value, model_name))
|
170 |
+
result = cur.fetchone()
|
171 |
+
if result and result[0]:
|
172 |
+
return torch.tensor(result[0])
|
173 |
+
except Exception as e:
|
174 |
+
logging.error(f"Ошибка при получении эмбеддинга из БД: {e}")
|
175 |
return None
|
176 |
|
177 |
def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32, embedding):
|
178 |
"""Вставляет эмбеддинг в базу данных."""
|
179 |
+
try:
|
180 |
+
with conn.cursor() as cur:
|
181 |
cur.execute(f"""
|
182 |
INSERT INTO {table_name}
|
183 |
(movie_id, embedding_crc32, string_crc32, model_name, embedding)
|
184 |
VALUES (%s, %s, %s, %s, %s)
|
185 |
ON CONFLICT (movie_id) DO NOTHING
|
186 |
""", (movie_id, embedding_crc32, string_crc32, model_name, embedding.tolist()))
|
187 |
+
conn.commit()
|
188 |
+
return True
|
189 |
+
except Exception as e:
|
190 |
+
logging.error(f"Ошибка при вставке эмбеддинга: {e}")
|
191 |
+
conn.rollback()
|
192 |
+
return False
|
193 |
|
194 |
def process_movies():
|
195 |
"""Обрабатывает фильмы, создавая для них эмбеддинги."""
|
196 |
global processing_complete
|
197 |
+
|
198 |
+
logging.info("Начало обработки фильмов.")
|
199 |
+
|
200 |
# Получаем список фильмов, которые нужно обработать
|
201 |
movies_to_process = get_movies_without_embeddings()
|
202 |
+
|
203 |
if not movies_to_process:
|
204 |
+
logging.info("Все фильмы уже обработаны.")
|
205 |
processing_complete = True
|
206 |
return
|
207 |
|
|
|
214 |
processing_complete = True
|
215 |
return
|
216 |
|
217 |
+
try:
|
218 |
+
while not movies_queue.empty():
|
219 |
+
if search_in_progress:
|
220 |
+
time.sleep(1)
|
221 |
+
continue
|
222 |
+
|
223 |
+
batch = []
|
224 |
+
while not movies_queue.empty() and len(batch) < batch_size:
|
225 |
+
try:
|
226 |
+
movie = movies_queue.get_nowait()
|
227 |
+
batch.append(movie)
|
228 |
+
except queue.Empty:
|
229 |
+
break
|
230 |
|
231 |
+
if not batch:
|
|
|
|
|
|
|
|
|
|
|
232 |
break
|
233 |
|
234 |
+
logging.info(f"Обработка пакета из {len(batch)} фильмов...")
|
235 |
+
|
236 |
+
for movie in batch:
|
237 |
+
embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
|
238 |
+
string_crc32 = calculate_crc32(embedding_string)
|
239 |
+
|
240 |
+
# Проверяем существующий эмбеддинг
|
241 |
+
existing_embedding = get_embedding_from_db(conn, embeddings_table, "string_crc32", string_crc32, model_name)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
242 |
|
243 |
+
if existing_embedding is None:
|
244 |
+
embedding = encode_string(embedding_string)
|
245 |
+
embedding_crc32 = calculate_crc32(str(embedding.tolist()))
|
246 |
+
|
247 |
+
if insert_embedding(conn, embeddings_table, movie['id'], embedding_crc32, string_crc32, embedding):
|
248 |
+
logging.info(f"Сохранен эмбеддинг для '{movie['name']}'")
|
249 |
+
else:
|
250 |
+
logging.error(f"Ошибка сохранения эмбеддинга для '{movie['name']}'")
|
251 |
+
else:
|
252 |
+
logging.info(f"Эмбеддинг для '{movie['name']}' уже существует")
|
253 |
+
except Exception as e:
|
254 |
+
logging.error(f"Ошибка при обработке фильмов: {e}")
|
255 |
+
finally:
|
256 |
+
conn.close()
|
257 |
+
processing_complete = True
|
258 |
+
logging.info("Обработка фильмов завершена")
|
259 |
|
260 |
def get_movie_embeddings(conn):
|
261 |
"""Загружает все эмбеддинги фильмов из базы данных."""
|
262 |
movie_embeddings = {}
|
263 |
+
try:
|
264 |
+
with conn.cursor() as cur:
|
265 |
+
cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
|
266 |
+
for movie_id, embedding in cur.fetchall():
|
267 |
+
# Находим название фильма по ID
|
268 |
+
for movie in movies_data:
|
269 |
+
if movie['id'] == movie_id:
|
270 |
+
movie_embeddings[movie['name']] = torch.tensor(embedding)
|
271 |
+
break
|
272 |
+
logging.info(f"Загружено {len(movie_embeddings)} эмбеддингов фильмов.")
|
273 |
+
except Exception as e:
|
274 |
+
logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
|
275 |
return movie_embeddings
|
276 |
|
277 |
def search_movies(query, top_k=10):
|
|
|
290 |
|
291 |
if query_embedding is None:
|
292 |
query_embedding = encode_string(query)
|
293 |
+
|
294 |
+
try:
|
295 |
+
with conn.cursor() as cur:
|
296 |
+
cur.execute(f"""
|
297 |
+
INSERT INTO {query_cache_table} (query_crc32, query, model_name, embedding)
|
298 |
+
VALUES (%s, %s, %s, %s)
|
299 |
+
ON CONFLICT (query_crc32) DO NOTHING
|
300 |
+
""", (query_crc32, query, model_name, query_embedding.tolist()))
|
301 |
conn.commit()
|
302 |
+
except Exception as e:
|
303 |
+
logging.error(f"Ошибка при сохранении эмбеддинга запроса: {e}")
|
304 |
+
conn.rollback()
|
305 |
|
306 |
# Используем косинусное расстояние для поиска
|
307 |
+
try:
|
308 |
+
with conn.cursor() as cur:
|
309 |
+
cur.execute(f"""
|
310 |
+
SELECT m.movie_id, m.embedding <=> %s as distance
|
311 |
+
FROM {embeddings_table} m
|
312 |
+
ORDER BY distance ASC
|
313 |
+
LIMIT %s
|
314 |
+
""", (query_embedding.tolist(), top_k))
|
315 |
+
|
316 |
+
results = cur.fetchall()
|
317 |
+
except Exception as e:
|
318 |
+
logging.error(f"Ошибка при выполнении поискового запроса: {e}")
|
319 |
+
results = []
|
320 |
|
321 |
results_html = "<ol>"
|
322 |
for movie_id, distance in results:
|
|
|
326 |
if movie['id'] == movie_id:
|
327 |
movie_title = movie['name']
|
328 |
break
|
329 |
+
|
330 |
if movie_title:
|
331 |
similarity = 1 - distance # Конвертируем расстояние в сходство
|
332 |
results_html += f"<li><strong>{movie_title}</strong> (Сходство: {similarity:.4f})</li>"
|
333 |
results_html += "</ol>"
|
334 |
|
335 |
search_time = time.time() - start_time
|
336 |
+
logging.info(f"Поиск выполнен за {search_time:.2f} секунд.")
|
337 |
+
|
338 |
return f"<p>Время поиска: {search_time:.2f} сек</p>{results_html}"
|
339 |
+
|
340 |
+
except Exception as e:
|
341 |
+
logging.error(f"Ошибка при выполнении поиска: {e}")
|
342 |
+
return "<p>Произошла ошибка при выполнении поиска.</p>"
|
343 |
+
|
344 |
finally:
|
345 |
+
if conn:
|
346 |
+
conn.close()
|
347 |
search_in_progress = False
|
348 |
|
349 |
# Запускаем обработку фильмов в отдельном потоке
|