137 lines
6.0 KiB
Python
137 lines
6.0 KiB
Python
# -*- 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]) |