Update app.py

app.py

@@ -12,7 +12,6 @@ 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')
@@ -32,20 +31,12 @@ 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"
@@ -298,27 +289,6 @@ 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
@@ -362,29 +332,23 @@ def search_movies(query, top_k=20):
                     FROM {embeddings_table} m, query_embedding
                     ORDER BY similarity DESC
                     LIMIT %s
-                """, (query_crc32, top_k * 2))  # Увеличиваем лимит для последующего реранкинга
+                """, (query_crc32, top_k))
 
                 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, score in reranked_results:
+        for movie_id, similarity in 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['description']}</p>\n"
-                output += f"<p><strong>Релевантность:</strong> {score:.4f}</p>\n"  # Используем score от реранкера
+                output += f"<p><strong>Релевантность:</strong> {similarity:.4f}</p>\n"
                 output += "<hr>\n"
 
         search_time = time.time() - start_time