opex792 commited on
Commit
551d569
·
verified ·
1 Parent(s): 8cbf970

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +175 -130
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
- print("Ошибка: Файл movies.json не найден.")
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
- print(f"Ошибка подключения к базе данных: {e}")
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
- with conn.cursor() as cur:
78
- # Создаем расширение pgvector если его нет
79
- cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
80
-
81
- # Удаляем существующие таблицы если они есть
82
- cur.execute(f"DROP TABLE IF EXISTS {embeddings_table}, {query_cache_table};")
83
-
84
- # Создаем таблицу для хранения эмбеддингов фильмов
85
- cur.execute(f"""
86
- CREATE TABLE {embeddings_table} (
87
- movie_id INTEGER PRIMARY KEY,
88
- embedding_crc32 BIGINT,
89
- string_crc32 BIGINT,
90
- model_name TEXT,
91
- embedding vector(1024)
92
- );
93
- CREATE INDEX ON {embeddings_table} USING ivfflat (embedding vector_cosine_ops);
94
- CREATE INDEX ON {embeddings_table} (string_crc32);
95
- """)
96
-
97
- # Создаем таблицу для кэширования запросов
98
- cur.execute(f"""
99
- CREATE TABLE {query_cache_table} (
100
- query_crc32 BIGINT PRIMARY KEY,
101
- query TEXT,
102
- model_name TEXT,
103
- embedding vector(1024),
104
- created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
105
- );
106
- CREATE INDEX ON {query_cache_table} USING ivfflat (embedding vector_cosine_ops);
107
- CREATE INDEX ON {query_cache_table} (query_crc32);
108
- CREATE INDEX ON {query_cache_table} (created_at);
109
- """)
110
-
111
- conn.commit()
112
- conn.close()
 
 
 
 
 
 
113
 
114
  # Настраиваем базу данных при запуске
115
  setup_database()
@@ -129,55 +142,66 @@ def get_movies_without_embeddings():
129
  return []
130
 
131
  movies_to_process = []
132
- with conn.cursor() as cur:
133
- # Получаем список ID фильмов, которые уже есть в базе
134
- cur.execute(f"SELECT movie_id FROM {embeddings_table}")
135
- existing_ids = {row[0] for row in cur.fetchall()}
136
-
137
- # Фильтруем только те фильмы, которых нет в базе
138
- for movie in movies_data:
139
- if movie['id'] not in existing_ids:
140
- movies_to_process.append(movie)
141
-
142
- conn.close()
 
 
 
 
 
 
143
  return movies_to_process
144
 
145
  def get_embedding_from_db(conn, table_name, crc32_column, crc32_value, model_name):
146
  """Получает эмбеддинг из базы данных."""
147
- with conn.cursor() as cur:
148
- cur.execute(f"SELECT embedding FROM {table_name} WHERE {crc32_column} = %s AND model_name = %s",
149
- (crc32_value, model_name))
150
- result = cur.fetchone()
151
- if result and result[0]:
152
- return torch.tensor(result[0])
 
 
 
153
  return None
154
 
155
  def insert_embedding(conn, table_name, movie_id, embedding_crc32, string_crc32, embedding):
156
  """Вставляет эмбеддинг в базу данных."""
157
- with conn.cursor() as cur:
158
- try:
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
- conn.commit()
166
- return True
167
- except Exception as e:
168
- print(f"Ошибка при вставке эмбеддинга: {e}")
169
- conn.rollback()
170
- return False
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
- print("Все фильмы уже обработаны.")
181
  processing_complete = True
182
  return
183
 
@@ -190,57 +214,64 @@ def process_movies():
190
  processing_complete = True
191
  return
192
 
193
- while True:
194
- if search_in_progress:
195
- time.sleep(1)
196
- continue
 
 
 
 
 
 
 
 
 
197
 
198
- batch = []
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
- if not batch:
207
- break
208
-
209
- print(f"Обработка пакета из {len(batch)} фильмов...")
210
-
211
- for movie in batch:
212
- embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
213
- string_crc32 = calculate_crc32(embedding_string)
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
- conn.close()
230
- processing_complete = True
231
- print("Обработка фильмов завершена")
 
 
 
 
 
 
 
 
 
 
 
 
 
232
 
233
  def get_movie_embeddings(conn):
234
  """Загружает все эмбеддинги фильмов из базы данных."""
235
  movie_embeddings = {}
236
- with conn.cursor() as cur:
237
- cur.execute(f"SELECT movie_id, embedding FROM {embeddings_table}")
238
- for movie_id, embedding in cur.fetchall():
239
- # Находим название фильма по ID
240
- for movie in movies_data:
241
- if movie['id'] == movie_id:
242
- movie_embeddings[movie['name']] = torch.tensor(embedding)
243
- break
 
 
 
 
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
- with conn.cursor() as cur:
264
- cur.execute(f"""
265
- INSERT INTO {query_cache_table} (query_crc32, query, model_name, embedding)
266
- VALUES (%s, %s, %s, %s)
267
- ON CONFLICT (query_crc32) DO NOTHING
268
- """, (query_crc32, query, model_name, query_embedding.tolist()))
 
269
  conn.commit()
 
 
 
270
 
271
  # Используем косинусное расстояние для поиска
272
- with conn.cursor() as cur:
273
- cur.execute(f"""
274
- SELECT m.movie_id, m.embedding <=> %s as distance
275
- FROM {embeddings_table} m
276
- ORDER BY distance ASC
277
- LIMIT %s
278
- """, (query_embedding.tolist(), top_k))
279
-
280
- results = cur.fetchall()
 
 
 
 
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
- conn.close()
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
  # Запускаем обработку фильмов в отдельном потоке