Update app.py

app.py

@@ -12,7 +12,7 @@ from urllib.parse import urlparse
 import logging
 from sklearn.preprocessing import normalize
 from concurrent.futures import ThreadPoolExecutor
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import requests
 
 # Настройка логирования
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -38,13 +38,10 @@ 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("Модель реранкера загружена успешно.")
+# Voyage AI API Key
+VOYAGE_API_KEY = os.environ.get("VOYAGE_API_KEY")
+if VOYAGE_API_KEY is None:
+    raise ValueError("VOYAGE_API_KEY environment variable not set.")
 
 # Имена таблиц
 embeddings_table = "movie_embeddings"
@@ -81,6 +78,9 @@ batch_size = 32
 # Количество потоков для параллельной обработки
 num_threads = 5
 
+# Количество потоков для параллельного реранкинга
+rerank_threads = 5 # Подберите оптимальное значение
+
 def get_db_connection():
     """Устанавливает соединение с базой данных."""
     try:
@@ -298,24 +298,84 @@ def get_movie_embeddings(conn):
         logging.error(f"Ошибка при загрузке эмбеддингов фильмов: {e}")
     return movie_embeddings
 
-def rerank_results(query, results):
-    """Переранжирует результаты поиска с помощью реранкера."""
-    logging.info(f"Начало переранжирования для запроса: '{query}'")
-    pairs = []
+def rerank_batch_voyage(query, batch):
+    """Переранжирует пакет результатов с помощью Voyage AI."""
+    url = "https://api.voyageai.com/v1/rerank"
+    headers = {
+        "Authorization": f"Bearer {VOYAGE_API_KEY}",
+        "content-type": "application/json"
+    }
+    
+    documents = []
     movie_ids = []
-    for i, (movie_id, _) in enumerate(results):
+    for movie_id, _ in batch:
         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])
+            documents.append(movie_info)
             movie_ids.append(movie_id)
-            logging.info(f"Обработка фильма для реранка {i+1}/{len(results)}: {movie['name']}")
 
-    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()
+    payload = {
+        "query": query,
+        "documents": documents,
+        "model": "rerank-2", # Можно использовать rerank-2-lite для более быстрой, но менее точной модели
+        "return_documents": False,
+        "truncation": True
+    }
+
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()  # Проверка на ошибки HTTP
+        response_json = response.json()
+        
+        reranked_results = []
+        for item in response_json['data']:
+            reranked_results.append((movie_ids[item['index']], item['relevance_score']))
+        
+        logging.info(f"Voyage AI: Успешно переранжирован батч. Задействовано токенов: {response_json['usage']['total_tokens']}")
+        return reranked_results
+    
+    except requests.exceptions.RequestException as e:
+        logging.error(f"Ошибка запроса к Voyage AI: {e}")
+        return []
+    except KeyError as e:
+        logging.error(f"Ошибка обработки ответа от Voyage AI: {e}. Полный ответ: {response_json}")
+        return []
+
+def rerank_results(query, results):
+    """Переранжирует результаты поиска с помощью Voyage AI."""
+    logging.info(f"Начало переранжирования для запроса: '{query}'")
+    reranked_results = []
+
+    with ThreadPoolExecutor(max_workers=rerank_threads) as executor:
+        futures = []
+        batch = []
+        batch_num = 0
+        for i, result in enumerate(results):
+            batch.append(result)
+            if len(batch) >= batch_size: # Отправл��ем на реранк батчами
+                logging.info(f"Отправка на переранжирование батча {batch_num+1} ({len(batch)} фильмов)")
+                future = executor.submit(rerank_batch_voyage, query, batch)
+                futures.append(future)
+                batch = []
+                batch_num += 1
+
+        # Обработка остатка
+        if batch:
+            logging.info(f"Отправка на переранжирование батча {batch_num+1} ({len(batch)} фильмов)")
+            future = executor.submit(rerank_batch_voyage, query, batch)
+            futures.append(future)
+
+        # Сбор результатов
+        for i, future in enumerate(futures):
+            try:
+                batch_result = future.result()
+                reranked_results.extend(batch_result)
+                logging.info(f"Завершен реранк батча {i+1}")
+            except Exception as e:
+                logging.error(f"Ошибка при переранжировании батча {i+1}: {e}")
 
-    reranked_results = sorted(zip(movie_ids, scores.tolist()), key=lambda x: x[1], reverse=True)
+    reranked_results = sorted(reranked_results, key=lambda x: x[1], reverse=True)
     logging.info("Переранжирование завершено.")
     return reranked_results
 
@@ -362,7 +422,7 @@ def search_movies(query, top_k=20):
                     FROM {embeddings_table} m, query_embedding
                     ORDER BY similarity DESC
                     LIMIT %s
-                """, (query_crc32, int(top_k * 2)))  # Уменьшаем лимит до * 1.1
+                """, (query_crc32, int(top_k * 1.1)))  # Уменьшаем лимит до * 1.1
 
                 results = cur.fetchall()
             logging.info(f"Найдено {len(results)} предварительных результатов поиска.")
@@ -381,7 +441,7 @@ def search_movies(query, top_k=20):
                 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>Релевантность (reranker score):</strong> {score:.4f}</p>\n"
+                output += f"<p><strong>Релевантность (Voyage AI reranker score):</strong> {score:.4f}</p>\n"
                 output += "<hr>\n"
 
         search_time = time.time() - start_time