From bfc9d2a901bfc3e1aa346a0a34956204e9d278e7 Mon Sep 17 00:00:00 2001
From: Markov Andrey <anmarkov@pptrf.ru>
Date: Tue, 30 Jun 2026 21:26:27 +0000
Subject: [PATCH] =?UTF-8?q?=D0=A0=D0=B5=D0=B4=D0=B0=D0=BA=D1=82=D0=B8?=
 =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B0=D1=82=D1=8C=20reranker=5Fservice.py?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 rag/services/reranker_service.py | 73 ++++++++++++++++++--------------
 1 file changed, 42 insertions(+), 31 deletions(-)

diff --git a/rag/services/reranker_service.py b/rag/services/reranker_service.py
index e5a8e27..702937d 100644
--- a/rag/services/reranker_service.py
+++ b/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])
\ No newline at end of file