diff --git a/core/services/reranker_service.py b/core/services/reranker_service.py new file mode 100644 index 0000000..e5a8e27 --- /dev/null +++ b/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]) \ No newline at end of file