KaLM-Embedding/KaLM-Reranker-V1-Nano-R2

1from __future__ import annotations2 3from typing import Any, Dict, List, Optional, Sequence, Tuple, Union4 5import numpy as np6import torch7from transformers import AutoModelForSeq2SeqLM, AutoTokenizer8 9try:10    from .kalm_reranker_utils import (11        DEFAULT_INSTRUCTION,12        DEFAULT_SYSTEM_INSTRUCTION,13        answer_token_id,14        build_decoder_text,15        cast_floating_parameters,16        extract_yes_no_logits,17        forward_reranker_model,18        get_encoder,19        pool_encoder_chunks,20        validate_text_pairs,21    )22except ImportError:  # Support ``from kalm_reranker import KaLMReranker``.23    from kalm_reranker_utils import (24        DEFAULT_INSTRUCTION,25        DEFAULT_SYSTEM_INSTRUCTION,26        answer_token_id,27        build_decoder_text,28        cast_floating_parameters,29        extract_yes_no_logits,30        forward_reranker_model,31        get_encoder,32        pool_encoder_chunks,33        validate_text_pairs,34    )35 36 37class KaLMReranker:38    """Score query-document relevance with a KaLM encoder-decoder reranker.39 40    The returned score is ``P(yes)`` after applying a two-class softmax to the41    model's ``yes`` and ``no`` logits.42    """43 44    def __init__(45        self,46        model_name_or_path: str,47        *,48        device: Optional[Union[str, torch.device]] = None,49        dtype: Optional[Union[str, torch.dtype]] = None,50        batch_size: int = 32,51        query_max_length: int = 512,52        max_length: int = 1024,53        chunk_size: Optional[int] = 4,54        instruction: str = DEFAULT_INSTRUCTION,55        system_instruction: str = DEFAULT_SYSTEM_INSTRUCTION,56        **model_kwargs: Any,57    ) -> None:58        if not isinstance(model_name_or_path, str) or not model_name_or_path:59            raise ValueError("model_name_or_path must be a non-empty string.")60        if batch_size <= 0:61            raise ValueError("batch_size must be positive.")62        if query_max_length <= 0 or max_length <= 0:63            raise ValueError("query_max_length and max_length must be positive.")64        if chunk_size is not None and chunk_size <= 0:65            raise ValueError("chunk_size must be positive or None.")66        if not isinstance(instruction, str) or not isinstance(system_instruction, str):67            raise TypeError("instruction and system_instruction must be strings.")68 69        self.device = self._resolve_device(device)70        self.dtype = self._resolve_dtype(dtype, self.device)71        self.batch_size = batch_size72        self.query_max_length = query_max_length73        self.max_length = max_length74        self.chunk_size = chunk_size75        self.instruction = instruction76        self.system_instruction = system_instruction77 78        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)79        if self.tokenizer.pad_token_id is None:80            if self.tokenizer.eos_token_id is None:81                raise ValueError(82                    "The tokenizer must define a pad token or an EOS token."83                )84            self.tokenizer.pad_token = self.tokenizer.eos_token85        # Final decoder-token indexing assumes right padding, matching training.86        self.tokenizer.padding_side = "right"87 88        self.model = AutoModelForSeq2SeqLM.from_pretrained(89            model_name_or_path,90            dtype=self.dtype,91            **model_kwargs,92        )93        cast_floating_parameters(self.model, self.dtype)94        self.model.to(device=self.device)95        self.model.eval()96 97        self.yes_token_id = self._answer_token_id("yes")98        self.no_token_id = self._answer_token_id("no")99 100    @staticmethod101    def _resolve_device(device: Optional[Union[str, torch.device]]) -> torch.device:102        if device is None:103            device = "cuda" if torch.cuda.is_available() else "cpu"104        resolved = torch.device(device)105        if resolved.type == "cuda" and not torch.cuda.is_available():106            raise RuntimeError("CUDA was requested, but no CUDA device is available.")107        return resolved108 109    @staticmethod110    def _resolve_dtype(111        dtype: Optional[Union[str, torch.dtype]], device: torch.device112    ) -> torch.dtype:113        if dtype is None:114            return torch.bfloat16 if device.type == "cuda" else torch.float32115        if isinstance(dtype, torch.dtype):116            return dtype117        if not isinstance(dtype, str):118            raise TypeError(119                "dtype must be a torch.dtype or a string such as 'bfloat16'."120            )121        normalized = dtype.lower().removeprefix("torch.")122        supported = {123            "bfloat16": torch.bfloat16,124            "bf16": torch.bfloat16,125            "float16": torch.float16,126            "fp16": torch.float16,127            "float32": torch.float32,128            "fp32": torch.float32,129        }130        if normalized not in supported:131            raise ValueError(f"Unsupported dtype: {dtype!r}.")132        return supported[normalized]133 134    def _answer_token_id(self, answer: str) -> int:135        return answer_token_id(self.tokenizer, answer)136 137    def _get_encoder(self):138        return get_encoder(self.model)139 140    @staticmethod141    def _pool_encoder_chunks(142        hidden_states: torch.Tensor,143        attention_mask: torch.Tensor,144        chunk_size: int,145    ) -> Tuple[torch.Tensor, torch.Tensor]:146        return pool_encoder_chunks(hidden_states, attention_mask, chunk_size)147 148    def _decoder_text(self, query: str, instruction: str) -> str:149        return build_decoder_text(150            self.tokenizer,151            query,152            instruction,153            self.system_instruction,154            self.query_max_length,155        )156 157    @staticmethod158    def _validate_pairs(159        pairs: Sequence[Tuple[str, str]],160    ) -> List[Tuple[str, str]]:161        return validate_text_pairs(pairs)162 163    @torch.inference_mode()164    def _predict_batch(165        self, pairs: Sequence[Tuple[str, str]], instruction: str166    ) -> List[float]:167        encoder_texts = [f"<Document>: {document}" for _, document in pairs]168        decoder_texts = [self._decoder_text(query, instruction) for query, _ in pairs]169 170        encoder_batch = self.tokenizer(171            encoder_texts,172            padding=True,173            truncation=True,174            max_length=self.max_length,175            add_special_tokens=False,176            return_tensors="pt",177        ).to(self.device)178        decoder_batch = self.tokenizer(179            decoder_texts,180            padding=True,181            pad_to_multiple_of=8,182            add_special_tokens=False,183            return_tensors="pt",184        ).to(self.device)185 186        outputs = forward_reranker_model(187            self.model,188            input_ids=encoder_batch["input_ids"],189            attention_mask=encoder_batch["attention_mask"],190            decoder_input_ids=decoder_batch["input_ids"],191            decoder_attention_mask=decoder_batch["attention_mask"],192            encoder_chunk_size=self.chunk_size,193        )194        yes_no_logits = extract_yes_no_logits(195            outputs.logits,196            decoder_batch["attention_mask"],197            self.yes_token_id,198            self.no_token_id,199        )200        return torch.softmax(yes_no_logits, dim=-1)[:, 0].cpu().tolist()201 202    def predict(203        self,204        pairs: Sequence[Tuple[str, str]],205        *,206        instruction: Optional[str] = None,207        batch_size: Optional[int] = None,208    ) -> List[float]:209        """Return ``P(yes)`` scores in the same order as ``pairs``."""210        validated_pairs = self._validate_pairs(pairs)211        if not validated_pairs:212            return []213        effective_instruction = self.instruction if instruction is None else instruction214        if not isinstance(effective_instruction, str):215            raise TypeError("instruction must be a string or None.")216        effective_batch_size = self.batch_size if batch_size is None else batch_size217        if not isinstance(effective_batch_size, int) or effective_batch_size <= 0:218            raise ValueError("batch_size must be a positive integer.")219 220        length_sorted_indices = np.argsort(221            [-(len(query) + len(document)) for query, document in validated_pairs]222        )223        sorted_pairs = [validated_pairs[index] for index in length_sorted_indices]224 225        tested_batch_size = effective_batch_size226        first_batch_scores: Optional[List[float]] = None227        while tested_batch_size > 1:228            try:229                first_batch_scores = self._predict_batch(230                    sorted_pairs[: min(len(sorted_pairs), tested_batch_size)],231                    effective_instruction,232                )233                break234            except torch.cuda.OutOfMemoryError:235                if torch.cuda.is_available():236                    torch.cuda.empty_cache()237                tested_batch_size = max(1, tested_batch_size * 3 // 4)238 239        if first_batch_scores is None:240            sorted_scores: List[float] = []241            loop_start = 0242        else:243            sorted_scores = list(first_batch_scores)244            loop_start = tested_batch_size245        try:246            for start in range(loop_start, len(sorted_pairs), tested_batch_size):247                sorted_scores.extend(248                    self._predict_batch(249                        sorted_pairs[start : start + tested_batch_size],250                        effective_instruction,251                    )252                )253        except torch.cuda.OutOfMemoryError as error:254            if torch.cuda.is_available():255                torch.cuda.empty_cache()256            raise RuntimeError(257                "CUDA ran out of memory during reranking. Retry with a smaller "258                "batch_size or shorter max_length."259            ) from error260        inverse_indices = np.argsort(length_sorted_indices)261        return [sorted_scores[index] for index in inverse_indices]262 263    def rank(264        self,265        query: str,266        documents: Sequence[str],267        *,268        instruction: Optional[str] = None,269        top_k: Optional[int] = None,270        batch_size: Optional[int] = None,271    ) -> List[Dict[str, Union[int, float]]]:272        """Rank documents and return ``corpus_id``/``score`` dictionaries."""273        if not isinstance(query, str):274            raise TypeError("query must be a string.")275        if isinstance(documents, (str, bytes)) or not isinstance(documents, Sequence):276            raise TypeError("documents must be a sequence of strings.")277        if any(not isinstance(document, str) for document in documents):278            raise TypeError("every document must be a string.")279        if top_k is not None and (not isinstance(top_k, int) or top_k < 0):280            raise ValueError("top_k must be a non-negative integer or None.")281 282        scores = self.predict(283            [(query, document) for document in documents],284            instruction=instruction,285            batch_size=batch_size,286        )287        rankings: List[Dict[str, Union[int, float]]] = [288            {"corpus_id": corpus_id, "score": score}289            for corpus_id, score in enumerate(scores)290        ]291        rankings.sort(key=lambda item: item["score"], reverse=True)292        return rankings if top_k is None else rankings[:top_k]293 294 295__all__ = ["KaLMReranker"]296