Text Ranking
sentence-transformers
Safetensors
Transformers
multilingual
t5gemma2
text2text-generation
reranker
encoder-decoder
FBNL
Retrieval
RAG

Add sentence-transformers CrossEncoder support

#2
by cosyy - opened
KaLM-Embedding org

Sentence Transformers CrossEncoder support

Summary

This PR adds model-side support for loading the reranker through the modular
Sentence Transformers CrossEncoder API:

import torch
from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "KaLM-Embedding/KaLM-Reranker-V1-Large",
    trust_remote_code=True,
    device="cuda",
    model_kwargs={"dtype": torch.bfloat16, "chunk_size": 4},
)

scores = model.predict([(query, document)])
rankings = model.rank(query, documents)

The document is encoded on the encoder side, while the task instruction and
query remain on the decoder side. The default Sigmoid activation returns
P(yes); passing activation_fn=torch.nn.Identity() returns
yes_logit - no_logit.

chunk_size=N performs mask-aware mean pooling over each consecutive group of
N encoder token states. Set chunk_size=None to disable encoder compression.

Files

The PR changes only these six adapter/configuration files:

  • kalm_reranker.py
  • kalm_reranker_utils.py
  • kalm_cross_encoder.py
  • kalm_cross_encoder_config.json
  • modules.json
  • config_sentence_transformers.json

It does not modify or upload the repository README, model weights, tokenizer
artifacts, generation configuration, or Transformers config.json.

Compatibility

  • sentence-transformers>=5.6,<6
  • transformers>=5.3,<6
  • PyTorch inference backend only

The adapter code is shared with the validated
KaLM-Embedding/KaLM-Reranker-V1-Nano CrossEncoder integration. Small and
Large retain their existing model-specific Transformers configurations and
weights.

KaLM-Embedding org

KaLM-Reranker-V1-Large validation

The six files in this PR were verified byte-for-byte against the locally tested
adapter files. The repository README at the PR revision is unchanged from
main.

H100 BF16 numerical smoke test

Item Result
sentence-transformers 5.6.0
transformers 5.3.0
torch 2.6.0+cu126
encoder chunk size 4
batch size 2
tolerance rtol=1e-3, atol=2e-3
max absolute error 2.9103830456733704e-11
ranking corpus IDs [0, 1]

Reference KaLMReranker scores and CrossEncoder scores agree within the BF16
tolerance for the five mixed English, Chinese, empty-query, and empty-document
inputs. The existing Large config.json, tokenizer, and model weights were
used without modification.

Immutable local checkpoint hashes observed during validation:

  • config.json: cb9cd9526c0669c3397e0e332215f9b2b642edca075dea43819002ce6debdfe2
  • model.safetensors: 94c59b9ce11b0f36c456419a3e43c9ca9c3a9ab4c86c1f97e51fdb896317e002
  • tokenizer.json: f5b325224482ec441ec5fbe2a5ac08c3758e0f9605f6e54368e31f736fcfb01d
cosyy changed pull request status to merged

Sign up or log in to comment