Add new file

core/services/reranker_service.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+"""
+Сервис переранжирования с использованием кросс-энкодеров.
+Загружает модель один раз при первом вызове и использует её для ранжирования фрагментов.
+"""
+
+import logging
+from typing import List, Tuple, Optional
+
+logger = logging.getLogger(__name__)
+
+# Импортируем sentence-transformers с обработкой ошибок
+try:
+    from sentence_transformers import CrossEncoder
+    HAS_CROSS_ENCODER = True
+except ImportError:
+    HAS_CROSS_ENCODER = False
+    logger.warning("sentence-transformers не установлен. Используйте LLM-переранжирование.")
+
+
+class RerankerService:
+    """
+    Сервис переранжирования на основе кросс-энкодеров.
+    Загружает модель при инициализации или при первом вызове.
+    """
+
+    # Рекомендуемая модель для русского языка
+    DEFAULT_MODEL = "intfloat/multilingual-e5-reranker"
+
+    def __init__(self, model_name: Optional[str] = None, device: Optional[str] = None):
+        """
+        Инициализация сервиса.
+
+        Аргументы:
+            model_name: имя модели (если None, используется DEFAULT_MODEL).
+            device: устройство для вычислений ('cpu' или 'cuda'). Если None, автоопределение.
+        """
+        self.model_name = model_name or self.DEFAULT_MODEL
+        self.device = device or ("cuda" if self._cuda_available() else "cpu")
+        self._model = None
+        logger.info(f"RerankerService инициализирован с моделью {self.model_name} (device={self.device})")
+
+    @staticmethod
+    def _cuda_available() -> bool:
+        """Проверяет доступность CUDA (для GPU)."""
+        try:
+            import torch
+            return torch.cuda.is_available()
+        except ImportError:
+            return False
+
+    def _load_model(self) -> None:
+        """Загружает модель кросс-энкодера (ленивая загрузка)."""
+        if not HAS_CROSS_ENCODER:
+            raise RuntimeError("sentence-transformers не установлен. Установите: pip install sentence-transformers")
+
+        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(
+        self,
+        query: str,
+        fragments: List[str],
+        top_k: Optional[int] = None
+    ) -> List[Tuple[str, float]]:
+        """
+        Ранжирует фрагменты по релевантности запросу.
+
+        Аргументы:
+            query: текст запроса.
+            fragments: список текстовых фрагментов.
+            top_k: максимальное количество возвращаемых фрагментов (если None – все).
+
+        Возвращает:
+            List[Tuple[str, float]]: список кортежей (фрагмент, оценка_релевантности),
+                                      отсортированный по убыванию оценки.
+        """
+        if not fragments:
+            return []
+
+        self._load_model()
+
+        # Формируем пары (запрос, фрагмент) для оценки
+        pairs = [(query, frag) for frag in fragments]
+
+        # Получаем оценки релевантности (числа от 0 до 1)
+        try:
+            scores = self._model.predict(pairs)
+        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(
+        self,
+        query: str,
+        context: str,
+        top_k: Optional[int] = None
+    ) -> str:
+        """
+        Удобная обёртка: принимает строку контекста с фрагментами, разделёнными '\n\n',
+        возвращает отранжированную строку.
+
+        Аргументы:
+            query: текст запроса.
+            context: исходный контекст (фрагменты разделены '\n\n').
+            top_k: максимальное количество возвращаемых фрагментов.
+
+        Возвращает:
+            str: отранжированный контекст (фрагменты разделены '\n\n').
+        """
+        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)
+
+        # Собираем обратно в строку (только текст, без оценок)
+        return "\n\n".join([frag for frag, _ in ranked])