A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 19:00 UTC
[Model] Let more models to support the score template. (#31335)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
wang.yuqi committed
Jan 5, 2026 at 19:54 UTC
911d38ed99e1d191d3654f42fbeee97736092da8
23 files changed
+767
-337
docs/models/supported_models.md
+19
-17
@@ -540,21 +540,28 @@ If your model is not in the above list, we will try to automatically convert the
540
Cross-encoder and reranker models are a subset of classification models that accept two prompts as input.
541
These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) API.
542
543
-| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
544
-|--------------|--------|-------------------|----------------------|---------------------------|
545
-| `BertForSequenceClassification` | BERT-based | `cross-encoder/ms-marco-MiniLM-L-6-v2`, etc. | | |
546
-| `GemmaForSequenceClassification` | Gemma-based | `BAAI/bge-reranker-v2-gemma` (see note), etc. | ✅︎ | ✅︎ |
547
-| `GteNewForSequenceClassification` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-reranker-base`, etc. | | |
548
-| `LlamaBidirectionalForSequenceClassification`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-rerank-1b-v2` (see note), etc. | ✅︎ | ✅︎ |
549
-| `Qwen2ForSequenceClassification`<sup>C</sup> | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2` (see note), etc. | ✅︎ | ✅︎ |
550
-| `Qwen3ForSequenceClassification`<sup>C</sup> | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B` (see note), etc. | ✅︎ | ✅︎ |
551
-| `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | | |
552
-| `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | | |
553
-| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
543
+| Architecture | Models | Example HF Models | Score template (see note) | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
544
+|--------------|--------|-------------------|---------------------------|-----------------------------|-----------------------------------------|
545
+| `BertForSequenceClassification` | BERT-based | `cross-encoder/ms-marco-MiniLM-L-6-v2`, etc. | N/A | | |
546
+| `GemmaForSequenceClassification` | Gemma-based | `BAAI/bge-reranker-v2-gemma`(see note), etc. | [bge-reranker-v2-gemma.jinja](../../examples/pooling/score/template/bge-reranker-v2-gemma.jinja) | ✅︎ | ✅︎ |
547
+| `GteNewForSequenceClassification` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-reranker-base`, etc. | N/A | | |
548
+| `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) | ✅︎ | ✅︎ |
549
+| `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) | ✅︎ | ✅︎ |
550
+| `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) | ✅︎ | ✅︎ |
551
+| `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | N/A | | |
552
+| `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | N/A | | |
553
+| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | N/A | \* | \* |
554
555
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./pooling_models.md#model-conversion))
556
\* Feature support is the same as that of the original model.
557
558
+!!! note
559
+ Some models require a specific prompt format to work correctly.
560
+
561
+ You can find Example HF Models's corresponding score template in [examples/pooling/score/template/](../../examples/pooling/score/template)
562
+
563
+ Examples : [examples/pooling/score/using_template_offline.py](../../examples/pooling/score/using_template_offline.py) [examples/pooling/score/using_template_online.py](../../examples/pooling/score/using_template_online.py)
564
+
565
!!! note
566
Load the official original `BAAI/bge-reranker-v2-gemma` by using the following command.
567
@@ -565,11 +572,6 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
572
!!! note
573
The second-generation GTE model (mGTE-TRM) is named `NewForSequenceClassification`. The name `NewForSequenceClassification` is too generic, you should set `--hf-overrides '{"architectures": ["GteNewForSequenceClassification"]}'` to specify the use of the `GteNewForSequenceClassification` architecture.
574
568
-!!! note
569
- `nvidia/llama-nemotron-rerank-1b-v2` require a specific prompt format to work correctly.
570
-
571
- Examples : [offline_using_template.py](../../examples/pooling/score/offline_using_template.py) [online_using_template.py](../../examples/pooling/score/online_using_template.py)
572
-
575
!!! note
576
Load the official original `mxbai-rerank-v2` by using the following command.
577
@@ -578,7 +580,7 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
580
```
581
582
!!! note
581
- Load the official original `Qwen3 Reranker` by using the following command. More information can be found at: [examples/pooling/score/offline_reranker.py](../../examples/pooling/score/offline_reranker.py).
583
+ Load the official original `Qwen3 Reranker` by using the following command. More information can be found at: [examples/pooling/score/qwen3_reranker_offline.py](../../examples/pooling/score/qwen3_reranker_offline.py) [examples/pooling/score/qwen3_reranker_online.py](../../examples/pooling/score/qwen3_reranker_online.py).
584
585
```bash
586
vllm serve Qwen/Qwen3-Reranker-0.6B --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}'
examples/pooling/score/convert_model_to_seq_cls.py
+93
-19
@@ -2,35 +2,70 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# ruff: noqa: E501
4
5
+"""
6
+Script to convert Large Language Models (LLMs) to Sequence Classification models.
7
+This is particularly useful for converting reranker models that use next-token
8
+prediction to a sequence classification format for compatibility with standard
9
+classification and rerank pipelines.
10
+
11
+Usage examples:
12
+- For BAAI/bge-reranker-v2-gemma:
13
+ python convert_model_to_seq_cls.py --model_name BAAI/bge-reranker-v2-gemma \
14
+ --classifier_from_tokens '["Yes"]' --method no_post_processing \
15
+ --path ./bge-reranker-v2-gemma-seq-cls
16
+
17
+- For mxbai-rerank-v2:
18
+ python convert_model_to_seq_cls.py --model_name mixedbread-ai/mxbai-rerank-base-v2 \
19
+ --classifier_from_tokens '["0", "1"]' --method from_2_way_softmax \
20
+ --path ./mxbai-rerank-base-v2-seq-cls
21
+
22
+- For Qwen3-Reranker:
23
+ python convert_model_to_seq_cls.py --model_name Qwen/Qwen3-Reranker-0.6B \
24
+ --classifier_from_tokens '["no", "yes"]' --method from_2_way_softmax \
25
+ --path ./Qwen3-Reranker-0.6B-seq-cls
26
+
27
+Note: For BAAI/bge-reranker-v2-gemma, "Yes" and "yes" are different tokens.
28
+"""
29
+
30
import argparse
31
import json
32
33
import torch
34
import transformers
35
11
-# Usage:
12
-# for BAAI/bge-reranker-v2-gemma
13
-# Caution: "Yes" and "yes" are two different tokens
14
-# python convert_model_to_seq_cls.py --model_name BAAI/bge-reranker-v2-gemma --classifier_from_tokens '["Yes"]' --method no_post_processing --path ./bge-reranker-v2-gemma-seq-cls
15
-# for mxbai-rerank-v2
16
-# python convert_model_to_seq_cls.py --model_name mixedbread-ai/mxbai-rerank-base-v2 --classifier_from_tokens '["0", "1"]' --method from_2_way_softmax --path ./mxbai-rerank-base-v2-seq-cls
17
-# for Qwen3-Reranker
18
-# python convert_model_to_seq_cls.py --model_name Qwen/Qwen3-Reranker-0.6B --classifier_from_tokens '["no", "yes"]' --method from_2_way_softmax --path ./Qwen3-Reranker-0.6B-seq-cls
19
-
36
37
def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
22
- # refer to https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
23
- assert len(tokens) == 2
38
+ """
39
+ This method extracts the difference between weights for 'true' and 'false' tokens
40
+ from the language model head to create a single classification weight vector.
41
+
42
+ Args:
43
+ causal_lm: The original causal language model
44
+ seq_cls_model: The target sequence classification model
45
+ tokenizer: Model tokenizer
46
+ tokens: List of two tokens representing [false_token, true_token]
47
+ device: Target device (cpu/cuda)
48
+
49
+ Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
50
+ """
51
+ assert len(tokens) == 2, (
52
+ "Method requires exactly two tokens for binary classification"
53
+ )
54
55
+ # Get the language model head weights (vocabulary_size x hidden_size)
56
lm_head_weights = causal_lm.lm_head.weight
57
58
+ # Convert token strings to their corresponding token IDs
59
false_id = tokenizer.convert_tokens_to_ids(tokens[0])
60
true_id = tokenizer.convert_tokens_to_ids(tokens[1])
61
62
+ # Compute the classification weight as the difference between true and false token weights
63
+ # This follows the approach in: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
64
score_weight = lm_head_weights[true_id].to(device).to(
65
torch.float32
66
) - lm_head_weights[false_id].to(device).to(torch.float32)
67
68
+ # Copy the computed weights to the sequence classification model
69
with torch.no_grad():
70
seq_cls_model.score.weight.copy_(score_weight.unsqueeze(0))
71
if seq_cls_model.score.bias is not None:
@@ -38,12 +73,29 @@ def from_2_way_softmax(causal_lm, seq_cls_model, tokenizer, tokens, device):
73
74
75
def no_post_processing(causal_lm, seq_cls_model, tokenizer, tokens, device):
76
+ """
77
+ Directly use token weights from the language model head for classification.
78
+
79
+ This method maps each classification label directly to a corresponding token
80
+ in the vocabulary without additional transformation.
81
+
82
+ Args:
83
+ causal_lm: The original causal language model
84
+ seq_cls_model: The target sequence classification model
85
+ tokenizer: Model tokenizer
86
+ tokens: List of tokens representing class labels
87
+ device: Target device (cpu/cuda)
88
+ """
89
+ # Get the language model head weights (vocabulary_size x hidden_size)
90
lm_head_weights = causal_lm.lm_head.weight
91
92
+ # Convert all tokens to their corresponding token IDs
93
token_ids = [tokenizer.convert_tokens_to_ids(t) for t in tokens]
94
95
+ # Extract weights for the specific tokens (num_tokens x hidden_size)
96
score_weight = lm_head_weights[token_ids].to(device)
97
98
+ # Copy the weights to the sequence classification model
99
with torch.no_grad():
100
seq_cls_model.score.weight.copy_(score_weight)
101
if seq_cls_model.score.bias is not None:
@@ -58,19 +110,33 @@ method_map = {
110
def converting(
111
model_name, classifier_from_tokens, path, method, use_pad_token=False, device="cpu"
112
):
61
- assert method in method_map
62
-
113
+ """
114
+ Main conversion function to transform a CausalLM model to SequenceClassification.
115
+
116
+ Args:
117
+ model_name: Name or path of the pretrained model
118
+ classifier_from_tokens: List of tokens used for classification
119
+ path: Output path to save the converted model
120
+ method: Conversion method ('from_2_way_softmax' or 'no_post_processing')
121
+ use_pad_token: Whether to use padding token in the sequence classification model
122
+ device: Device to load the model on ('cpu' or 'cuda')
123
+ """
124
+ assert method in method_map, f"Unknown method: {method}"
125
+
126
+ # Determine number of labels based on conversion method
127
if method == "from_2_way_softmax":
128
assert len(classifier_from_tokens) == 2
129
num_labels = 1
130
else:
131
num_labels = len(classifier_from_tokens)
132
133
+ # Load tokenizer and original causal language model
134
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
135
causal_lm = transformers.AutoModelForCausalLM.from_pretrained(
136
model_name, device_map=device
137
)
138
139
+ # Load an empty sequence classification model with the same architecture
140
seq_cls_model = transformers.AutoModelForSequenceClassification.from_pretrained(
141
model_name,
142
num_labels=num_labels,
@@ -78,14 +144,17 @@ def converting(
144
device_map=device,
145
)
146
147
+ # Apply the selected conversion method to transfer weights
148
method_map[method](
149
causal_lm, seq_cls_model, tokenizer, classifier_from_tokens, device
150
)
151
85
- # `llm as reranker` defaults to not using pad_token
152
+ # Configure padding token settings
153
+ # Note: Reranker models typically don't use padding tokens by default
154
seq_cls_model.config.use_pad_token = use_pad_token
155
seq_cls_model.config.pad_token_id = tokenizer.pad_token_id
156
157
+ # Save the converted model and tokenizer
158
seq_cls_model.save_pretrained(path)
159
tokenizer.save_pretrained(path)
160
@@ -99,25 +168,30 @@ def parse_args():
168
"--model_name",
169
type=str,
170
default="BAAI/bge-reranker-v2-gemma",
102
- help="Model name",
171
+ help="HuggingFace model name or local path",
172
)
173
parser.add_argument(
174
"--classifier_from_tokens",
175
type=str,
176
default='["Yes"]',
108
- help="classifier from tokens",
177
+ help="JSON string of tokens used for classification labels",
178
)
179
parser.add_argument(
111
- "--method", type=str, default="no_post_processing", help="Converting converting"
180
+ "--method",
181
+ type=str,
182
+ default="no_post_processing",
183
+ help="Conversion method to use",
184
)
185
parser.add_argument(
114
- "--use-pad-token", action="store_true", help="Whether to use pad_token"
186
+ "--use-pad-token",
187
+ action="store_true",
188
+ help="Enable padding token in the sequence classification model",
189
)
190
parser.add_argument(
191
"--path",
192
type=str,
193
default="./bge-reranker-v2-gemma-seq-cls",
120
- help="Path to save converted model",
194
+ help="Output directory to save the converted model",
195
)
196
return parser.parse_args()
197
examples/pooling/score/offline_reranker.py
deleted
-89
@@ -1,89 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# ruff: noqa: E501
4
-
5
-from vllm import LLM
6
-
7
-model_name = "Qwen/Qwen3-Reranker-0.6B"
8
-
9
-# What is the difference between the official original version and one
10
-# that has been converted into a sequence classification model?
11
-# Qwen3-Reranker is a language model that doing reranker by using the
12
-# logits of "no" and "yes" tokens.
13
-# It needs to computing 151669 tokens logits, making this method extremely
14
-# inefficient, not to mention incompatible with the vllm score API.
15
-# A method for converting the original model into a sequence classification
16
-# model was proposed. See：https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
17
-# Models converted offline using this method can not only be more efficient
18
-# and support the vllm score API, but also make the init parameters more
19
-# concise, for example.
20
-# llm = LLM(model="tomaarsen/Qwen3-Reranker-0.6B-seq-cls", runner="pooling")
21
-
22
-# If you want to load the official original version, the init parameters are
23
-# as follows.
24
-
25
-
26
-def get_llm() -> LLM:
27
- """Initializes and returns the LLM model for Qwen3-Reranker."""
28
- return LLM(
29
- model=model_name,
30
- runner="pooling",
31
- hf_overrides={
32
- "architectures": ["Qwen3ForSequenceClassification"],
33
- "classifier_from_token": ["no", "yes"],
34
- "is_original_qwen3_reranker": True,
35
- },
36
- )
37
-
38
-
39
-# Why do we need hf_overrides for the official original version:
40
-# vllm converts it to Qwen3ForSequenceClassification when loaded for
41
-# better performance.
42
-# - Firstly, we need using `"architectures": ["Qwen3ForSequenceClassification"],`
43
-# to manually route to Qwen3ForSequenceClassification.
44
-# - Then, we will extract the vector corresponding to classifier_from_token
45
-# from lm_head using `"classifier_from_token": ["no", "yes"]`.
46
-# - Third, we will convert these two vectors into one vector. The use of
47
-# conversion logic is controlled by `using "is_original_qwen3_reranker": True`.
48
-
49
-# Please use the query_template and document_template to format the query and
50
-# document for better reranker results.
51
-
52
-prefix = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
53
-suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
54
-
55
-query_template = "{prefix}<Instruct>: {instruction}\n<Query>: {query}\n"
56
-document_template = "<Document>: {doc}{suffix}"
57
-
58
-
59
-def main() -> None:
60
- instruction = (
61
- "Given a web search query, retrieve relevant passages that answer the query"
62
- )
63
-
64
- queries = [
65
- "What is the capital of China?",
66
- "Explain gravity",
67
- ]
68
-
69
- documents = [
70
- "The capital of China is Beijing.",
71
- "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
72
- ]
73
-
74
- queries = [
75
- query_template.format(prefix=prefix, instruction=instruction, query=query)
76
- for query in queries
77
- ]
78
- documents = [document_template.format(doc=doc, suffix=suffix) for doc in documents]
79
-
80
- llm = get_llm()
81
- outputs = llm.score(queries, documents)
82
-
83
- print("-" * 30)
84
- print([output.outputs.score for output in outputs])
85
- print("-" * 30)
86
-
87
-
88
-if __name__ == "__main__":
89
- main()
examples/pooling/score/offline_using_template.py
deleted
-27
@@ -1,27 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# ruff: noqa: E501
4
-from pathlib import Path
5
-
6
-from vllm import LLM
7
-
8
-model_name = "nvidia/llama-nemotron-rerank-1b-v2"
9
-
10
-# Path to template file
11
-template_path = Path(__file__).parent / "template" / "nemotron-rerank.jinja"
12
-chat_template = template_path.read_text()
13
-
14
-llm = LLM(model=model_name, runner="pooling", trust_remote_code=True)
15
-
16
-query = "how much protein should a female eat?"
17
-documents = [
18
- "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
19
- "Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
20
- "Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
21
-]
22
-
23
-outputs = llm.score(query, documents, chat_template=chat_template)
24
-
25
-print("-" * 30)
26
-print([output.outputs.score for output in outputs])
27
-print("-" * 30)
examples/pooling/score/online_using_template.py
deleted
-46
@@ -1,46 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# ruff: noqa: E501
4
-"""
5
-Example of using the rerank API with template.
6
-
7
-run:
8
- vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja
9
-"""
10
-
11
-import json
12
-
13
-import requests
14
-
15
-url = "http://127.0.0.1:8000/rerank"
16
-
17
-headers = {"accept": "application/json", "Content-Type": "application/json"}
18
-
19
-query = "how much protein should a female eat?"
20
-documents = [
21
- "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
22
- "Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
23
- "Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
24
-]
25
-
26
-data = {
27
- "model": "nvidia/llama-nemotron-rerank-1b-v2",
28
- "query": query,
29
- "documents": documents,
30
-}
31
-
32
-
33
-def main():
34
- response = requests.post(url, headers=headers, json=data)
35
-
36
- # Check the response
37
- if response.status_code == 200:
38
- print("Request successful!")
39
- print(json.dumps(response.json(), indent=2))
40
- else:
41
- print(f"Request failed with status code: {response.status_code}")
42
- print(response.text)
43
-
44
-
45
-if __name__ == "__main__":
46
- main()
examples/pooling/score/qwen3_reranker_offline.py
new
+104
@@ -0,0 +1,104 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+
5
+"""
6
+What is the difference between the official original version and one
7
+that has been converted into a sequence classification model?
8
+
9
+Qwen3-Reranker is a language model that doing reranker by using the
10
+logits of "no" and "yes" tokens.
11
+This requires computing logits for all 151,669 tokens in the vocabulary,
12
+making it inefficient and incompatible with vLLM's score() API.
13
+
14
+A conversion method has been proposed to transform the original model into a
15
+sequence classification model. This converted model:
16
+1. Is significantly more efficient
17
+2. Fully supports vLLM's score() API
18
+3. Simplifies initialization parameters
19
+Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
20
+Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
21
+
22
+For the converted model, initialization would simply be:
23
+llm = LLM(model="tomaarsen/Qwen3-Reranker-0.6B-seq-cls", runner="pooling")
24
+
25
+This example demonstrates loading the ORIGINAL model with special overrides
26
+to make it compatible with vLLM's score API.
27
+"""
28
+
29
+from pathlib import Path
30
+
31
+from vllm import LLM
32
+
33
+model_name = "Qwen/Qwen3-Reranker-0.6B"
34
+
35
+
36
+def get_llm() -> LLM:
37
+ """
38
+ Initializes and returns the LLM model for Qwen3-Reranker.
39
+
40
+ Returns:
41
+ LLM: Configured vLLM instance for reranking tasks.
42
+
43
+ Note:
44
+ This function loads the ORIGINAL Qwen3-Reranker model with specific
45
+ overrides to make it compatible with vLLM's score API.
46
+ """
47
+ return LLM(
48
+ # Specify the original model from HuggingFace
49
+ model=model_name,
50
+ # Use pooling runner for score task
51
+ runner="pooling",
52
+ # HuggingFace model configuration overrides required for compatibility
53
+ hf_overrides={
54
+ # Manually route to sequence classification architecture
55
+ # This tells vLLM to use Qwen3ForSequenceClassification instead of
56
+ # the default Qwen3ForCausalLM
57
+ "architectures": ["Qwen3ForSequenceClassification"],
58
+ # Specify which token logits to extract from the language model head
59
+ # The original reranker uses "no" and "yes" token logits for scoring
60
+ "classifier_from_token": ["no", "yes"],
61
+ # Enable special handling for original Qwen3-Reranker models
62
+ # This flag triggers conversion logic that transforms the two token
63
+ # vectors into a single classification vector
64
+ "is_original_qwen3_reranker": True,
65
+ },
66
+ )
67
+
68
+
69
+def main() -> None:
70
+ # Load the Jinja template for formatting query-document pairs
71
+ # The template ensures proper formatting for the reranker model
72
+ template_home = Path(__file__).parent / "template"
73
+ template_path = "qwen3_reranker.jinja"
74
+ chat_template = (template_home / template_path).read_text()
75
+
76
+ # Sample queries for testing the reranker
77
+ queries = [
78
+ "What is the capital of China?",
79
+ "Explain gravity",
80
+ ]
81
+
82
+ # Corresponding documents to be scored against each query
83
+ documents = [
84
+ "The capital of China is Beijing.",
85
+ "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
86
+ ]
87
+
88
+ # Initialize the LLM model with the original Qwen3-Reranker configuration
89
+ llm = get_llm()
90
+
91
+ # Compute relevance scores for each query-document pair
92
+ # The score() method returns a relevance score for each pair
93
+ # Higher scores indicate better relevance
94
+ outputs = llm.score(queries, documents, chat_template=chat_template)
95
+
96
+ # Extract and print the relevance scores from the outputs
97
+ # Each output contains a score representing query-document relevance
98
+ print("-" * 30)
99
+ print("Relevance scores:", [output.outputs.score for output in outputs])
100
+ print("-" * 30)
101
+
102
+
103
+if __name__ == "__main__":
104
+ main()
examples/pooling/score/qwen3_reranker_online.py
new
+80
@@ -0,0 +1,80 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+"""
5
+What is the difference between the official original version and one
6
+that has been converted into a sequence classification model?
7
+
8
+Qwen3-Reranker is a language model that doing reranker by using the
9
+logits of "no" and "yes" tokens.
10
+This requires computing logits for all 151,669 tokens in the vocabulary,
11
+making it inefficient and incompatible with vLLM's score() API.
12
+
13
+A conversion method has been proposed to transform the original model into a
14
+sequence classification model. This converted model:
15
+1. Is significantly more efficient
16
+2. Fully supports vLLM's score() API
17
+3. Simplifies initialization parameters
18
+Reference: https://huggingface.co/Qwen/Qwen3-Reranker-0.6B/discussions/3
19
+Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
20
+
21
+For the converted model, initialization would simply be:
22
+ vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --runner pooling --chat-template examples/pooling/score/template/qwen3_reranker.jinja
23
+
24
+This example demonstrates loading the ORIGINAL model with special overrides
25
+to make it compatible with vLLM's score API.
26
+ 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
27
+"""
28
+
29
+import json
30
+
31
+import requests
32
+
33
+# URL of the vLLM server's score endpoint
34
+# Default vLLM server runs on localhost port 8000
35
+url = "http://127.0.0.1:8000/score"
36
+
37
+# HTTP headers for the request
38
+headers = {"accept": "application/json", "Content-Type": "application/json"}
39
+
40
+# Example queries & documents
41
+queries = [
42
+ "What is the capital of China?",
43
+ "Explain gravity",
44
+]
45
+documents = [
46
+ "The capital of China is Beijing.",
47
+ "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
48
+]
49
+
50
+# Request payload for the score API
51
+data = {
52
+ "model": "Qwen/Qwen3-Reranker-0.6B",
53
+ "text_1": queries,
54
+ "text_2": documents,
55
+}
56
+
57
+
58
+def main():
59
+ """Main function to send a score request to the vLLM server.
60
+
61
+ This function sends a POST request to the /score endpoint with
62
+ the query and documents, then prints the relevance scores.
63
+ """
64
+ # Send POST request to the vLLM server's score endpoint
65
+ response = requests.post(url, headers=headers, json=data)
66
+
67
+ # Check if the request was successful
68
+ if response.status_code == 200:
69
+ print("Request successful!")
70
+ # Pretty print the JSON response containing relevance scores
71
+ # The response includes scores for each document's relevance to the query
72
+ print(json.dumps(response.json(), indent=2))
73
+ else:
74
+ # Handle request failure
75
+ print(f"Request failed with status code: {response.status_code}")
76
+ print(response.text)
77
+
78
+
79
+if __name__ == "__main__":
80
+ main()
examples/pooling/score/template/bge-reranker-v2-gemma.jinja
new
+3
@@ -0,0 +1,3 @@
1
+A: {{ (messages | selectattr("role", "eq", "query") | first).content }}
2
+B: {{ (messages | selectattr("role", "eq", "document") | first).content }}
3
+Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'.
\ No newline at end of file
examples/pooling/score/template/mxbai_rerank_v2.jinja
new
+8
@@ -0,0 +1,8 @@
1
+<|im_start|>system
2
+You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
3
+<|im_start|>user
4
+query: {{ (messages | selectattr("role", "eq", "query") | first).content }}
5
+document: {{ (messages | selectattr("role", "eq", "document") | first).content }}
6
+You are a search relevance expert who evaluates how well documents match search queries. For each query-document pair, carefully analyze the semantic relationship between them, then provide your binary relevance judgment (0 for not relevant, 1 for relevant).
7
+Relevance:<|im_end|>
8
+<|im_start|>assistant
examples/pooling/score/template/qwen3_reranker.jinja
new
+11
@@ -0,0 +1,11 @@
1
+<|im_start|>system
2
+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".<|im_end|>
3
+<|im_start|>user
4
+<Instruct>: {{ messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query") }}
5
+<Query>: {{ messages | selectattr("role", "eq", "query") | map(attribute="content") | first }}
6
+<Document>: {{ messages | selectattr("role", "eq", "document") | map(attribute="content") | first }}<|im_end|>
7
+<|im_start|>assistant
8
+<think>
9
+
10
+</think>
11
+
examples/pooling/score/using_template_offline.py
new
+159
@@ -0,0 +1,159 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+from argparse import Namespace
5
+from pathlib import Path
6
+from typing import Any
7
+
8
+from vllm import LLM, EngineArgs
9
+from vllm.utils.argparse_utils import FlexibleArgumentParser
10
+
11
+
12
+def parse_args():
13
+ """Parse command line arguments for the reranking example.
14
+
15
+ This function sets up the argument parser with default values
16
+ specific to reranking models, including the model name and
17
+ runner type.
18
+ """
19
+ parser = FlexibleArgumentParser()
20
+ # Add all EngineArgs command line arguments to the parser
21
+ parser = EngineArgs.add_cli_args(parser)
22
+
23
+ # Set default values specific to this reranking example
24
+ # These defaults ensure the script works out-of-the-box for reranking tasks
25
+ parser.set_defaults(
26
+ model="nvidia/llama-nemotron-rerank-1b-v2", # Default reranking model
27
+ runner="pooling", # Required for cross-encoder/reranking models
28
+ trust_remote_code=True, # Allow loading models with custom code
29
+ )
30
+ return parser.parse_args()
31
+
32
+
33
+def get_chat_template(model: str) -> str:
34
+ """Load the appropriate chat template for the specified model.
35
+
36
+ Reranking models require specific prompt templates to format
37
+ query-document pairs correctly. This function maps model names
38
+ to their corresponding template files.
39
+ """
40
+ # Directory containing all chat template files
41
+ template_home = Path(__file__).parent / "template"
42
+
43
+ # Mapping from model names to their corresponding template files
44
+ # Each reranking model has its own specific prompt format
45
+ model_name_to_template_path_map = {
46
+ "BAAI/bge-reranker-v2-gemma": "bge-reranker-v2-gemma.jinja",
47
+ "Qwen/Qwen3-Reranker-0.6B": "qwen3_reranker.jinja",
48
+ "Qwen/Qwen3-Reranker-4B": "qwen3_reranker.jinja",
49
+ "Qwen/Qwen3-Reranker-8B": "qwen3_reranker.jinja",
50
+ "tomaarsen/Qwen3-Reranker-0.6B-seq-cls": "qwen3_reranker.jinja",
51
+ "tomaarsen/Qwen3-Reranker-4B-seq-cls": "qwen3_reranker.jinja",
52
+ "tomaarsen/Qwen3-Reranker-8B-seq-cls": "qwen3_reranker.jinja",
53
+ "mixedbread-ai/mxbai-rerank-base-v2": "mxbai_rerank_v2.jinja",
54
+ "mixedbread-ai/mxbai-rerank-large-v2": "mxbai_rerank_v2.jinja",
55
+ "nvidia/llama-nemotron-rerank-1b-v2": "nemotron-rerank.jinja",
56
+ }
57
+
58
+ # Get the template filename for the specified model
59
+ template_path = model_name_to_template_path_map.get(model)
60
+
61
+ if template_path is None:
62
+ raise ValueError(f"This demo does not support model name: {model}.")
63
+
64
+ # Read and return the template content
65
+ return (template_home / template_path).read_text()
66
+
67
+
68
+def get_hf_overrides(model: str) -> dict[str, Any]:
69
+ """Convert Large Language Models (LLMs) to Sequence Classification models.
70
+
71
+ note:
72
+ Some reranking models require special configuration overrides to work
73
+ correctly with vLLM's score API.
74
+ Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_offline.py
75
+ Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
76
+ """
77
+
78
+ model_name_to_hf_overrides_map = {
79
+ "BAAI/bge-reranker-v2-gemma": {
80
+ "architectures": ["GemmaForSequenceClassification"],
81
+ "classifier_from_token": ["Yes"],
82
+ "method": "no_post_processing",
83
+ },
84
+ "Qwen/Qwen3-Reranker-0.6B": {
85
+ "architectures": ["Qwen3ForSequenceClassification"],
86
+ "classifier_from_token": ["no", "yes"],
87
+ "is_original_qwen3_reranker": True,
88
+ },
89
+ "Qwen/Qwen3-Reranker-4B": {
90
+ "architectures": ["Qwen3ForSequenceClassification"],
91
+ "classifier_from_token": ["no", "yes"],
92
+ "is_original_qwen3_reranker": True,
93
+ },
94
+ "Qwen/Qwen3-Reranker-8B": {
95
+ "architectures": ["Qwen3ForSequenceClassification"],
96
+ "classifier_from_token": ["no", "yes"],
97
+ "is_original_qwen3_reranker": True,
98
+ },
99
+ "tomaarsen/Qwen3-Reranker-0.6B-seq-cls": {},
100
+ "tomaarsen/Qwen3-Reranker-4B-seq-cls": {},
101
+ "tomaarsen/Qwen3-Reranker-8B-seq-cls": {},
102
+ "mixedbread-ai/mxbai-rerank-base-v2": {
103
+ "architectures": ["Qwen2ForSequenceClassification"],
104
+ "classifier_from_token": ["0", "1"],
105
+ "method": "from_2_way_softmax",
106
+ },
107
+ "mixedbread-ai/mxbai-rerank-large-v2": {
108
+ "architectures": ["Qwen2ForSequenceClassification"],
109
+ "classifier_from_token": ["0", "1"],
110
+ "method": "from_2_way_softmax",
111
+ },
112
+ "nvidia/llama-nemotron-rerank-1b-v2": {},
113
+ }
114
+
115
+ hf_overrides = model_name_to_hf_overrides_map.get(model)
116
+
117
+ if hf_overrides is None:
118
+ raise ValueError(f"This demo does not support model name: {model}.")
119
+
120
+ return hf_overrides
121
+
122
+
123
+def main(args: Namespace):
124
+ """Main execution function for the reranking example."""
125
+
126
+ # Get the overrides for the specified model
127
+ args.hf_overrides = get_hf_overrides(args.model)
128
+
129
+ # Initialize the LLM with all provided arguments
130
+ llm = LLM(**vars(args))
131
+
132
+ # Example query for demonstration
133
+ query = "how much protein should a female eat?"
134
+
135
+ # Example documents to be reranked based on relevance to the query
136
+ documents = [
137
+ "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
138
+ "Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
139
+ "Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
140
+ ]
141
+
142
+ # Load the appropriate chat template for the selected model
143
+ # The template formats query-document pairs for the reranking model
144
+ chat_template = get_chat_template(args.model)
145
+
146
+ # Score documents based on relevance to the query
147
+ # The score method returns relevance scores for each document
148
+ outputs = llm.score(query, documents, chat_template=chat_template)
149
+
150
+ # Display the relevance scores
151
+ # Higher scores indicate more relevant documents
152
+ print("-" * 30)
153
+ print([output.outputs.score for output in outputs])
154
+ print("-" * 30)
155
+
156
+
157
+if __name__ == "__main__":
158
+ args = parse_args()
159
+ main(args)
examples/pooling/score/using_template_online.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+"""
5
+Example of using the rerank API with template.
6
+
7
+This script demonstrates how to interact with a vLLM server running
8
+a reranking model via the REST API.
9
+Before running this script, start the vLLM server with one of the
10
+supported reranking models using the commands below.
11
+
12
+note:
13
+ Some reranking models require special configuration overrides to work correctly
14
+ with vLLM's score API.
15
+ Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/qwen3_reranker_online.py
16
+ Reference: https://github.com/vllm-project/vllm/blob/main/examples/pooling/score/convert_model_to_seq_cls.py
17
+
18
+run:
19
+ vllm serve BAAI/bge-reranker-v2-gemma --hf_overrides '{"architectures": ["GemmaForSequenceClassification"],"classifier_from_token": ["Yes"],"method": "no_post_processing"}' --chat-template examples/pooling/score/template/bge-reranker-v2-gemma.jinja
20
+ vllm serve tomaarsen/Qwen3-Reranker-0.6B-seq-cls --chat-template examples/pooling/score/template/qwen3_reranker.jinja
21
+ 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
22
+ vllm serve nvidia/llama-nemotron-rerank-1b-v2 --runner pooling --trust-remote-code --chat-template examples/pooling/score/template/nemotron-rerank.jinja
23
+ 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
24
+"""
25
+
26
+import json
27
+
28
+import requests
29
+
30
+# URL of the vLLM server's rerank endpoint
31
+# Default vLLM server runs on localhost port 8000
32
+url = "http://127.0.0.1:8000/rerank"
33
+
34
+# HTTP headers for the request
35
+headers = {"accept": "application/json", "Content-Type": "application/json"}
36
+
37
+# Example query & documents
38
+query = "how much protein should a female eat?"
39
+documents = [
40
+ "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
41
+ "Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments.",
42
+ "Calorie intake should not fall below 1,200 a day in women or 1,500 a day in men, except under the supervision of a health professional.",
43
+]
44
+
45
+# Request payload for the rerank API
46
+data = {
47
+ "model": "nvidia/llama-nemotron-rerank-1b-v2", # Model to use for reranking
48
+ "query": query, # The query to score documents against
49
+ "documents": documents, # List of documents to be scored
50
+}
51
+
52
+
53
+def main():
54
+ """Main function to send a rerank request to the vLLM server.
55
+
56
+ This function sends a POST request to the /rerank endpoint with
57
+ the query and documents, then prints the relevance scores.
58
+ """
59
+ # Send POST request to the vLLM server's rerank endpoint
60
+ response = requests.post(url, headers=headers, json=data)
61
+
62
+ # Check if the request was successful
63
+ if response.status_code == 200:
64
+ print("Request successful!")
65
+ # Pretty print the JSON response containing relevance scores
66
+ # The response includes scores for each document's relevance to the query
67
+ print(json.dumps(response.json(), indent=2))
68
+ else:
69
+ # Handle request failure
70
+ print(f"Request failed with status code: {response.status_code}")
71
+ print(response.text)
72
+
73
+
74
+if __name__ == "__main__":
75
+ main()
tests/conftest.py
+6
-2
@@ -45,7 +45,11 @@ from transformers import (
45
)
46
from transformers.models.auto.auto_factory import _BaseAutoModelClass
47
48
-from tests.models.utils import TokensTextLogprobs, TokensTextLogprobsPromptLogprobs
48
+from tests.models.utils import (
49
+ TokensTextLogprobs,
50
+ TokensTextLogprobsPromptLogprobs,
51
+ softmax,
52
+)
53
from vllm import LLM, SamplingParams, envs
54
from vllm.assets.audio import AudioAsset
55
from vllm.assets.image import ImageAsset
@@ -513,7 +517,7 @@ class HfRunner:
517
elif problem_type == "multi_label_classification":
518
logits = output.logits.sigmoid()[0].tolist()
519
else:
516
- logits = output.logits.softmax(dim=-1)[0].tolist()
520
+ logits = softmax(output.logits)[0].tolist()
521
outputs.append(logits)
522
523
return outputs
tests/models/language/pooling_mteb_test/mteb_score_utils.py
+69
-26
@@ -3,13 +3,16 @@
3
4
import tempfile
5
from pathlib import Path
6
+from typing import Any
7
8
import mteb
9
import numpy as np
10
import requests
11
+import torch
12
from mteb.models import ModelMeta
13
from torch.utils.data import DataLoader
14
15
+from tests.conftest import HfRunner
16
from tests.models.utils import (
17
RerankModelInfo,
18
get_vllm_extra_kwargs,
@@ -67,6 +70,12 @@ class VllmMtebCrossEncoder(MtebCrossEncoderMixin):
70
queries = [text for batch in inputs1 for text in batch["text"]]
71
corpus = [text for batch in inputs2 for text in batch["text"]]
72
73
+ # Hoping to discover potential scheduling
74
+ # issues by randomizing the order.
75
+ r = self.rng.permutation(len(queries))
76
+ queries = [queries[i] for i in r]
77
+ corpus = [corpus[i] for i in r]
78
+
79
outputs = self.llm.score(
80
queries,
81
corpus,
@@ -75,6 +84,7 @@ class VllmMtebCrossEncoder(MtebCrossEncoderMixin):
84
chat_template=self.chat_template,
85
)
86
scores = np.array(outputs)
87
+ scores = scores[np.argsort(r)]
88
return scores
89
90
@@ -84,7 +94,6 @@ class ScoreClientMtebEncoder(MtebCrossEncoderMixin):
94
def __init__(self, model_name: str, url):
95
self.model_name = model_name
96
self.url = url
87
- self.rng = np.random.default_rng(seed=42)
97
98
def predict(
99
self,
@@ -130,6 +139,50 @@ class RerankClientMtebEncoder(ScoreClientMtebEncoder):
139
return response["results"][0]["relevance_score"]
140
141
142
+class HFMtebCrossEncoder(MtebCrossEncoderMixin, HfRunner):
143
+ chat_template: str | None = None
144
+
145
+ def __init__(self, model_name: str, dtype: str = "auto", **kwargs: Any) -> None:
146
+ HfRunner.__init__(
147
+ self, model_name=model_name, is_cross_encoder=True, dtype=dtype, **kwargs
148
+ )
149
+
150
+ @torch.no_grad
151
+ def predict(
152
+ self,
153
+ inputs1: DataLoader[mteb.types.BatchedInput],
154
+ inputs2: DataLoader[mteb.types.BatchedInput],
155
+ *args,
156
+ **kwargs,
157
+ ) -> np.ndarray:
158
+ queries = [text for batch in inputs1 for text in batch["text"]]
159
+ corpus = [text for batch in inputs2 for text in batch["text"]]
160
+
161
+ if self.chat_template is not None:
162
+ tokenizer = self.model.tokenizer
163
+ prompts = []
164
+ for query, document in zip(queries, corpus):
165
+ conversation = [
166
+ {"role": "query", "content": query},
167
+ {"role": "document", "content": document},
168
+ ]
169
+
170
+ prompt = tokenizer.apply_chat_template(
171
+ conversation=conversation,
172
+ tools=None,
173
+ chat_template=self.chat_template,
174
+ tokenize=False,
175
+ )
176
+ prompts.append(prompt)
177
+ outputs_list = HfRunner.classify(self, prompts)
178
+ scores = np.array(outputs_list).squeeze(-1)
179
+ return scores
180
+ else:
181
+ prompts = list(zip(queries, corpus))
182
+ outputs_tensor = HfRunner.predict(self, prompts, show_progress_bar=False)
183
+ return outputs_tensor.cpu().numpy()
184
+
185
+
186
def run_mteb_rerank(cross_encoder: mteb.CrossEncoderProtocol, tasks, languages):
187
with tempfile.TemporaryDirectory() as prediction_folder:
188
bm25s = mteb.get_model("bm25s")
@@ -168,31 +221,21 @@ def run_mteb_rerank(cross_encoder: mteb.CrossEncoderProtocol, tasks, languages):
221
return main_score
222
223
171
-def mteb_test_rerank_models_hf(
172
- hf_runner, model_name, hf_dtype="float32", hf_model_callback=None
173
-):
174
- with hf_runner(model_name, is_cross_encoder=True, dtype=hf_dtype) as hf_model:
175
- if hf_model_callback is not None:
176
- hf_model_callback(hf_model)
177
-
178
- st_main_score = run_mteb_rerank(
179
- hf_model, tasks=MTEB_RERANK_TASKS, languages=MTEB_RERANK_LANGS
180
- )
181
- st_dtype = next(hf_model.model.model.parameters()).dtype
182
- return st_main_score, st_dtype
183
-
184
-
224
def mteb_test_rerank_models(
186
- hf_runner,
225
vllm_runner,
226
model_info: RerankModelInfo,
227
+ hf_runner=HFMtebCrossEncoder,
228
vllm_extra_kwargs=None,
190
- hf_model_callback=None,
229
vllm_mteb_encoder=VllmMtebCrossEncoder,
230
atol=MTEB_RERANK_TOL,
231
):
232
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
233
234
+ # Maybe load chat_template.
235
+ chat_template: str | None = None
236
+ if model_info.chat_template_name is not None:
237
+ chat_template = (template_home / model_info.chat_template_name).read_text()
238
+
239
with vllm_runner(
240
model_info.name,
241
runner="pooling",
@@ -201,6 +244,7 @@ def mteb_test_rerank_models(
244
**vllm_extra_kwargs,
245
) as vllm_model:
246
model_config = vllm_model.llm.llm_engine.model_config
247
+ vllm_model.chat_template = chat_template
248
249
# Confirm whether vllm is using the correct architecture
250
if model_info.architecture:
@@ -209,12 +253,6 @@ def mteb_test_rerank_models(
253
# Score API is only enabled for num_labels == 1
254
assert model_config.hf_config.num_labels == 1
255
212
- # Maybe load chat_template.
213
- chat_template: str | None = None
214
- if model_info.chat_template_name is not None:
215
- chat_template = (template_home / model_info.chat_template_name).read_text()
216
- vllm_model.chat_template = chat_template
217
-
256
# Confirm whether the important configs in model_config are correct.
257
if model_info.pooling_type is not None:
258
assert model_config.pooler_config.pooling_type == model_info.pooling_type
@@ -242,9 +280,14 @@ def mteb_test_rerank_models(
280
# Accelerate mteb test by setting
281
# SentenceTransformers mteb score to a constant
282
if model_info.mteb_score is None:
245
- st_main_score, st_dtype = mteb_test_rerank_models_hf(
246
- hf_runner, model_info.name, model_info.hf_dtype, hf_model_callback
247
- )
283
+ with hf_runner(model_info.name, dtype=model_info.hf_dtype) as hf_model:
284
+ hf_model.chat_template = chat_template
285
+ st_main_score = run_mteb_rerank(
286
+ hf_model,
287
+ tasks=MTEB_RERANK_TASKS,
288
+ languages=MTEB_RERANK_LANGS,
289
+ )
290
+ st_dtype = next(hf_model.model.model.parameters()).dtype
291
else:
292
st_main_score = model_info.mteb_score
293
st_dtype = "Constant"
tests/models/language/pooling_mteb_test/test_baai.py
+2
-4
@@ -112,7 +112,5 @@ def test_embed_models_correctness(
112
113
114
@pytest.mark.parametrize("model_info", RERANK_MODELS)
115
-def test_rerank_models_mteb(
116
- hf_runner, vllm_runner, model_info: RerankModelInfo
117
-) -> None:
118
- mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
115
+def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
116
+ mteb_test_rerank_models(vllm_runner, model_info)
tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
+29
-40
@@ -11,40 +11,60 @@ from torch.utils.data import DataLoader
11
from tests.conftest import HfRunner
12
from tests.models.utils import RerankModelInfo
13
14
-from .mteb_score_utils import VllmMtebCrossEncoder, mteb_test_rerank_models
14
+from .mteb_score_utils import (
15
+ MtebCrossEncoderMixin,
16
+ mteb_test_rerank_models,
17
+)
18
19
RERANK_MODELS = [
20
RerankModelInfo(
21
"BAAI/bge-reranker-v2-gemma",
22
architecture="GemmaForSequenceClassification",
20
- mteb_score=0.33757,
23
hf_overrides={
24
"architectures": ["GemmaForSequenceClassification"],
25
"classifier_from_token": ["Yes"],
26
"method": "no_post_processing",
27
},
28
+ mteb_score=0.33757,
29
pooling_type="LAST",
30
attn_type="decoder",
31
is_prefix_caching_supported=True,
32
is_chunked_prefill_supported=True,
33
+ chat_template_name="bge-reranker-v2-gemma.jinja",
34
),
35
]
36
37
PROMPT = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'." # noqa: E501
38
39
36
-class GemmaRerankerHfRunner(HfRunner):
40
+class GemmaRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
41
def __init__(
42
self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any
43
) -> None:
44
from transformers import AutoModelForCausalLM, AutoTokenizer
45
42
- super().__init__(model_name, dtype, auto_cls=AutoModelForCausalLM)
46
+ HfRunner.__init__(
47
+ self,
48
+ model_name=model_name,
49
+ auto_cls=AutoModelForCausalLM,
50
+ dtype=dtype,
51
+ **kwargs,
52
+ )
53
+
54
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
55
self.yes_loc = self.tokenizer.convert_tokens_to_ids("Yes")
56
46
- @torch.no_grad()
47
- def predict(self, prompts: list[list[str]], *args, **kwargs) -> torch.Tensor:
57
+ @torch.no_grad
58
+ def predict(
59
+ self,
60
+ inputs1: DataLoader[mteb.types.BatchedInput],
61
+ inputs2: DataLoader[mteb.types.BatchedInput],
62
+ *args,
63
+ **kwargs,
64
+ ) -> np.ndarray:
65
+ queries = [text for batch in inputs1 for text in batch["text"]]
66
+ corpus = [text for batch in inputs2 for text in batch["text"]]
67
+
68
def get_inputs(pairs, tokenizer, prompt=None):
69
if prompt is None:
70
prompt = PROMPT
@@ -89,8 +109,8 @@ class GemmaRerankerHfRunner(HfRunner):
109
)
110
111
scores = []
92
- for query, doc, *_ in prompts:
93
- pairs = [(query, doc)]
112
+ for query, document in zip(queries, corpus):
113
+ pairs = [(query, document)]
114
inputs = get_inputs(pairs, self.tokenizer)
115
inputs = inputs.to(self.model.device)
116
_n_tokens = inputs["input_ids"].shape[1]
@@ -107,41 +127,10 @@ class GemmaRerankerHfRunner(HfRunner):
127
return torch.Tensor(scores)
128
129
110
-class GemmaMtebEncoder(VllmMtebCrossEncoder):
111
- def __init__(self, *args, **kwargs):
112
- super().__init__(*args, **kwargs)
113
- self.query_template = "A: {query}\n"
114
- self.document_template = "B: {doc}\n{prompt}"
115
-
116
- def predict(
117
- self,
118
- inputs1: DataLoader[mteb.types.BatchedInput],
119
- inputs2: DataLoader[mteb.types.BatchedInput],
120
- *args,
121
- **kwargs,
122
- ) -> np.ndarray:
123
- queries = [
124
- self.query_template.format(query=text)
125
- for batch in inputs1
126
- for text in batch["text"]
127
- ]
128
- corpus = [
129
- self.document_template.format(doc=text, prompt=PROMPT)
130
- for batch in inputs2
131
- for text in batch["text"]
132
- ]
133
- outputs = self.llm.score(
134
- queries, corpus, truncate_prompt_tokens=-1, use_tqdm=False
135
- )
136
- scores = np.array(outputs)
137
- return scores
138
-
139
-
130
@pytest.mark.parametrize("model_info", RERANK_MODELS)
131
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
132
mteb_test_rerank_models(
143
- GemmaRerankerHfRunner,
133
vllm_runner,
134
model_info,
146
- vllm_mteb_encoder=GemmaMtebEncoder,
135
+ hf_runner=GemmaRerankerHfRunner,
136
)
tests/models/language/pooling_mteb_test/test_cross_encoder.py
+5
-6
@@ -11,27 +11,26 @@ from .mteb_score_utils import mteb_test_rerank_models
11
RERANK_MODELS = [
12
RerankModelInfo(
13
"cross-encoder/ms-marco-TinyBERT-L-2-v2",
14
- mteb_score=0.32898,
14
architecture="BertForSequenceClassification",
15
pooling_type="CLS",
16
attn_type="encoder_only",
17
is_prefix_caching_supported=False,
18
is_chunked_prefill_supported=False,
19
+ mteb_score=0.32898,
20
),
21
RerankModelInfo(
22
"tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
23
- mteb_score=0.25736,
23
architecture="Qwen3ForSequenceClassification",
24
pooling_type="LAST",
25
attn_type="decoder",
26
is_prefix_caching_supported=True,
27
is_chunked_prefill_supported=True,
28
+ chat_template_name="qwen3_reranker.jinja",
29
+ mteb_score=0.33459,
30
),
31
]
32
33
34
@pytest.mark.parametrize("model_info", RERANK_MODELS)
34
-def test_rerank_models_mteb(
35
- hf_runner, vllm_runner, model_info: RerankModelInfo
36
-) -> None:
37
- mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
35
+def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
36
+ mteb_test_rerank_models(vllm_runner, model_info)
tests/models/language/pooling_mteb_test/test_gte.py
+2
-4
@@ -143,7 +143,5 @@ def test_embed_models_correctness(
143
144
145
@pytest.mark.parametrize("model_info", RERANK_MODELS)
146
-def test_rerank_models_mteb(
147
- hf_runner, vllm_runner, model_info: RerankModelInfo
148
-) -> None:
149
- mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
146
+def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
147
+ mteb_test_rerank_models(vllm_runner, model_info)
tests/models/language/pooling_mteb_test/test_jina.py
+2
-4
@@ -72,10 +72,8 @@ def test_embed_models_correctness(
72
73
74
@pytest.mark.parametrize("model_info", RERANK_MODELS)
75
-def test_rerank_models_mteb(
76
- hf_runner, vllm_runner, model_info: RerankModelInfo
77
-) -> None:
78
- mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
75
+def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
76
+ mteb_test_rerank_models(vllm_runner, model_info)
77
78
79
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
+45
-23
@@ -2,13 +2,16 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import Any
4
5
+import mteb
6
+import numpy as np
7
import pytest
8
import torch
9
+from torch.utils.data import DataLoader
10
11
from tests.conftest import HfRunner
12
from tests.models.utils import RerankModelInfo
13
11
-from .mteb_score_utils import mteb_test_rerank_models
14
+from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models
15
16
mxbai_rerank_hf_overrides = {
17
"architectures": ["Qwen2ForSequenceClassification"],
@@ -21,50 +24,69 @@ RERANK_MODELS = [
24
"mixedbread-ai/mxbai-rerank-base-v2",
25
architecture="Qwen2ForSequenceClassification",
26
hf_overrides=mxbai_rerank_hf_overrides,
24
- mteb_score=0.273,
27
pooling_type="LAST",
28
attn_type="decoder",
29
is_prefix_caching_supported=True,
30
is_chunked_prefill_supported=True,
31
+ chat_template_name="mxbai_rerank_v2.jinja",
32
+ mteb_score=0.33651,
33
enable_test=True,
34
),
35
RerankModelInfo(
36
"mixedbread-ai/mxbai-rerank-large-v2",
37
architecture="Qwen2ForSequenceClassification",
38
hf_overrides=mxbai_rerank_hf_overrides,
39
+ chat_template_name="mxbai_rerank_v2.jinja",
40
enable_test=False,
41
),
42
]
43
44
40
-class MxbaiRerankerHfRunner(HfRunner):
45
+class MxbaiRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
46
def __init__(
47
self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any
48
) -> None:
49
from transformers import AutoModelForCausalLM, AutoTokenizer
50
46
- super().__init__(model_name, dtype, auto_cls=AutoModelForCausalLM)
51
+ HfRunner.__init__(
52
+ self,
53
+ model_name=model_name,
54
+ auto_cls=AutoModelForCausalLM,
55
+ dtype=dtype,
56
+ **kwargs,
57
+ )
58
59
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
60
self.yes_loc = self.tokenizer.convert_tokens_to_ids("1")
61
self.no_loc = self.tokenizer.convert_tokens_to_ids("0")
62
52
- def predict(self, prompts: list[list[str]], *args, **kwargs) -> torch.Tensor:
53
- def process_inputs(pairs):
54
- inputs = self.tokenizer(
55
- pairs,
56
- padding=False,
57
- truncation="longest_first",
58
- return_attention_mask=False,
63
+ @torch.no_grad
64
+ def predict(
65
+ self,
66
+ inputs1: DataLoader[mteb.types.BatchedInput],
67
+ inputs2: DataLoader[mteb.types.BatchedInput],
68
+ *args,
69
+ **kwargs,
70
+ ) -> np.ndarray:
71
+ queries = [text for batch in inputs1 for text in batch["text"]]
72
+ corpus = [text for batch in inputs2 for text in batch["text"]]
73
+
74
+ tokenizer = self.tokenizer
75
+ prompts = []
76
+ for query, document in zip(queries, corpus):
77
+ conversation = [
78
+ {"role": "query", "content": query},
79
+ {"role": "document", "content": document},
80
+ ]
81
+
82
+ prompt = tokenizer.apply_chat_template(
83
+ conversation=conversation,
84
+ tools=None,
85
+ chat_template=self.chat_template,
86
+ tokenize=False,
87
)
60
- for i, ele in enumerate(inputs["input_ids"]):
61
- inputs["input_ids"][i] = ele
62
- inputs = self.tokenizer.pad(inputs, padding=True, return_tensors="pt")
63
- for key in inputs:
64
- inputs[key] = inputs[key].to(self.model.device)
65
- return inputs
66
-
67
- @torch.no_grad()
88
+ prompts.append(prompt)
89
+
90
def compute_logits(inputs):
91
logits = self.model(**inputs).logits[:, -1, :]
92
yes_logits = logits[:, self.yes_loc]
@@ -74,9 +96,9 @@ class MxbaiRerankerHfRunner(HfRunner):
96
return scores
97
98
scores = []
77
- for query, doc, *_ in prompts:
78
- pairs = [(query, doc)]
79
- inputs = process_inputs(pairs)
99
+ for prompt in prompts:
100
+ inputs = tokenizer([prompt], return_tensors="pt")
101
+ inputs = self.wrap_device(inputs)
102
score = compute_logits(inputs)
103
scores.append(score[0].item())
104
return torch.Tensor(scores)
@@ -84,4 +106,4 @@ class MxbaiRerankerHfRunner(HfRunner):
106
107
@pytest.mark.parametrize("model_info", RERANK_MODELS)
108
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
87
- mteb_test_rerank_models(MxbaiRerankerHfRunner, vllm_runner, model_info)
109
+ mteb_test_rerank_models(vllm_runner, model_info, hf_runner=MxbaiRerankerHfRunner)
tests/models/language/pooling_mteb_test/test_nemotron.py
+2
-4
@@ -46,7 +46,5 @@ def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -
46
47
48
@pytest.mark.parametrize("model_info", RERANK_MODELS)
49
-def test_rerank_models_mteb(
50
- hf_runner, vllm_runner, model_info: RerankModelInfo
51
-) -> None:
52
- mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
49
+def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
50
+ mteb_test_rerank_models(vllm_runner, model_info)
tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
+52
-25
@@ -1,15 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
from typing import Any
5
6
+import mteb
7
+import numpy as np
8
import pytest
9
import torch
10
+from torch.utils.data import DataLoader
11
12
from tests.conftest import HfRunner
13
from tests.models.utils import RerankModelInfo
14
from tests.utils import multi_gpu_test
15
12
-from .mteb_score_utils import mteb_test_rerank_models
16
+from .mteb_score_utils import MtebCrossEncoderMixin, mteb_test_rerank_models
17
18
qwen3_reranker_hf_overrides = {
19
"architectures": ["Qwen3ForSequenceClassification"],
@@ -21,51 +25,71 @@ RERANK_MODELS = [
25
RerankModelInfo(
26
"Qwen/Qwen3-Reranker-0.6B",
27
architecture="Qwen3ForSequenceClassification",
24
- mteb_score=0.25736,
28
hf_overrides=qwen3_reranker_hf_overrides,
29
+ chat_template_name="qwen3_reranker.jinja",
30
pooling_type="LAST",
31
attn_type="decoder",
32
is_prefix_caching_supported=True,
33
is_chunked_prefill_supported=True,
34
+ mteb_score=0.33459,
35
enable_test=True,
36
),
37
RerankModelInfo(
38
"Qwen/Qwen3-Reranker-4B",
39
architecture="Qwen3ForSequenceClassification",
40
+ chat_template_name="qwen3_reranker.jinja",
41
hf_overrides=qwen3_reranker_hf_overrides,
42
enable_test=False,
43
),
44
]
45
46
41
-class Qwen3RerankerHfRunner(HfRunner):
47
+class Qwen3RerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
48
def __init__(
49
self, model_name: str, dtype: str = "auto", *args: Any, **kwargs: Any
50
) -> None:
51
from transformers import AutoModelForCausalLM, AutoTokenizer
52
47
- super().__init__(model_name, dtype, auto_cls=AutoModelForCausalLM)
53
+ HfRunner.__init__(
54
+ self,
55
+ model_name=model_name,
56
+ auto_cls=AutoModelForCausalLM,
57
+ dtype=dtype,
58
+ **kwargs,
59
+ )
60
61
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
62
self.token_false_id = self.tokenizer.convert_tokens_to_ids("no")
63
self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
52
-
53
- def predict(self, prompts: list[list[str]], *args, **kwargs) -> torch.Tensor:
54
- def process_inputs(pairs):
55
- inputs = self.tokenizer(
56
- pairs,
57
- padding=False,
58
- truncation="longest_first",
59
- return_attention_mask=False,
64
+ self.max_length = 40960
65
+
66
+ @torch.no_grad
67
+ def predict(
68
+ self,
69
+ inputs1: DataLoader[mteb.types.BatchedInput],
70
+ inputs2: DataLoader[mteb.types.BatchedInput],
71
+ *args,
72
+ **kwargs,
73
+ ) -> np.ndarray:
74
+ queries = [text for batch in inputs1 for text in batch["text"]]
75
+ corpus = [text for batch in inputs2 for text in batch["text"]]
76
+
77
+ tokenizer = self.tokenizer
78
+ prompts = []
79
+ for query, document in zip(queries, corpus):
80
+ conversation = [
81
+ {"role": "query", "content": query},
82
+ {"role": "document", "content": document},
83
+ ]
84
+
85
+ prompt = tokenizer.apply_chat_template(
86
+ conversation=conversation,
87
+ tools=None,
88
+ chat_template=self.chat_template,
89
+ tokenize=False,
90
)
61
- for i, ele in enumerate(inputs["input_ids"]):
62
- inputs["input_ids"][i] = ele
63
- inputs = self.tokenizer.pad(inputs, padding=True, return_tensors="pt")
64
- for key in inputs:
65
- inputs[key] = inputs[key].to(self.model.device)
66
- return inputs
67
-
68
- @torch.no_grad()
91
+ prompts.append(prompt)
92
+
93
def compute_logits(inputs):
94
batch_scores = self.model(**inputs).logits[:, -1, :]
95
true_vector = batch_scores[:, self.token_true_id]
@@ -76,9 +100,9 @@ class Qwen3RerankerHfRunner(HfRunner):
100
return scores
101
102
scores = []
79
- for query, doc, *_ in prompts:
80
- pairs = [(query, doc)]
81
- inputs = process_inputs(pairs)
103
+ for prompt in prompts:
104
+ inputs = tokenizer([prompt], return_tensors="pt")
105
+ inputs = self.wrap_device(inputs)
106
score = compute_logits(inputs)
107
scores.append(score[0].item())
108
return torch.Tensor(scores)
@@ -86,7 +110,7 @@ class Qwen3RerankerHfRunner(HfRunner):
110
111
@pytest.mark.parametrize("model_info", RERANK_MODELS)
112
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
89
- mteb_test_rerank_models(Qwen3RerankerHfRunner, vllm_runner, model_info)
113
+ mteb_test_rerank_models(vllm_runner, model_info, hf_runner=Qwen3RerankerHfRunner)
114
115
116
@pytest.mark.parametrize("model_info", RERANK_MODELS)
@@ -99,5 +123,8 @@ def test_rerank_models_mteb_tp(vllm_runner, model_info: RerankModelInfo) -> None
123
}
124
125
mteb_test_rerank_models(
102
- Qwen3RerankerHfRunner, vllm_runner, model_info, vllm_extra_kwargs
126
+ vllm_runner,
127
+ model_info,
128
+ vllm_extra_kwargs=vllm_extra_kwargs,
129
+ hf_runner=Qwen3RerankerHfRunner,
130
)
vllm/model_executor/models/config.py
+1
-1
@@ -251,7 +251,7 @@ class Qwen3ForSequenceClassificationConfig(VerifyAndUpdateConfig):
251
tokens = getattr(config, "classifier_from_token", None)
252
assert tokens is not None and len(tokens) == 2, (
253
"Try loading the original Qwen3 Reranker?, see: "
254
- "https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/offline_reranker.py"
254
+ "https://github.com/vllm-project/vllm/tree/main/examples/pooling/score/qwen3_reranker_offline.py"
255
)
256
model_config.hf_config.method = "from_2_way_softmax"
257