Update app.py

app.py

@@ -12,6 +12,7 @@ from urllib.parse import urlparse
 import logging
 from sklearn.preprocessing import normalize
 from concurrent.futures import ThreadPoolExecutor
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -31,12 +32,20 @@ db_params = {
     "sslmode": "require"
 }
 
-# Загружаем модель
+# Загружаем модель-энкодер
 model_name = "BAAI/bge-m3"
 logging.info(f"Загрузка модели {model_name}...")
 model = SentenceTransformer(model_name)
 logging.info("Модель загружена успешно.")
 
+# Загружаем модель-реранкер
+reranker_name = 'BAAI/bge-reranker-v2-m3'
+logging.info(f"Загрузка модели реранкера {reranker_name}...")
+reranker_tokenizer = AutoTokenizer.from_pretrained(reranker_name)
+reranker_model = AutoModelForSequenceClassification.from_pretrained(reranker_name)
+reranker_model.eval()
+logging.info("Модель реранкера загружена успешно.")
+
 # Имена таблиц
 embeddings_table = "movie_embeddings"
 query_cache_table = "query_cache"
@@ -207,7 +216,7 @@ def process_batch(batch):
 
     try:
         for movie in batch:
-            embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genreslist']}\nОписание: {movie['description']}"
+            embedding_string = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
             string_crc32 = calculate_crc32(embedding_string)
 
             # Проверяем существующий эмбеддинг
@@ -289,6 +298,27 @@ def get_movie_embeddings(conn):
         logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
     return movie_embeddings
 
+def rerank_results(query, results, conn):
+    """Ранжирует результаты поиска с помощью реранкера."""
+    if not results:
+        return []
+    
+    pairs = []
+    movie_ids = []
+    for movie_id, _ in results:
+        movie = next((m for m in movies_data if m['id'] == movie_id), None)
+        if movie:
+            movie_info = f"Название: {movie['name']}\nГод: {movie['year']}\nЖанры: {movie['genresList']}\nОписание: {movie['description']}"
+            pairs.append([query, movie_info])
+            movie_ids.append(movie_id)
+
+    with torch.no_grad():
+        inputs = reranker_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
+        scores = reranker_model(**inputs, return_dict=True).logits.view(-1, ).float().cpu().numpy()
+    
+    reranked_results = sorted(zip(movie_ids, scores), key=lambda x: x[1], reverse=True)
+    return reranked_results
+
 def search_movies(query, top_k=20):
     """Выполняет поиск фильмов по запросу."""
     global search_in_progress
@@ -332,23 +362,29 @@ def search_movies(query, top_k=20):
                     FROM {embeddings_table} m, query_embedding
                     ORDER BY similarity DESC
                     LIMIT %s
-                """, (query_crc32, top_k))
+                """, (query_crc32, top_k * 2))  # Увеличиваем лимит для последующего реранкинга
 
                 results = cur.fetchall()
-            logging.info(f"Найдено {len(results)} результатов поиска.")
+            logging.info(f"Найдено {len(results)} предварительных результатов поиска.")
         except Exception as e:
             logging.error(f"Ошибка при выполнении поискового запроса: {e}")
             results = []
+        
+        # Применяем реранкер
+        reranked_results = rerank_results(query, results, conn)
+
+        # Ограничиваем количество результатов после реранкинга
+        reranked_results = reranked_results[:top_k]
 
         output = ""
-        for movie_id, similarity in results:
+        for movie_id, score in reranked_results:
             # Находим фильм по ID
             movie = next((m for m in movies_data if m['id'] == movie_id), None)
             if movie:
                 output += f"<h3>{movie['name']} ({movie['year']})</h3>\n"
-                output += f"<p><strong>Жанры:</strong> {movie['genreslist']}</p>\n"
+                output += f"<p><strong>Жанры:</strong> {movie['genresList']}</p>\n"
                 output += f"<p><strong>Описание:</strong> {movie['description']}</p>\n"
-                output += f"<p><strong>Релевантность:</strong> {similarity:.4f}</p>\n"
+                output += f"<p><strong>Релевантность:</strong> {score:.4f}</p>\n"  # Используем score от реранкера
                 output += "<hr>\n"
 
         search_time = time.time() - start_time