diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index 1c33c6852d..0198810cbf 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -49,7 +49,7 @@ The three supported scoring functions are as illustrated in the figure below. | `LlamaBidirectionalForSequenceClassification`C | 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`C | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2`(see note), etc. | [mxbai_rerank_v2.jinja](../../../examples/pooling/score/template/mxbai_rerank_v2.jinja) | ✅︎ | ✅︎ | -| `Qwen3ForSequenceClassification`C | 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`C | 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`C, `*ForCausalLM`C, 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 diff --git a/examples/pooling/score/template/zerank2.jinja b/examples/pooling/score/template/zerank2.jinja new file mode 100644 index 0000000000..8bd16e4301 --- /dev/null +++ b/examples/pooling/score/template/zerank2.jinja @@ -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 + '\n' + 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' -}} diff --git a/examples/pooling/score/using_template_offline.py b/examples/pooling/score/using_template_offline.py index abe8157d4c..1b7fe9f7c0 100644 --- a/examples/pooling/score/using_template_offline.py +++ b/examples/pooling/score/using_template_offline.py @@ -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) diff --git a/examples/pooling/score/using_template_online.py b/examples/pooling/score/using_template_online.py index f7dc109dd8..9cf0703ca4 100644 --- a/examples/pooling/score/using_template_online.py +++ b/examples/pooling/score/using_template_online.py @@ -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 """ diff --git a/tests/models/language/pooling_mteb_test/test_zerank.py b/tests/models/language/pooling_mteb_test/test_zerank.py new file mode 100644 index 0000000000..c90c80bc4d --- /dev/null +++ b/tests/models/language/pooling_mteb_test/test_zerank.py @@ -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)