Add sentence-transformers CrossEncoder integration

config_sentence_transformers.json

@@ -0,0 +1,13 @@
+{
+  "model_type": "CrossEncoder",
+  "__version__": {
+    "sentence_transformers": "5.6.0",
+    "transformers": "5.3.0",
+    "pytorch": "2.6.0"
+  },
+  "prompts": {
+    "retrieval": "Given a query, retrieve documents that answer the query."
+  },
+  "default_prompt_name": "retrieval",
+  "activation_fn": "torch.nn.modules.activation.Sigmoid"
+}

kalm_cross_encoder.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+from typing import Any, ClassVar
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+import torch
+from sentence_transformers.base.modules import InputModule
+from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer
+
+from .kalm_reranker_utils import (
+    DEFAULT_INSTRUCTION,
+    DEFAULT_SYSTEM_INSTRUCTION,
+    answer_token_id,
+    build_decoder_text,
+    cast_floating_parameters,
+    extract_yes_no_logits,
+    forward_reranker_model,
+    normalize_requested_dtype,
+    validate_text_pairs,
+)
+
+
+class KaLMCrossEncoderModule(InputModule):
+    """Sentence Transformers input module for KaLM encoder-decoder rerankers."""
+
+    config_file_name = "kalm_cross_encoder_config.json"
+    config_keys: ClassVar[list[str]] = [
+        "query_max_length",
+        "document_max_length",
+        "encoder_chunk_size",
+        "system_instruction",
+    ]
+    save_in_root = True
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        *,
+        query_max_length: int = 512,
+        document_max_length: int = 1024,
+        encoder_chunk_size: int | None = 4,
+        system_instruction: str = DEFAULT_SYSTEM_INSTRUCTION,
+        model_kwargs: dict[str, Any] | None = None,
+        processor_kwargs: dict[str, Any] | None = None,
+        config_kwargs: dict[str, Any] | None = None,
+        backend: str = "torch",
+    ) -> None:
+        super().__init__()
+        if backend != "torch":
+            raise ValueError(
+                "KaLMCrossEncoderModule only supports backend='torch'; "
+                f"received {backend!r}."
+            )
+        if not isinstance(model_name_or_path, str) or not model_name_or_path:
+            raise ValueError("model_name_or_path must be a non-empty string.")
+        if not isinstance(query_max_length, int) or query_max_length <= 0:
+            raise ValueError("query_max_length must be a positive integer.")
+        if not isinstance(document_max_length, int) or document_max_length <= 0:
+            raise ValueError("document_max_length must be a positive integer.")
+        if encoder_chunk_size is not None and (
+            not isinstance(encoder_chunk_size, int) or encoder_chunk_size <= 0
+        ):
+            raise ValueError("encoder_chunk_size must be a positive integer or None.")
+        if not isinstance(system_instruction, str):
+            raise TypeError("system_instruction must be a string.")
+
+        self.query_max_length = query_max_length
+        self.max_seq_length = document_max_length
+        self.encoder_chunk_size = encoder_chunk_size
+        self.system_instruction = system_instruction
+        self.backend = backend
+
+        model_kwargs = dict(model_kwargs or {})
+        processor_kwargs = dict(processor_kwargs or {})
+        config_kwargs = dict(config_kwargs or {})
+
+        num_labels = config_kwargs.pop("num_labels", 1)
+        if num_labels != 1:
+            raise ValueError(
+                "KaLM reranking produces one relevance score; num_labels must be 1."
+            )
+
+        config = AutoConfig.from_pretrained(model_name_or_path, **config_kwargs)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path, **processor_kwargs
+        )
+        if self.tokenizer.pad_token_id is None:
+            if self.tokenizer.eos_token_id is None:
+                raise ValueError(
+                    "The tokenizer must define a pad token or an EOS token."
+                )
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "right"
+        self.processor = self.tokenizer
+
+        requested_dtype = normalize_requested_dtype(
+            model_kwargs.get("dtype", model_kwargs.get("torch_dtype"))
+        )
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_name_or_path,
+            config=config,
+            **model_kwargs,
+        )
+        cast_floating_parameters(self.model, requested_dtype)
+
+        self.yes_token_id = answer_token_id(self.tokenizer, "yes")
+        self.no_token_id = answer_token_id(self.tokenizer, "no")
+
+    @property
+    def document_max_length(self) -> int:
+        return self.max_seq_length
+
+    @document_max_length.setter
+    def document_max_length(self, value: int) -> None:
+        if not isinstance(value, int) or value <= 0:
+            raise ValueError("document_max_length must be a positive integer.")
+        self.max_seq_length = value
+
+    @property
+    def encoder_chunk_size(self) -> int | None:
+        return self._encoder_chunk_size
+
+    @encoder_chunk_size.setter
+    def encoder_chunk_size(self, value: int | None) -> None:
+        if value is not None and (not isinstance(value, int) or value <= 0):
+            raise ValueError("encoder_chunk_size must be a positive integer or None.")
+        self._encoder_chunk_size = value
+
+    @property
+    def chunk_size(self) -> int | None:
+        """Alias for the encoder token mean-pooling compression rate."""
+        return self.encoder_chunk_size
+
+    @chunk_size.setter
+    def chunk_size(self, value: int | None) -> None:
+        self.encoder_chunk_size = value
+
+    def preprocess(
+        self,
+        inputs: list[Any],
+        prompt: str | None = None,
+        **kwargs: Any,
+    ) -> dict[str, torch.Tensor]:
+        pairs = validate_text_pairs(inputs)
+        if not pairs:
+            return {}
+
+        instruction = DEFAULT_INSTRUCTION if prompt is None else prompt
+        if not isinstance(instruction, str):
+            raise TypeError("prompt must be a string or None.")
+
+        encoder_texts = [f"<Document>: {document}" for _, document in pairs]
+        decoder_texts = [
+            build_decoder_text(
+                self.tokenizer,
+                query,
+                instruction,
+                self.system_instruction,
+                self.query_max_length,
+            )
+            for query, _ in pairs
+        ]
+
+        encoder_batch = self.tokenizer(
+            encoder_texts,
+            padding=True,
+            truncation=True,
+            max_length=self.document_max_length,
+            add_special_tokens=False,
+            return_tensors="pt",
+        )
+        decoder_batch = self.tokenizer(
+            decoder_texts,
+            padding=True,
+            pad_to_multiple_of=8,
+            add_special_tokens=False,
+            return_tensors="pt",
+        )
+        return {
+            "input_ids": encoder_batch["input_ids"],
+            "attention_mask": encoder_batch["attention_mask"],
+            "decoder_input_ids": decoder_batch["input_ids"],
+            "decoder_attention_mask": decoder_batch["attention_mask"],
+        }
+
+    def forward(
+        self,
+        features: dict[str, torch.Tensor | Any],
+        **kwargs: Any,
+    ) -> dict[str, torch.Tensor | Any]:
+        outputs = forward_reranker_model(
+            self.model,
+            input_ids=features["input_ids"],
+            attention_mask=features["attention_mask"],
+            decoder_input_ids=features["decoder_input_ids"],
+            decoder_attention_mask=features["decoder_attention_mask"],
+            encoder_chunk_size=self.chunk_size,
+        )
+        yes_no_logits = extract_yes_no_logits(
+            outputs.logits,
+            features["decoder_attention_mask"],
+            self.yes_token_id,
+            self.no_token_id,
+        )
+        features["scores"] = (yes_no_logits[:, 0] - yes_no_logits[:, 1]).unsqueeze(1)
+        return features
+
+    def save(
+        self,
+        output_path: str,
+        *args: Any,
+        safe_serialization: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        self.model.save_pretrained(output_path, safe_serialization=safe_serialization)
+        self.tokenizer.save_pretrained(output_path)
+        self.save_config(output_path)
+
+    @classmethod
+    def load(
+        cls,
+        model_name_or_path: str,
+        subfolder: str = "",
+        token: bool | str | None = None,
+        cache_folder: str | None = None,
+        revision: str | None = None,
+        local_files_only: bool = False,
+        trust_remote_code: bool = False,
+        model_kwargs: dict[str, Any] | None = None,
+        processor_kwargs: dict[str, Any] | None = None,
+        config_kwargs: dict[str, Any] | None = None,
+        backend: str = "torch",
+        **kwargs: Any,
+    ) -> Self:
+        module_config = cls.load_config(
+            model_name_or_path,
+            subfolder=subfolder,
+            token=token,
+            cache_folder=cache_folder,
+            revision=revision,
+            local_files_only=local_files_only,
+        )
+
+        supplied_model_kwargs = dict(model_kwargs or {})
+        supplied_config_kwargs = dict(config_kwargs or {})
+        supplied_module_kwargs = dict(kwargs)
+        chunk_size_values: list[tuple[str, int | None]] = []
+        for source_name, source in (
+            ("model_kwargs", supplied_model_kwargs),
+            ("config_kwargs", supplied_config_kwargs),
+            ("module kwargs", supplied_module_kwargs),
+        ):
+            for key in ("chunk_size", "encoder_chunk_size"):
+                if key in source:
+                    chunk_size_values.append((f"{source_name}.{key}", source.pop(key)))
+        if chunk_size_values:
+            first_name, first_value = chunk_size_values[0]
+            for current_name, current_value in chunk_size_values[1:]:
+                if current_value != first_value:
+                    raise ValueError(
+                        "Conflicting encoder chunk sizes: "
+                        f"{first_name}={first_value!r}, "
+                        f"{current_name}={current_value!r}."
+                    )
+            module_config["encoder_chunk_size"] = first_value
+
+        hub_kwargs = {
+            "subfolder": subfolder,
+            "token": token,
+            "cache_dir": cache_folder,
+            "revision": revision,
+            "local_files_only": local_files_only,
+            "trust_remote_code": trust_remote_code,
+        }
+        effective_model_kwargs = {**hub_kwargs, **supplied_model_kwargs}
+        effective_processor_kwargs = {**hub_kwargs, **(processor_kwargs or {})}
+        effective_config_kwargs = {**hub_kwargs, **supplied_config_kwargs}
+
+        if "model_max_length" in effective_processor_kwargs:
+            module_config["document_max_length"] = effective_processor_kwargs[
+                "model_max_length"
+            ]
+
+        return cls(
+            model_name_or_path,
+            model_kwargs=effective_model_kwargs,
+            processor_kwargs=effective_processor_kwargs,
+            config_kwargs=effective_config_kwargs,
+            backend=backend,
+            **module_config,
+        )
+
+
+__all__ = ["KaLMCrossEncoderModule"]

kalm_cross_encoder_config.json

@@ -0,0 +1,6 @@
+{
+  "query_max_length": 512,
+  "document_max_length": 1024,
+  "encoder_chunk_size": 4,
+  "system_instruction": "Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\"."
+}

modules.json

@@ -0,0 +1,8 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "kalm_cross_encoder.KaLMCrossEncoderModule"
+  }
+]