[Model] Add zerank-2 reranker (#60415)

Signed-off-by: rdlhio <contact@rdlh.io>
2026-10-09 03:01:51 +00:00
parent 702ce313f4
commit a1cd0ec650
5 changed files with 119 additions and 1 deletions
@@ -49,7 +49,7 @@ The three supported scoring functions are as illustrated in the figure below.
 | `LlamaBidirectionalForSequenceClassification`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-rerank-1b-v2`, etc. | [nemotron-rerank.jinja](../../../examples/pooling/score/template/nemotron-rerank.jinja) | ✅︎ | ✅︎ |
 | `ModernBertForSequenceClassification` | ModernBERT-based | `Alibaba-NLP/gte-reranker-modernbert-base`, etc. | N/A | | |
 | `Qwen2ForSequenceClassification`<sup>C</sup> | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2`(see note), etc. | [mxbai_rerank_v2.jinja](../../../examples/pooling/score/template/mxbai_rerank_v2.jinja) | ✅︎ | ✅︎ |
-| `Qwen3ForSequenceClassification`<sup>C</sup> | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B`(see note), etc. | [qwen3_reranker.jinja](../../../examples/pooling/score/template/qwen3_reranker.jinja) | ✅︎ | ✅︎ |
+| `Qwen3ForSequenceClassification`<sup>C</sup> | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B`(see note), `polaria-tech/zerank-2-reranker-vllm`(see note), etc. | [qwen3_reranker.jinja](../../../examples/pooling/score/template/qwen3_reranker.jinja), [zerank2.jinja](../../../examples/pooling/score/template/zerank2.jinja) | ✅︎ | ✅︎ |
 | `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | N/A | ✅︎ | |
 | `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | N/A | ✅︎ | |
 | `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | N/A | \* | \* |
@@ -88,6 +88,13 @@ The three supported scoring functions are as illustrated in the figure below.
    vllm serve Qwen/Qwen3-Reranker-0.6B --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}'
    ```

+!!! note
+    `polaria-tech/zerank-2-reranker-vllm` is [zeroentropy/zerank-2-reranker](https://huggingface.co/zeroentropy/zerank-2-reranker) converted to `Qwen3ForSequenceClassification`; it returns zerank-2's calibrated score `sigmoid(logit / 5)`. Serve it with [zerank2.jinja](../../../examples/pooling/score/template/zerank2.jinja), which also supports the `instruction` request field.
+
+    ```bash
+    vllm serve polaria-tech/zerank-2-reranker-vllm --chat-template examples/pooling/score/template/zerank2.jinja
+    ```
+
 #### Multimodal Models

 !!! note
@@ -0,0 +1,13 @@
+{#- zerank-2 (zeroentropy/zerank-2-reranker): query in the system turn, document in the user turn.
+    The assistant suffix is emitted unconditionally, since score templates are rendered
+    without add_generation_prompt. An optional `instruction` is wrapped as recommended by
+    ZeroEntropy; without it, the prompt is the official zerank-2 prompt. -#}
+{%- set query = messages | selectattr("role", "eq", "query") | map(attribute="content") | list -%}
+{%- set document = messages | selectattr("role", "eq", "document") | map(attribute="content") | list -%}
+{%- set query = query[0] if query else '' -%}
+{%- if instruction is defined and instruction is string and instruction | trim -%}
+{%- set query = '<query>' + query + '</query>\n<instruction>' + instruction + '</instruction>' -%}
+{%- endif -%}
+{{- '<|im_start|>system\n' + query + '<|im_end|>\n' -}}
+{{- '<|im_start|>user\n' + (document[0] if document else '') + '<|im_end|>\n' -}}
+{{- '<|im_start|>assistant\n' -}}
@@ -53,6 +53,7 @@ def get_chat_template(model: str) -> str:
        "mixedbread-ai/mxbai-rerank-base-v2": "mxbai_rerank_v2.jinja",
        "mixedbread-ai/mxbai-rerank-large-v2": "mxbai_rerank_v2.jinja",
        "nvidia/llama-nemotron-rerank-1b-v2": "nemotron-rerank.jinja",
+        "polaria-tech/zerank-2-reranker-vllm": "zerank2.jinja",
    }

    # Get the template filename for the specified model
@@ -110,6 +111,7 @@ def get_hf_overrides(model: str) -> dict[str, Any]:
            "method": "from_2_way_softmax",
        },
        "nvidia/llama-nemotron-rerank-1b-v2": {},
+        "polaria-tech/zerank-2-reranker-vllm": {},
    }

    hf_overrides = model_name_to_hf_overrides_map.get(model)
@@ -19,6 +19,7 @@ run:
    vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --chat-template examples/pooling/score/template/qwen3_reranker.jinja
    vllm serve mixedbread-ai/mxbai-rerank-base-v2 --hf_overrides '{"architectures": ["Qwen2ForSequenceClassification"],"classifier_from_token": ["0", "1"], "method": "from_2_way_softmax"}' --chat-template examples/pooling/score/template/mxbai_rerank_v2.jinja
    vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja
+    vllm serve polaria-tech/zerank-2-reranker-vllm --runner pooling --chat-template examples/pooling/score/template/zerank2.jinja
    vllm serve Qwen/Qwen3-Reranker-0.6B --runner pooling --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}' --chat-template examples/pooling/score/template/qwen3_reranker.jinja

 """
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
+import mteb
+import numpy as np
+import pytest
+import torch
+from torch.utils.data import DataLoader
+
+from tests.conftest import HfRunner
+from tests.models.utils import RerankModelInfo
+
+from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models
+
+RERANK_MODELS = [
+    RerankModelInfo(
+        "polaria-tech/zerank-2-reranker-vllm",
+        architecture="Qwen3ForSequenceClassification",
+        chat_template_name="zerank2.jinja",
+        seq_pooling_type="LAST",
+        attn_type="decoder",
+        is_prefix_caching_supported=True,
+        is_chunked_prefill_supported=True,
+        mteb_score=0.34305,
+        mteb_tol=1e-2,
+        enable_test=True,
+    ),
+]
+
+
+class ZerankHfRunner(MtebCrossEncoderMixin, HfRunner):
+    # The converted checkpoint has no lm_head, so the reference is the original
+    # causal LM: the logit of the "Yes" token at the last position. The converted
+    # checkpoint returns sigmoid(logit / 5), which ranks documents identically.
+    original_model = "zeroentropy/zerank-2-reranker"
+
+    def __init__(
+        self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any
+    ) -> None:
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        HfRunner.__init__(
+            self,
+            model_name=self.original_model,
+            auto_cls=AutoModelForCausalLM,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.original_model, padding_side="left"
+        )
+        self.token_true_id = self.tokenizer.convert_tokens_to_ids("Yes")
+
+    @torch.no_grad
+    def predict(
+        self,
+        inputs1: DataLoader[mteb.types.BatchedInput],
+        inputs2: DataLoader[mteb.types.BatchedInput],
+        *args,
+        **kwargs,
+    ) -> np.ndarray:
+        queries = [text for batch in inputs1 for text in batch["text"]]
+        corpus = [text for batch in inputs2 for text in batch["text"]]
+
+        tokenizer = self.tokenizer
+        prompts = []
+        for query, document in zip(queries, corpus):
+            conversation = [
+                {"role": "query", "content": query},
+                {"role": "document", "content": document},
+            ]
+
+            prompt = tokenizer.apply_chat_template(
+                conversation=conversation,
+                tools=None,
+                chat_template=self.chat_template,
+                tokenize=False,
+            )
+            prompts.append(prompt)
+
+        scores = []
+        for prompt in prompts:
+            inputs = tokenizer([prompt], return_tensors="pt")
+            inputs = self.wrap_device(inputs)
+            logits = self.model(**inputs).logits[:, -1, :]
+            scores.append(logits[0, self.token_true_id].item())
+        return torch.Tensor(scores)
+
+
+@pytest.mark.flaky(reruns=2)
+@pytest.mark.parametrize("model_info", RERANK_MODELS)
+def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
+    mteb_test_rerank_models(vllm_runner, model_info, hf_runner=ZerankHfRunner)