Редактировать reranker_service.py

rag/services/reranker_service.py

@@ -1,9 +1,10 @@
 # -*- coding: utf-8 -*-
 """
 Сервис переранжирования с использованием кросс-энкодеров.
-Загружает модель один раз при первом вызове и использует её для ранжирования фрагментов.
+Загружает модель асинхронно при первом вызове с использованием asyncio.Lock.
 """
 
+import asyncio
 import logging
 from typing import List, Tuple, Optional
 
@@ -21,7 +22,7 @@ except ImportError:
 class RerankerService:
     """
     Сервис переранжирования на основе кросс-энкодеров.
-    Загружает модель при инициализации или при первом вызове.
+    Загружает модель асинхронно при инициализации или при первом вызове.
     """
 
     # Рекомендуемая модель для русского языка
@@ -38,6 +39,9 @@ class RerankerService:
         self.model_name = model_name or self.DEFAULT_MODEL
         self.device = device or ("cuda" if self._cuda_available() else "cpu")
         self._model = None
+        self._load_lock = asyncio.Lock()
+        self._model_loaded = asyncio.Event()
+        self._loading = False
         logger.info(f"RerankerService инициализирован с моделью {self.model_name} (device={self.device})")
 
     @staticmethod
@@ -49,17 +53,34 @@ class RerankerService:
         except ImportError:
             return False
 
-    def _load_model(self) -> None:
-        """Загружает модель кросс-энкодера (ленивая загрузка)."""
+    async def _load_model_async(self) -> None:
+        """Загружает модель кросс-энкодера асинхронно (ленивая загрузка)."""
+        async with self._load_lock:
+            if self._model_loaded.is_set():
+                return
+            if self._loading:
+                await self._model_loaded.wait()
+                return
+            self._loading = True
+            try:
+                loop = asyncio.get_running_loop()
+                self._model = await loop.run_in_executor(None, self._load_model_sync)
+                self._model_loaded.set()
+                logger.info(f"Модель {self.model_name} загружена")
+            except Exception as e:
+                logger.error(f"Ошибка загрузки модели: {e}")
+                self._loading = False
+                raise
+            finally:
+                self._loading = False
+
+    def _load_model_sync(self):
+        """Синхронная загрузка модели (выполняется в отдельном потоке)."""
         if not HAS_CROSS_ENCODER:
             raise RuntimeError("sentence-transformers не установлен. Установите: pip install sentence-transformers")
+        return CrossEncoder(self.model_name, device=self.device)
 
-        if self._model is None:
-            logger.info(f"Загрузка модели {self.model_name} на устройство {self.device}...")
-            self._model = CrossEncoder(self.model_name, device=self.device)
-            logger.info(f"Модель {self.model_name} загружена")
-
-    def rerank(
+    async def rerank(
         self,
         query: str,
         fragments: List[str],
@@ -80,31 +101,25 @@ class RerankerService:
         if not fragments:
             return []
 
-        self._load_model()
+        await self._load_model_async()
 
         # Формируем пары (запрос, фрагмент) для оценки
         pairs = [(query, frag) for frag in fragments]
 
-        # Получаем оценки релевантности (числа от 0 до 1)
         try:
-            scores = self._model.predict(pairs)
+            # Выполняем предсказание синхронно в отдельном потоке
+            loop = asyncio.get_running_loop()
+            scores = await loop.run_in_executor(None, self._model.predict, pairs)
+            results = list(zip(fragments, scores))
+            results.sort(key=lambda x: x[1], reverse=True)
+            if top_k:
+                results = results[:top_k]
+            return results
         except Exception as e:
             logger.error(f"Ошибка при оценке кросс-энкодером: {e}")
-            # В случае ошибки возвращаем фрагменты в исходном порядке с оценкой 0.0
             return [(frag, 0.0) for frag in fragments]
 
-        # Собираем результаты и сортируем по убыванию оценки
-        results = list(zip(fragments, scores))
-        results.sort(key=lambda x: x[1], reverse=True)
-
-        # Ограничиваем количество, если указано
-        if top_k is not None:
-            results = results[:top_k]
-
-        logger.debug(f"Переранжирование: из {len(fragments)} фрагментов оставлено {len(results)}")
-        return results
-
-    def rerank_context(
+    async def rerank_context(
         self,
         query: str,
         context: str,
@@ -125,13 +140,9 @@ class RerankerService:
         if not context:
             return ""
 
-        # Разбиваем контекст на фрагменты
         fragments = [f.strip() for f in context.split("\n\n") if f.strip()]
         if not fragments:
             return ""
 
-        # Ранжируем
-        ranked = self.rerank(query, fragments, top_k)
-
-        # Собираем обратно в строку (только текст, без оценок)
+        ranked = await self.rerank(query, fragments, top_k)
         return "\n\n".join([frag for frag, _ in ranked])