aramesh-backend/chatbot/reranker.py

"""مدل ری‌رنکر BGE — با لود تنبل و امتیازدهی دسته‌ای."""

import logging
import os

from functools import lru_cache

from .config import settings

log = logging.getLogger(__name__)


class _CrossEncoder:
    """غلاف سبک روی AutoModelForSequenceClassification — API شبیه rank_bm25."""
    def __init__(self) -> None:
        import torch
        from transformers import AutoModelForSequenceClassification, AutoTokenizer

        os.makedirs(settings.hf_cache_dir, exist_ok=True)
        self._torch = torch
        self.tokenizer = AutoTokenizer.from_pretrained(settings.reranker_model, cache_dir=settings.hf_cache_dir)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            settings.reranker_model, cache_dir=settings.hf_cache_dir
        )
        device = settings.reranker_device if settings.reranker_device == "cuda" and torch.cuda.is_available() else "cpu"
        self.model.to(device)
        self.model.eval()
        self.device = device
        log.info("Reranker loaded on %s", device)

    def score(self, pairs: list[list[str]], batch_size: int = 16) -> list[float]:
        torch = self._torch
        out: list[float] = []
        with torch.no_grad():
            for i in range(0, len(pairs), batch_size):
                inputs = self.tokenizer(
                    pairs[i:i + batch_size],
                    padding=True,
                    truncation=True,
                    max_length=512,
                    return_tensors="pt",
                ).to(self.device)
                logits = self.model(**inputs).logits.squeeze(-1)
                out.extend(logits.cpu().tolist())
        return out


@lru_cache(maxsize=1)
def get_reranker() -> _CrossEncoder:
    return _CrossEncoder()