Add sentence-transformers CrossEncoder support

# KaLM-Reranker-V1-Nano CrossEncoder validation

Validation date: 2026-07-29.

## Immutable checkpoint baseline

- `config.json`: `75926d34f90b64a3119e62535211bcd2305e8ea0ae4339773cb20dc84c8991d4`
- `model.safetensors`: `98078ef7b21ec1ee82afa1b5cbdb3923d9548d5f3bed78d290b605c65cc3f49d`
- `tokenizer.json`: `f5b325224482ec441ec5fbe2a5ac08c3758e0f9605f6e54368e31f736fcfb01d`

## Pre-adaptation CPU FP32 baseline

Reference implementation, batch size 2, query length 512, document length
1024, encoder chunk size 4:

```text
[0.9998076558113098,
0.000004796196662937291,
0.2090071737766266,
0.021608587354421616,
0.00004020768392365426]
```

## Validation matrix

| Target | Device/dtype | Status | Max absolute error | Notes |
|---|---|---:|---:|---|
| sentence-transformers 5.6.0, chunk size 4 | CPU FP32 | passed | `1.4901161193847656e-08` | `rtol=1e-5`, `atol=1e-6`; ranking `[0, 1]` |
| local 5.7.0.dev0/main, chunk size 4 | CPU FP32 | passed | `1.4901161193847656e-08` | Imported from the local source tree; ranking `[0, 1]` |
| sentence-transformers 5.6.0, pooling disabled | CPU FP32 | passed | `1.862645149230957e-09` | Reference and CrossEncoder both used `encoder_chunk_size=None` |
| sentence-transformers 5.6.0, chunk size 4 | CUDA BF16 | passed | `3.637978807091713e-12` | H100 MIG 3g.40gb; `rtol=1e-3`, `atol=2e-3` |
| local 5.7.0.dev0/main, chunk size 4 | CUDA BF16 | passed | `3.637978807091713e-12` | Imported from the local source tree; ranking `[0, 1]` |
| sentence-transformers 5.6.0, pooling disabled | CUDA BF16 | passed | `1.862645149230957e-09` | Reference and CrossEncoder both used `chunk_size=None` |
| Hub PR revision, fresh cache | CPU/CUDA | blocked | — | Hub pre-upload returned 403; no PR was created |

CUDA error 803 was resolved by preferring the host driver libraries in
`LD_LIBRARY_PATH` and setting the NVIDIA compute/utility capabilities before
starting Python.

## Environment

- Python 3.10.20
- PyTorch 2.6.0+cu126
- Transformers 5.3.0
- Stable environment: sentence-transformers 5.6.0
- Main environment: sentence-transformers 5.7.0.dev0 imported from the local
`sentence-transformers` source tree
- GPU: NVIDIA H100 80GB HBM3 MIG 3g.40gb; CUDA BF16 is supported
- CUDA recovery environment:

```bash
export NVIDIA_DRIVER_CAPABILITIES="${NVIDIA_DRIVER_CAPABILITIES:-compute,utility}"
export LD_LIBRARY_PATH="/lib/x86_64-linux-gnu:/usr/lib/x86_64-linux-gnu${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
export CUDA_VISIBLE_DEVICES=0
```

## Interface and failure-path coverage

| Check | Status | Evidence |
|---|---:|---|
| `predict()` NumPy and Tensor output; batch sizes 1, 2, and 3; incomplete final batch; input-order restoration | passed | Five mixed-length English/Chinese/empty-input pairs match the reference at the same batch size |
| Individual `(query, document)` pair | passed | Scalar result matches the corresponding batch result |
| `rank()` default/custom instruction, `prompt_name`, `top_k`, and `return_documents` | passed | Ordering and returned document IDs/text verified |
| Empty list, empty query, empty document, malformed pair | passed | Empty shape `(0,)`; malformed input raises a typed error |
| Query/document truncation | passed | 700-query-token/1400-document-token case yields a 1024-token encoder input and an 8-aligned decoder input |
| Configurable encoder compression | passed | `model_kwargs={"chunk_size": 4}` is applied at load time; each masked group of four states is mean-pooled; `None` disables pooling |
| Default Sigmoid and Identity margin | passed | `sigmoid(yes_logit - no_logit)` matches the default score within `1e-7` |
| Single-token answers | passed | Nano tokenizer resolves `yes=4443`, `no=1904`; multi-token answer is rejected |
| Per-sample decoder endpoint and non-finite logits | passed | Unit tests cover different valid lengths and explicit NaN rejection |
| Unsupported backend | passed | `backend="onnx"` is rejected before model loading |
| `save_pretrained()` and local reload | passed | Saved custom code/config/tokenizer/weights reload and reproduce scores |

Lightweight pytest results are `7 passed, 2 skipped` in both the stable 5.6.0
environment and the local 5.7.0.dev0/main environment. Checkpoint-backed
validation was run sequentially to keep peak host memory bounded. The complete
stable-5.6.0 CUDA BF16 integration suite finished with `9 passed` in 46.97
seconds.

BF16 scores can vary slightly when the batch size changes because the padded
tensor shape changes. Batch-size coverage therefore compares each CrossEncoder
result with the reference implementation using the same batch size; all such
comparisons pass the BF16 tolerance.

## Integrity and repository isolation

- The three immutable hashes above are unchanged after all tests.
- The local sentence-transformers Git worktree remains clean; no adapter change
was made to sentence-transformers source.
- The Hub PR manifest contains only six adapter/configuration files. It excludes
the repository `README.md`, `model.safetensors`, tokenizer artifacts, and
Transformers `config.json`. The README usage section is provided separately
for manual review and editing.

## Remaining external validations

1. Grant account `cosyy` effective contributor/write access to the target
repository, or provide a token from an account that already has that access,
then run
`python create_hub_pr.py --token-file /tmp/hf_token --create-pr`.
2. Run `verify_hub_revision.py --revision refs/pr/<PR number>` for CPU FP32 and
CUDA BF16. The script uses blank model and remote-code caches and asserts
that the custom module came from that cache rather than local `PYTHONPATH`.

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\"."
+}

kalm_reranker.py

@@ -4,16 +4,34 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
-import torch.nn.functional as F
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from transformers.modeling_outputs import BaseModelOutput
 
-
-DEFAULT_INSTRUCTION = "Given a query, retrieve documents that answer the query."
-DEFAULT_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".'
-)
+try:
+    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,
+        get_encoder,
+        pool_encoder_chunks,
+        validate_text_pairs,
+    )
+except ImportError:  # Support ``from kalm_reranker import KaLMReranker``.
+    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,
+        get_encoder,
+        pool_encoder_chunks,
+        validate_text_pairs,
+    )
 
 
 class KaLMReranker:
@@ -60,8 +78,11 @@ class KaLMReranker:
         self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
         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.")
+                raise ValueError(
+                    "The tokenizer must define a pad token or an EOS token."
+                )
             self.tokenizer.pad_token = self.tokenizer.eos_token
+        # Final decoder-token indexing assumes right padding, matching training.
         self.tokenizer.padding_side = "right"
 
         self.model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -69,9 +90,7 @@ class KaLMReranker:
             dtype=self.dtype,
             **model_kwargs,
         )
-        for parameter in self.model.parameters():
-            if parameter.is_floating_point() and parameter.dtype != self.dtype:
-                parameter.data = parameter.data.to(dtype=self.dtype)
+        cast_floating_parameters(self.model, self.dtype)
         self.model.to(device=self.device)
         self.model.eval()
 
@@ -96,7 +115,9 @@ class KaLMReranker:
         if isinstance(dtype, torch.dtype):
             return dtype
         if not isinstance(dtype, str):
-            raise TypeError("dtype must be a torch.dtype or a string such as 'bfloat16'.")
+            raise TypeError(
+                "dtype must be a torch.dtype or a string such as 'bfloat16'."
+            )
         normalized = dtype.lower().removeprefix("torch.")
         supported = {
             "bfloat16": torch.bfloat16,
@@ -111,17 +132,10 @@ class KaLMReranker:
         return supported[normalized]
 
     def _answer_token_id(self, answer: str) -> int:
-        token_ids = self.tokenizer(answer, add_special_tokens=False)["input_ids"]
-        if not token_ids:
-            raise ValueError(f"Failed to tokenize the answer {answer!r}.")
-        return token_ids[-1]
+        return answer_token_id(self.tokenizer, answer)
 
     def _get_encoder(self):
-        if hasattr(self.model, "get_encoder"):
-            return self.model.get_encoder()
-        if hasattr(self.model, "encoder"):
-            return self.model.encoder
-        raise AttributeError(f"Cannot find the encoder on {type(self.model).__name__}.")
+        return get_encoder(self.model)
 
     @staticmethod
     def _pool_encoder_chunks(
@@ -129,64 +143,22 @@ class KaLMReranker:
         attention_mask: torch.Tensor,
         chunk_size: int,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        batch_size, sequence_length, hidden_size = hidden_states.shape
-        num_chunks = (sequence_length + chunk_size - 1) // chunk_size
-        padded_length = num_chunks * chunk_size
-        pad_length = padded_length - sequence_length
-
-        if pad_length:
-            hidden_states = F.pad(hidden_states, (0, 0, 0, pad_length))
-            attention_mask = F.pad(attention_mask, (0, pad_length))
-
-        hidden_states = hidden_states.view(
-            batch_size, num_chunks, chunk_size, hidden_size
-        )
-        chunk_mask = attention_mask.view(batch_size, num_chunks, chunk_size)
-        expanded_mask = chunk_mask.unsqueeze(-1).to(hidden_states.dtype)
-        pooled_hidden = (hidden_states * expanded_mask).sum(dim=2)
-        pooled_hidden = pooled_hidden / chunk_mask.sum(dim=2).clamp(min=1).unsqueeze(-1)
-        pooled_mask = (chunk_mask.sum(dim=2) > 0).to(attention_mask.dtype)
-        return pooled_hidden, pooled_mask
+        return pool_encoder_chunks(hidden_states, attention_mask, chunk_size)
 
     def _decoder_text(self, query: str, instruction: str) -> str:
-        query_ids = self.tokenizer(
+        return build_decoder_text(
+            self.tokenizer,
             query,
-            add_special_tokens=False,
-            truncation=True,
-            max_length=self.query_max_length,
-        )["input_ids"]
-        truncated_query = self.tokenizer.decode(
-            query_ids,
-            skip_special_tokens=False,
-            clean_up_tokenization_spaces=False,
-        )
-        return (
-            "<bos><start_of_turn>user\n"
-            f"{self.system_instruction}\n\n"
-            f"<Instruct>: {instruction}\n"
-            f"<Query>: {truncated_query}<end_of_turn>\n"
-            "<start_of_turn>model\n\n\n\n"
+            instruction,
+            self.system_instruction,
+            self.query_max_length,
         )
 
     @staticmethod
     def _validate_pairs(
         pairs: Sequence[Tuple[str, str]],
     ) -> List[Tuple[str, str]]:
-        if isinstance(pairs, (str, bytes)) or not isinstance(pairs, Sequence):
-            raise TypeError("pairs must be a sequence of (query, document) pairs.")
-        validated: List[Tuple[str, str]] = []
-        for index, pair in enumerate(pairs):
-            if (
-                isinstance(pair, (str, bytes))
-                or not isinstance(pair, Sequence)
-                or len(pair) != 2
-            ):
-                raise ValueError(f"pairs[{index}] must contain exactly two strings.")
-            query, document = pair
-            if not isinstance(query, str) or not isinstance(document, str):
-                raise TypeError(f"pairs[{index}] must contain exactly two strings.")
-            validated.append((query, document))
-        return validated
+        return validate_text_pairs(pairs)
 
     @torch.inference_mode()
     def _predict_batch(
@@ -211,49 +183,20 @@ class KaLMReranker:
             return_tensors="pt",
         ).to(self.device)
 
-        if self.chunk_size is None:
-            outputs = self.model(
-                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"],
-                return_dict=True,
-            )
-        else:
-            encoder_outputs = self._get_encoder()(
-                input_ids=encoder_batch["input_ids"],
-                attention_mask=encoder_batch["attention_mask"],
-                return_dict=True,
-            )
-            pooled_hidden, pooled_mask = self._pool_encoder_chunks(
-                encoder_outputs.last_hidden_state,
-                encoder_batch["attention_mask"],
-                self.chunk_size,
-            )
-            outputs = self.model(
-                encoder_outputs=BaseModelOutput(last_hidden_state=pooled_hidden),
-                attention_mask=pooled_mask,
-                decoder_input_ids=decoder_batch["input_ids"],
-                decoder_attention_mask=decoder_batch["attention_mask"],
-                return_dict=True,
-            )
-
-        sequence_lengths = decoder_batch["attention_mask"].sum(dim=1) - 1
-        batch_indices = torch.arange(outputs.logits.shape[0], device=self.device)
-        last_logits = outputs.logits[batch_indices, sequence_lengths]
-        yes_no_logits = torch.stack(
-            (
-                last_logits[:, self.yes_token_id],
-                last_logits[:, self.no_token_id],
-            ),
-            dim=-1,
-        ).float()
-        if not torch.isfinite(yes_no_logits).all():
-            bad_count = (~torch.isfinite(yes_no_logits).all(dim=-1)).sum().item()
-            raise RuntimeError(
-                f"The model produced non-finite yes/no logits for {bad_count} input(s). "
-                "Use bfloat16 or float32 instead of float16."
-            )
+        outputs = forward_reranker_model(
+            self.model,
+            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"],
+            encoder_chunk_size=self.chunk_size,
+        )
+        yes_no_logits = extract_yes_no_logits(
+            outputs.logits,
+            decoder_batch["attention_mask"],
+            self.yes_token_id,
+            self.no_token_id,
+        )
         return torch.softmax(yes_no_logits, dim=-1)[:, 0].cpu().tolist()
 
     def predict(
@@ -293,11 +236,6 @@ class KaLMReranker:
                     torch.cuda.empty_cache()
                 tested_batch_size = max(1, tested_batch_size * 3 // 4)
 
-        # The while loop's condition (`> 1`) means batch size 1 is never
-        # actually probed. If every size down to 2 OOMs, it exits without a
-        # successful probe. Only skip ahead to `tested_batch_size` when the
-        # probe actually ran; otherwise fall back to starting at 0 like the
-        # loop below always did originally, or the first item(s) get dropped.
         if first_batch_scores is None:
             sorted_scores: List[float] = []
             loop_start = 0
@@ -316,8 +254,8 @@ class KaLMReranker:
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
             raise RuntimeError(
-                "CUDA ran out of memory during reranking. Retry with a smaller batch_size "
-                "or shorter max_length."
+                "CUDA ran out of memory during reranking. Retry with a smaller "
+                "batch_size or shorter max_length."
             ) from error
         inverse_indices = np.argsort(length_sorted_indices)
         return [sorted_scores[index] for index in inverse_indices]

kalm_reranker_utils.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from transformers.modeling_outputs import BaseModelOutput
+
+DEFAULT_INSTRUCTION = "Given a query, retrieve documents that answer the query."
+DEFAULT_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".'
+)
+
+
+def validate_text_pairs(inputs: Sequence[Sequence[str]]) -> list[tuple[str, str]]:
+    """Validate and normalize a batch of ``(query, document)`` pairs."""
+    if isinstance(inputs, (str, bytes)) or not isinstance(inputs, Sequence):
+        raise TypeError("inputs must be a sequence of (query, document) pairs.")
+
+    validated: list[tuple[str, str]] = []
+    for index, pair in enumerate(inputs):
+        if (
+            isinstance(pair, (str, bytes))
+            or not isinstance(pair, Sequence)
+            or len(pair) != 2
+        ):
+            raise ValueError(f"inputs[{index}] must contain exactly two strings.")
+        query, document = pair
+        if not isinstance(query, str) or not isinstance(document, str):
+            raise TypeError(f"inputs[{index}] must contain exactly two strings.")
+        validated.append((query, document))
+    return validated
+
+
+def answer_token_id(tokenizer: Any, answer: str) -> int:
+    """Return the single vocabulary token used to score an answer."""
+    token_ids = tokenizer(answer, add_special_tokens=False)["input_ids"]
+    if len(token_ids) != 1:
+        raise ValueError(
+            f"The answer {answer!r} must tokenize to exactly one token, "
+            f"got {token_ids!r}."
+        )
+    return token_ids[0]
+
+
+def build_decoder_text(
+    tokenizer: Any,
+    query: str,
+    instruction: str,
+    system_instruction: str,
+    query_max_length: int,
+) -> str:
+    """Build the decoder-side instruction/query prompt used during training."""
+    query_ids = tokenizer(
+        query,
+        add_special_tokens=False,
+        truncation=True,
+        max_length=query_max_length,
+    )["input_ids"]
+    truncated_query = tokenizer.decode(
+        query_ids,
+        skip_special_tokens=False,
+        clean_up_tokenization_spaces=False,
+    )
+    return (
+        "<bos><start_of_turn>user\n"
+        f"{system_instruction}\n\n"
+        f"<Instruct>: {instruction}\n"
+        f"<Query>: {truncated_query}<end_of_turn>\n"
+        "<start_of_turn>model\n\n\n\n"
+    )
+
+
+def get_encoder(model: torch.nn.Module) -> torch.nn.Module:
+    if hasattr(model, "get_encoder"):
+        return model.get_encoder()
+    if hasattr(model, "encoder"):
+        return model.encoder
+    raise AttributeError(f"Cannot find the encoder on {type(model).__name__}.")
+
+
+def pool_encoder_chunks(
+    hidden_states: torch.Tensor,
+    attention_mask: torch.Tensor,
+    chunk_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Mean-pool consecutive encoder tokens while respecting padding."""
+    if chunk_size <= 0:
+        raise ValueError("chunk_size must be positive.")
+
+    batch_size, sequence_length, hidden_size = hidden_states.shape
+    num_chunks = (sequence_length + chunk_size - 1) // chunk_size
+    padded_length = num_chunks * chunk_size
+    pad_length = padded_length - sequence_length
+
+    if pad_length:
+        hidden_states = F.pad(hidden_states, (0, 0, 0, pad_length))
+        attention_mask = F.pad(attention_mask, (0, pad_length))
+
+    hidden_states = hidden_states.view(batch_size, num_chunks, chunk_size, hidden_size)
+    chunk_mask = attention_mask.view(batch_size, num_chunks, chunk_size)
+    expanded_mask = chunk_mask.unsqueeze(-1).to(hidden_states.dtype)
+    pooled_hidden = (hidden_states * expanded_mask).sum(dim=2)
+    pooled_hidden = pooled_hidden / chunk_mask.sum(dim=2).clamp(min=1).unsqueeze(-1)
+    pooled_mask = (chunk_mask.sum(dim=2) > 0).to(attention_mask.dtype)
+    return pooled_hidden, pooled_mask
+
+
+def forward_reranker_model(
+    model: torch.nn.Module,
+    *,
+    input_ids: torch.Tensor,
+    attention_mask: torch.Tensor,
+    decoder_input_ids: torch.Tensor,
+    decoder_attention_mask: torch.Tensor,
+    encoder_chunk_size: int | None,
+):
+    """Run the encoder-decoder model with optional encoder token compression."""
+    if encoder_chunk_size is None:
+        return model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            return_dict=True,
+        )
+
+    encoder_outputs = get_encoder(model)(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        return_dict=True,
+    )
+    pooled_hidden, pooled_mask = pool_encoder_chunks(
+        encoder_outputs.last_hidden_state,
+        attention_mask,
+        encoder_chunk_size,
+    )
+    return model(
+        encoder_outputs=BaseModelOutput(last_hidden_state=pooled_hidden),
+        attention_mask=pooled_mask,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        return_dict=True,
+    )
+
+
+def extract_yes_no_logits(
+    logits: torch.Tensor,
+    decoder_attention_mask: torch.Tensor,
+    yes_token_id: int,
+    no_token_id: int,
+) -> torch.Tensor:
+    """Extract float32 yes/no logits at each sample's final non-padding token."""
+    if decoder_attention_mask.ndim != 2:
+        raise ValueError("decoder_attention_mask must have shape [batch, sequence].")
+    sequence_lengths = decoder_attention_mask.sum(dim=1) - 1
+    if (sequence_lengths < 0).any():
+        raise ValueError(
+            "Every decoder input must contain at least one non-padding token."
+        )
+
+    batch_indices = torch.arange(logits.shape[0], device=logits.device)
+    last_logits = logits[batch_indices, sequence_lengths]
+    yes_no_logits = torch.stack(
+        (last_logits[:, yes_token_id], last_logits[:, no_token_id]), dim=-1
+    ).float()
+    if not torch.isfinite(yes_no_logits).all():
+        bad_count = (~torch.isfinite(yes_no_logits).all(dim=-1)).sum().item()
+        raise RuntimeError(
+            f"The model produced non-finite yes/no logits for {bad_count} input(s). "
+            "Use bfloat16 or float32 instead of float16."
+        )
+    return yes_no_logits
+
+
+def normalize_requested_dtype(dtype: Any) -> torch.dtype | None:
+    """Normalize a caller-provided dtype without changing the ``auto`` behavior."""
+    if dtype is None or dtype == "auto":
+        return None
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    if not isinstance(dtype, str):
+        return None
+    normalized = dtype.lower().removeprefix("torch.")
+    return {
+        "bfloat16": torch.bfloat16,
+        "bf16": torch.bfloat16,
+        "float16": torch.float16,
+        "fp16": torch.float16,
+        "float32": torch.float32,
+        "fp32": torch.float32,
+    }.get(normalized)
+
+
+def cast_floating_parameters(model: torch.nn.Module, dtype: torch.dtype | None) -> None:
+    """Cast model parameters while preserving checkpoint buffer dtypes."""
+    if dtype is None:
+        return
+    for parameter in model.parameters():
+        if parameter.is_floating_point() and parameter.dtype != dtype:
+            parameter.data = parameter.data.to(dtype=dtype)
+
+
+__all__ = [
+    "DEFAULT_INSTRUCTION",
+    "DEFAULT_SYSTEM_INSTRUCTION",
+    "answer_token_id",
+    "build_decoder_text",
+    "cast_floating_parameters",
+    "extract_yes_no_logits",
+    "forward_reranker_model",
+    "get_encoder",
+    "normalize_requested_dtype",
+    "pool_encoder_chunks",
+    "validate_text_pairs",
+]

modules.json

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