Add sentence-transformers CrossEncoder support

by cosyy - opened 7 days ago

base: refs/heads/main

←

from: refs/pr/2

Discussion Files changed

+602

-121

cosyy

KaLM-Embedding org 7 days ago

•

edited 7 days ago

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:

[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

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

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.
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.

Add sentence-transformers CrossEncoder supportd43b4ac9

cosyy

KaLM-Embedding org 7 days ago

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:

[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 `refs/pr/2`	CPU FP32	passed	`1.4901161193847656e-08`	Custom module loaded from isolated remote-code cache
Hub PR revision `refs/pr/2`	CUDA BF16	passed	`3.637978807091713e-12`	H100 MIG 3g.40gb; isolated model and remote-code caches

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:

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

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.

Hub PR

PR: https://huggingface.co/KaLM-Embedding/KaLM-Reranker-V1-Nano/discussions/2
Commit: d43b4ac9e141b11cbcfcae5f7e5b64b73ae31f26
Revision tested: refs/pr/2
Uploaded files: six adapter/configuration files; no README, weights,
tokenizer artifacts, or Transformers config.json
CPU FP32 and CUDA BF16 were each downloaded and executed from a separate
empty model and remote-code cache.

No validation work remains for the PR revision. Merging the PR and manually
editing the existing model README remain subject to repository-owner review.

cosyy changed pull request status to merged 7 days ago

Upload images, audio, and videos by dragging in the text input, pasting, or clicking here.

Tap or paste here to upload images

· Sign up or log in to comment