148 lines
6.2 KiB
Python
148 lines
6.2 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Сервис переранжирования с использованием кросс-энкодеров.
|
||
Загружает модель асинхронно при первом вызове с использованием asyncio.Lock.
|
||
"""
|
||
|
||
import asyncio
|
||
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
|
||
self._load_lock = asyncio.Lock()
|
||
self._model_loaded = asyncio.Event()
|
||
self._loading = False
|
||
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
|
||
|
||
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)
|
||
|
||
async 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 []
|
||
|
||
await self._load_model_async()
|
||
|
||
# Формируем пары (запрос, фрагмент) для оценки
|
||
pairs = [(query, frag) for frag in fragments]
|
||
|
||
try:
|
||
# Выполняем предсказание синхронно в отдельном потоке
|
||
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}")
|
||
return [(frag, 0.0) for frag in fragments]
|
||
|
||
async 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 = await self.rerank(query, fragments, top_k)
|
||
return "\n\n".join([frag for frag, _ in ranked]) |