@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 23daef548d
Fork 0

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 11, 00:00 UTC

[Frontend] Support using chat template as custom score template for reranking models (#30550)

Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>

Jakub Zakrzewski committed Dec 23, 2025 at 12:19 UTC 23daef548dd1b33ba6ecb00c8c65e69f17102d13
19 files changed +663 -46
docs/models/supported_models.md
+9 -2
@@ -490,6 +490,7 @@ These models primarily support the [`LLM.embed`](./pooling_models.md#llmembed) A
490 | `GteNewModel`<sup>C</sup> | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | |
491 | `ModernBertModel`<sup>C</sup> | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | |
492 | `NomicBertModel`<sup>C</sup> | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | |
493 +| `LlamaBidirectionalModel`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ |
494 | `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
495 | `Qwen2Model`<sup>C</sup>, `Qwen2ForCausalLM`<sup>C</sup> | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ |
496 | `Qwen3Model`<sup>C</sup>, `Qwen3ForCausalLM`<sup>C</sup> | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ |
@@ -543,8 +544,9 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
544 | `BertForSequenceClassification` | BERT-based | `cross-encoder/ms-marco-MiniLM-L-6-v2`, etc. | | |
545 | `GemmaForSequenceClassification` | Gemma-based | `BAAI/bge-reranker-v2-gemma` (see note), etc. | ✅︎ | ✅︎ |
546 | `GteNewForSequenceClassification` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-reranker-base`, etc. | | |
546 -| `Qwen2ForSequenceClassification` | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2` (see note), etc. | ✅︎ | ✅︎ |
547 -| `Qwen3ForSequenceClassification` | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B` (see note), etc. | ✅︎ | ✅︎ |
547 +| `LlamaBidirectionalForSequenceClassification`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-rerank-1b-v2` (see note), etc. | ✅︎ | ✅︎ |
548 +| `Qwen2ForSequenceClassification`<sup>C</sup> | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2` (see note), etc. | ✅︎ | ✅︎ |
549 +| `Qwen3ForSequenceClassification`<sup>C</sup> | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B` (see note), etc. | ✅︎ | ✅︎ |
550 | `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | | |
551 | `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | | |
552 | `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
@@ -562,6 +564,11 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
564 !!! note
565 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.
566
567 +!!! note
568 + `nvidia/llama-nemotron-rerank-1b-v2` require a specific prompt format to work correctly.
569 +
570 + Examples : [offline_using_template.py](../../examples/pooling/score/offline_using_template.py) [online_using_template.py](../../examples/pooling/score/online_using_template.py)
571 +
572 !!! note
573 Load the official original `mxbai-rerank-v2` by using the following command.
574
docs/serving/openai_compatible_server.md
+15
@@ -669,6 +669,21 @@ You can find the documentation for cross encoder models at [sbert.net](https://w
669
670 Code example: [examples/pooling/score/openai_cross_encoder_score.py](../../examples/pooling/score/openai_cross_encoder_score.py)
671
672 +#### Score Template
673 +
674 +Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)).
675 +
676 +Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template.
677 +
678 +Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter:
679 +
680 +- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}`
681 +- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}`
682 +
683 +This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future.
684 +
685 +Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja)
686 +
687 #### Single inference
688
689 You can pass a string to both `text_1` and `text_2`, forming a single sentence pair.
examples/pooling/score/offline_using_template.py new
+27
@@ -0,0 +1,27 @@
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 new
+46
@@ -0,0 +1,46 @@
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/template/nemotron-rerank.jinja new
+3
@@ -0,0 +1,3 @@
1 +question:{{ (messages | selectattr("role", "eq", "query") | first).content }}
2 +
3 + passage:{{ (messages | selectattr("role", "eq", "document") | first).content }}
\ No newline at end of file
tests/entrypoints/pooling/score/test_utils.py new
+352
@@ -0,0 +1,352 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from unittest.mock import patch
5 +
6 +import pytest
7 +
8 +from vllm.config import ModelConfig
9 +from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
10 +from vllm.entrypoints.score_utils import get_score_prompt
11 +from vllm.inputs import TokensPrompt
12 +from vllm.tokenizers import get_tokenizer
13 +
14 +# A cross-encoder model for testing
15 +CROSS_ENCODER_MODEL_ID = "cross-encoder/ms-marco-MiniLM-L-6-v2"
16 +
17 +
18 +def assert_prompt_tokenization_consistent(
19 + tokenizer, full_prompt, engine_prompt, add_special_tokens=True
20 +):
21 + """Verify that engine_prompt token_ids match tokenizing full_prompt."""
22 + expected_ids = tokenizer(full_prompt, add_special_tokens=add_special_tokens)[
23 + "input_ids"
24 + ]
25 + actual_ids = engine_prompt["prompt_token_ids"]
26 + assert actual_ids == expected_ids, (
27 + f"Token IDs don't match.\nExpected: {expected_ids}\nActual: {actual_ids}"
28 + )
29 +
30 +
31 +@pytest.fixture(scope="module")
32 +def cross_encoder_model_config():
33 + return ModelConfig(
34 + CROSS_ENCODER_MODEL_ID,
35 + runner="pooling",
36 + )
37 +
38 +
39 +@pytest.fixture(scope="module")
40 +def cross_encoder_tokenizer(cross_encoder_model_config):
41 + return get_tokenizer(
42 + CROSS_ENCODER_MODEL_ID,
43 + trust_remote_code=cross_encoder_model_config.trust_remote_code,
44 + )
45 +
46 +
47 +@pytest.fixture(scope="module")
48 +def llm_reranker_model_config():
49 + """Model config for LLM-as-reranker style (no pad token)."""
50 + config = ModelConfig(
51 + CROSS_ENCODER_MODEL_ID,
52 + runner="pooling",
53 + )
54 + # use_pad_token is a property that reads from hf_config,
55 + # so we set it there to override the default (True)
56 + config.hf_config.use_pad_token = False
57 + return config
58 +
59 +
60 +@pytest.fixture
61 +def tokenization_kwargs():
62 + """Common tokenization kwargs used across tests."""
63 + return {"add_special_tokens": True, "return_tensors": None}
64 +
65 +
66 +@pytest.fixture
67 +def mock_model_with_score_template():
68 + """Mock model class that supports score template and tracks post_process calls."""
69 +
70 + class MockModelWithScoreTemplate:
71 + supports_score_template = True
72 + post_process_called: list[TokensPrompt] = []
73 +
74 + @staticmethod
75 + def get_score_template(p1: str, p2: str) -> str:
76 + return f"[QUERY]{p1}[SEP][DOC]{p2}"
77 +
78 + @staticmethod
79 + def post_process_tokens(prompt: TokensPrompt) -> None:
80 + MockModelWithScoreTemplate.post_process_called.append(prompt)
81 +
82 + return MockModelWithScoreTemplate
83 +
84 +
85 +@pytest.fixture
86 +def mock_model_no_score_template():
87 + """Mock model class that does not support score template."""
88 +
89 + class MockModelNoScoreTemplate:
90 + supports_score_template = False
91 +
92 + return MockModelNoScoreTemplate
93 +
94 +
95 +class TestGetScorePrompt:
96 + """Tests for the get_score_prompt function."""
97 +
98 + def test_tokenization_kwargs_passed_through(
99 + self,
100 + llm_reranker_model_config,
101 + cross_encoder_tokenizer,
102 + ):
103 + """Test that tokenization kwargs are properly passed through."""
104 + data_1 = "Query text"
105 + data_2 = "Document text"
106 +
107 + # Test with truncation - custom kwargs for this test
108 + custom_tokenization_kwargs = {
109 + "add_special_tokens": True,
110 + "return_tensors": None,
111 + "truncation": True,
112 + "max_length": 20,
113 + }
114 +
115 + full_prompt, engine_prompt = get_score_prompt(
116 + llm_reranker_model_config,
117 + cross_encoder_tokenizer,
118 + custom_tokenization_kwargs,
119 + data_1,
120 + data_2,
121 + )
122 +
123 + assert isinstance(full_prompt, str)
124 + assert "prompt_token_ids" in engine_prompt
125 + # With max_length=20 and truncation, should not exceed this
126 + assert len(engine_prompt["prompt_token_ids"]) <= 20
127 + # Since truncation was applied, token_ids should be a prefix of full encoding
128 + full_ids = cross_encoder_tokenizer(full_prompt, add_special_tokens=True)[
129 + "input_ids"
130 + ]
131 + actual_ids = engine_prompt["prompt_token_ids"]
132 + assert full_ids[: len(actual_ids)] == actual_ids, (
133 + f"Token IDs are not a prefix of full encoding.\n"
134 + f"Full IDs: {full_ids}\n"
135 + f"Actual IDs: {actual_ids}"
136 + )
137 +
138 + def test_model_supports_score_template(
139 + self,
140 + cross_encoder_model_config,
141 + cross_encoder_tokenizer,
142 + tokenization_kwargs,
143 + mock_model_with_score_template,
144 + ):
145 + """Test when model supports score template (no score_template arg)."""
146 + with patch(
147 + "vllm.model_executor.model_loader.get_model_cls",
148 + return_value=mock_model_with_score_template,
149 + ):
150 + full_prompt, engine_prompt = get_score_prompt(
151 + cross_encoder_model_config,
152 + cross_encoder_tokenizer,
153 + tokenization_kwargs,
154 + "query text",
155 + "document text",
156 + )
157 +
158 + assert full_prompt == "[QUERY]query text[SEP][DOC]document text"
159 + assert "prompt_token_ids" in engine_prompt
160 + assert len(engine_prompt["prompt_token_ids"]) > 0
161 + assert_prompt_tokenization_consistent(
162 + cross_encoder_tokenizer, full_prompt, engine_prompt
163 + )
164 +
165 + def test_model_supports_score_template_but_custom_template_provided(
166 + self,
167 + cross_encoder_model_config,
168 + cross_encoder_tokenizer,
169 + tokenization_kwargs,
170 + mock_model_with_score_template,
171 + ):
172 + """Test when model supports score template but custom template is provided."""
173 + template = (
174 + 'TEMPLATE_USED {{ messages[0]["content"] }} {{ messages[1]["content"] }}'
175 + )
176 + with (
177 + patch(
178 + "vllm.model_executor.model_loader.get_model_cls",
179 + return_value=mock_model_with_score_template,
180 + ),
181 + ):
182 + full_prompt, engine_prompt = get_score_prompt(
183 + cross_encoder_model_config,
184 + cross_encoder_tokenizer,
185 + tokenization_kwargs,
186 + "query",
187 + "doc",
188 + score_template=template, # Providing a template
189 + )
190 +
191 + assert "prompt_token_ids" in engine_prompt
192 + assert full_prompt == "TEMPLATE_USED query doc"
193 +
194 + assert_prompt_tokenization_consistent(
195 + cross_encoder_tokenizer, full_prompt, engine_prompt
196 + )
197 +
198 + def test_not_using_default_template(
199 + self,
200 + llm_reranker_model_config,
201 + cross_encoder_tokenizer,
202 + tokenization_kwargs,
203 + mock_model_no_score_template,
204 + ):
205 + # FIXME: Models implementing SupportsScoreTemplate must use their custom
206 + # template implementation by default to preserve existing functionality.
207 + # Attempting to use tokenizer_config.json templates would most likely break
208 + # these models, as often they just inherit the template from the original LLM.
209 + # CLI --chat-template overrides are still supported.
210 + with (
211 + patch(
212 + "vllm.model_executor.model_loader.get_model_cls",
213 + return_value=mock_model_no_score_template,
214 + ),
215 + patch(
216 + "vllm.entrypoints.score_utils.apply_hf_chat_template",
217 + return_value="test querytest doc",
218 + ),
219 + ):
220 + full_prompt, engine_prompt = get_score_prompt(
221 + llm_reranker_model_config,
222 + cross_encoder_tokenizer,
223 + tokenization_kwargs,
224 + "test query",
225 + "test doc",
226 + )
227 +
228 + assert full_prompt == "test querytest doc"
229 + assert "prompt_token_ids" in engine_prompt
230 + assert_prompt_tokenization_consistent(
231 + cross_encoder_tokenizer, full_prompt, engine_prompt
232 + )
233 +
234 + def test_fallback_with_pad_token(
235 + self,
236 + cross_encoder_model_config,
237 + cross_encoder_tokenizer,
238 + tokenization_kwargs,
239 + mock_model_no_score_template,
240 + ):
241 + """Test fallback path when ChatTemplateResolutionError
242 + and use_pad_token=True."""
243 + with (
244 + patch(
245 + "vllm.model_executor.model_loader.get_model_cls",
246 + return_value=mock_model_no_score_template,
247 + ),
248 + patch(
249 + "vllm.entrypoints.score_utils.apply_hf_chat_template",
250 + side_effect=ChatTemplateResolutionError("No template"),
251 + ),
252 + ):
253 + full_prompt, engine_prompt = get_score_prompt(
254 + cross_encoder_model_config, # use_pad_token=True
255 + cross_encoder_tokenizer,
256 + tokenization_kwargs,
257 + "query",
258 + "document",
259 + )
260 +
261 + assert "prompt_token_ids" in engine_prompt
262 + # Should have token_type_ids from text_pair encoding
263 + assert "token_type_ids" in engine_prompt
264 + assert "query" in full_prompt
265 + assert "document" in full_prompt
266 + assert full_prompt != "querydocument"
267 + assert (
268 + engine_prompt["prompt_token_ids"]
269 + == cross_encoder_tokenizer(
270 + "query", text_pair="document", add_special_tokens=True
271 + )["input_ids"]
272 + )
273 +
274 + # FIXME(?): add_special_tokens=False is needed because in this case
275 + # full_prompt is obtained by decoding the tokenized prompt, which includes
276 + # special tokens and we would get duplicated special tokens otherwise.
277 + # This is inconsistent with other cases.
278 + assert_prompt_tokenization_consistent(
279 + cross_encoder_tokenizer,
280 + full_prompt,
281 + engine_prompt,
282 + add_special_tokens=False,
283 + )
284 +
285 + def test_fallback_without_pad_token(
286 + self,
287 + llm_reranker_model_config,
288 + cross_encoder_tokenizer,
289 + tokenization_kwargs,
290 + mock_model_no_score_template,
291 + ):
292 + """Test fallback path when ChatTemplateResolutionError
293 + and use_pad_token=False."""
294 + with (
295 + patch(
296 + "vllm.model_executor.model_loader.get_model_cls",
297 + return_value=mock_model_no_score_template,
298 + ),
299 + patch(
300 + "vllm.entrypoints.score_utils.apply_hf_chat_template",
301 + side_effect=ChatTemplateResolutionError("No template"),
302 + ),
303 + ):
304 + full_prompt, engine_prompt = get_score_prompt(
305 + llm_reranker_model_config, # use_pad_token=False
306 + cross_encoder_tokenizer,
307 + tokenization_kwargs,
308 + "query",
309 + "document",
310 + )
311 +
312 + assert full_prompt == "querydocument"
313 + assert "prompt_token_ids" in engine_prompt
314 + assert_prompt_tokenization_consistent(
315 + cross_encoder_tokenizer, full_prompt, engine_prompt
316 + )
317 +
318 + def test_post_process_tokens_called(
319 + self,
320 + cross_encoder_model_config,
321 + cross_encoder_tokenizer,
322 + tokenization_kwargs,
323 + mock_model_with_score_template,
324 + ):
325 + """Test that post_process_tokens is called on the engine prompt."""
326 + # Reset the call tracker
327 + mock_model_with_score_template.post_process_called.clear()
328 +
329 + with (
330 + patch(
331 + "vllm.model_executor.model_loader.get_model_cls",
332 + return_value=mock_model_with_score_template,
333 + ),
334 + patch(
335 + "vllm.entrypoints.score_utils.apply_hf_chat_template",
336 + side_effect=ChatTemplateResolutionError("No template"),
337 + ),
338 + ):
339 + full_prompt, engine_prompt = get_score_prompt(
340 + cross_encoder_model_config,
341 + cross_encoder_tokenizer,
342 + tokenization_kwargs,
343 + "query",
344 + "doc",
345 + )
346 +
347 + # post_process_tokens should have been called once
348 + assert len(mock_model_with_score_template.post_process_called) == 1
349 + assert mock_model_with_score_template.post_process_called[0] is engine_prompt
350 + assert_prompt_tokenization_consistent(
351 + cross_encoder_tokenizer, full_prompt, engine_prompt
352 + )
tests/models/language/pooling_mteb_test/mteb_utils.py
+40 -24
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import tempfile
5 +from pathlib import Path
6
7 import mteb
8 import numpy as np
@@ -19,6 +20,11 @@ from tests.models.utils import (
20 get_vllm_extra_kwargs,
21 )
22
23 +template_home = (
24 + Path(__file__).parent.parent.parent.parent.parent
25 + / "examples/pooling/score/template"
26 +)
27 +
28 # Most embedding models on the STS12 task (See #17175):
29 # - Model implementation and minor changes in tensor dtype
30 # results in differences less than 1e-4
@@ -102,30 +108,6 @@ class VllmMtebEncoder(mteb.EncoderProtocol):
108 return sim
109
110
105 -class VllmMtebCrossEncoder(mteb.CrossEncoderProtocol):
106 - mteb_model_meta = _empty_model_meta
107 -
108 - def __init__(self, vllm_model):
109 - self.llm = vllm_model
110 - self.rng = np.random.default_rng(seed=42)
111 -
112 - def predict(
113 - self,
114 - inputs1: DataLoader[mteb.types.BatchedInput],
115 - inputs2: DataLoader[mteb.types.BatchedInput],
116 - *args,
117 - **kwargs,
118 - ) -> np.ndarray:
119 - queries = [text for batch in inputs1 for text in batch["text"]]
120 - corpus = [text for batch in inputs2 for text in batch["text"]]
121 -
122 - outputs = self.llm.score(
123 - queries, corpus, truncate_prompt_tokens=-1, use_tqdm=False
124 - )
125 - scores = np.array(outputs)
126 - return scores
127 -
128 -
111 class OpenAIClientMtebEncoder(VllmMtebEncoder):
112 def __init__(self, model_name: str, client):
113 self.model_name = model_name
@@ -153,6 +135,35 @@ class OpenAIClientMtebEncoder(VllmMtebEncoder):
135 return embeds
136
137
138 +class VllmMtebCrossEncoder(mteb.CrossEncoderProtocol):
139 + mteb_model_meta = _empty_model_meta
140 +
141 + def __init__(self, vllm_model):
142 + self.llm = vllm_model
143 + self.rng = np.random.default_rng(seed=42)
144 + self.chat_template: str | None = getattr(vllm_model, "chat_template", None)
145 +
146 + def predict(
147 + self,
148 + inputs1: DataLoader[mteb.types.BatchedInput],
149 + inputs2: DataLoader[mteb.types.BatchedInput],
150 + *args,
151 + **kwargs,
152 + ) -> np.ndarray:
153 + queries = [text for batch in inputs1 for text in batch["text"]]
154 + corpus = [text for batch in inputs2 for text in batch["text"]]
155 +
156 + outputs = self.llm.score(
157 + queries,
158 + corpus,
159 + truncate_prompt_tokens=-1,
160 + use_tqdm=False,
161 + chat_template=self.chat_template,
162 + )
163 + scores = np.array(outputs)
164 + return scores
165 +
166 +
167 class ScoreClientMtebEncoder(mteb.CrossEncoderProtocol):
168 mteb_model_meta = _empty_model_meta
169
@@ -387,6 +398,11 @@ def mteb_test_rerank_models(
398 == model_info.default_pooling_type
399 )
400
401 + chat_template: str | None = None
402 + if model_info.chat_template_name is not None:
403 + chat_template = (template_home / model_info.chat_template_name).read_text()
404 + vllm_model.chat_template = chat_template
405 +
406 vllm_main_score = run_mteb_rerank(
407 vllm_mteb_encoder(vllm_model),
408 tasks=MTEB_RERANK_TASKS,
tests/models/language/pooling_mteb_test/test_nemotron.py new
+42
@@ -0,0 +1,42 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from tests.models.utils import (
7 + EmbedModelInfo,
8 + LASTPoolingEmbedModelInfo,
9 + LASTPoolingRerankModelInfo,
10 + RerankModelInfo,
11 +)
12 +
13 +from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
14 +
15 +EMBEDDING_MODELS = [
16 + LASTPoolingEmbedModelInfo(
17 + "nvidia/llama-nemotron-embed-1b-v2",
18 + architecture="LlamaBidirectionalModel",
19 + mteb_score=0.689164662128673,
20 + )
21 +]
22 +
23 +RERANK_MODELS = [
24 + LASTPoolingRerankModelInfo(
25 + "nvidia/llama-nemotron-rerank-1b-v2",
26 + architecture="LlamaBidirectionalForSequenceClassification",
27 + chat_template_name="nemotron-rerank.jinja",
28 + mteb_score=0.33994,
29 + ),
30 +]
31 +
32 +
33 +@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
34 +def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
35 + mteb_test_embed_models(hf_runner, vllm_runner, model_info)
36 +
37 +
38 +@pytest.mark.parametrize("model_info", RERANK_MODELS)
39 +def test_rerank_models_mteb(
40 + hf_runner, vllm_runner, model_info: RerankModelInfo
41 +) -> None:
42 + mteb_test_rerank_models(hf_runner, vllm_runner, model_info)
tests/models/registry.py
+6
@@ -488,6 +488,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
488 ),
489 "JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"),
490 "LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
491 + "LlamaBidirectionalModel": _HfExamplesInfo(
492 + "nvidia/llama-nemotron-embed-1b-v2", trust_remote_code=True
493 + ),
494 "MistralModel": _HfExamplesInfo("intfloat/e5-mistral-7b-instruct"),
495 "ModernBertModel": _HfExamplesInfo(
496 "Alibaba-NLP/gte-modernbert-base", trust_remote_code=True
@@ -554,6 +557,9 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
557 trust_remote_code=True,
558 hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
559 ),
560 + "LlamaBidirectionalForSequenceClassification": _HfExamplesInfo(
561 + "nvidia/llama-nemotron-rerank-1b-v2", trust_remote_code=True
562 + ),
563 "ModernBertForSequenceClassification": _HfExamplesInfo(
564 "Alibaba-NLP/gte-reranker-modernbert-base"
565 ),
tests/models/utils.py
+1
@@ -399,6 +399,7 @@ class LASTPoolingEmbedModelInfo(EmbedModelInfo):
399 @dataclass
400 class RerankModelInfo(ModelInfo):
401 mteb_score: float | None = None
402 + chat_template_name: str | None = None
403
404
405 @dataclass
vllm/entrypoints/chat_utils.py
+10 -1
@@ -67,6 +67,15 @@ else:
67
68 logger = init_logger(__name__)
69
70 +
71 +class ChatTemplateResolutionError(ValueError):
72 + """Raised when chat template resolution fails.
73 +
74 + This is a subclass of ValueError for backward compatibility with
75 + existing exception handlers.
76 + """
77 +
78 +
79 MODALITY_PLACEHOLDERS_MAP = {
80 "image": "<##IMAGE##>",
81 "audio": "<##AUDIO##>",
@@ -1814,7 +1823,7 @@ def apply_hf_chat_template(
1823 )
1824
1825 if hf_chat_template is None:
1817 - raise ValueError(
1826 + raise ChatTemplateResolutionError(
1827 "As of transformers v4.44, default chat template is no longer "
1828 "allowed, so you must provide a chat template if the tokenizer "
1829 "does not define one."
vllm/entrypoints/llm.py
+11
@@ -1280,6 +1280,7 @@ class LLM:
1280 pooling_params: PoolingParams | None = None,
1281 lora_request: list[LoRARequest] | LoRARequest | None = None,
1282 tokenization_kwargs: dict[str, Any] | None = None,
1283 + score_template: str | None = None,
1284 ) -> list[ScoringRequestOutput]:
1285 model_config = self.model_config
1286
@@ -1313,6 +1314,7 @@ class LLM:
1314 data_2=d,
1315 tokenizer=tokenizer,
1316 tokenization_kwargs=tokenization_kwargs,
1317 + score_template=score_template,
1318 )
1319
1320 if token_type_ids := engine_prompt.pop("token_type_ids", None):
@@ -1347,6 +1349,7 @@ class LLM:
1349 use_tqdm: bool | Callable[..., tqdm] = True,
1350 pooling_params: PoolingParams | None = None,
1351 lora_request: list[LoRARequest] | LoRARequest | None = None,
1352 + chat_template: str | None = None,
1353 ) -> list[ScoringRequestOutput]:
1354 """Generate similarity scores for all pairs `<text,text_pair>` or
1355 `<multi-modal data, multi-modal data pair>`.
@@ -1379,6 +1382,8 @@ class LLM:
1382 lora_request: LoRA request to use for generation, if any.
1383 pooling_params: The pooling parameters for pooling. If None, we
1384 use the default pooling parameters.
1385 + chat_template: The chat template to use for the scoring. If None, we
1386 + use the model's default chat template.
1387 Returns:
1388 A list of `ScoringRequestOutput` objects containing the
1389 generated scores in the same order as the input prompts.
@@ -1406,6 +1411,11 @@ class LLM:
1411 ):
1412 raise ValueError("Score API is only enabled for num_labels == 1.")
1413
1414 + if not model_config.is_cross_encoder and chat_template is not None:
1415 + raise ValueError(
1416 + "chat_template is only supported for cross-encoder models."
1417 + )
1418 +
1419 # the tokenizer for models such as
1420 # "cross-encoder/ms-marco-MiniLM-L-6-v2" doesn't support passing
1421 # lists of tokens to the `text` and `text_pair` kwargs
@@ -1475,6 +1485,7 @@ class LLM:
1485 use_tqdm,
1486 pooling_params,
1487 lora_request,
1488 + score_template=chat_template,
1489 )
1490 else:
1491 return self._embedding_score(
vllm/entrypoints/openai/api_server.py
+1
@@ -1145,6 +1145,7 @@ async def init_app_state(
1145 engine_client,
1146 state.openai_serving_models,
1147 request_logger=request_logger,
1148 + score_template=resolved_chat_template,
1149 log_error_stack=args.log_error_stack,
1150 )
1151 if ("embed" in supported_tasks or "score" in supported_tasks)
vllm/entrypoints/openai/run_batch.py
+1
@@ -495,6 +495,7 @@ async def run_batch(
495 engine_client,
496 openai_serving_models,
497 request_logger=request_logger,
498 + score_template=None,
499 )
500 if ("embed" in supported_tasks or enable_serving_reranking)
501 else None
vllm/entrypoints/pooling/score/serving.py
+3
@@ -52,6 +52,7 @@ class ServingScores(OpenAIServing):
52 models: OpenAIServingModels,
53 *,
54 request_logger: RequestLogger | None,
55 + score_template: str | None = None,
56 log_error_stack: bool = False,
57 ) -> None:
58 super().__init__(
@@ -60,6 +61,7 @@ class ServingScores(OpenAIServing):
61 request_logger=request_logger,
62 log_error_stack=log_error_stack,
63 )
64 + self.score_template = score_template
65
66 async def _embedding_score(
67 self,
@@ -169,6 +171,7 @@ class ServingScores(OpenAIServing):
171 data_2=data_2,
172 tokenizer=tokenizer,
173 tokenization_kwargs=tokenization_kwargs,
174 + score_template=self.score_template,
175 )
176 self._validate_input(request, engine_prompt["prompt_token_ids"], full_prompt)
177 if request.mm_processor_kwargs is not None:
vllm/entrypoints/score_utils.py
+46 -16
@@ -11,9 +11,11 @@ from vllm.entrypoints.chat_utils import (
11 ChatCompletionContentPartImageEmbedsParam,
12 ChatCompletionContentPartImageParam,
13 ChatCompletionContentPartTextParam,
14 + ChatTemplateResolutionError,
15 MultiModalItemTracker,
16 _ContentPart,
17 _parse_chat_message_content_part,
18 + apply_hf_chat_template,
19 )
20 from vllm.inputs import TokensPrompt
21 from vllm.model_executor.models.interfaces import supports_score_template
@@ -139,10 +141,8 @@ def _parse_score_content(
141 return next(iter(mm_placeholder_storage.values()))[0]
142
143
142 -def apply_score_template(
143 - model_config: ModelConfig,
144 - prompt_1: str,
145 - prompt_2: str,
144 +def _apply_model_score_template(
145 + model_config: ModelConfig, prompt_1: str, prompt_2: str
146 ) -> str:
147 # NOTE(Simon): lazy import to avoid bring in all dependencies (e.g. gguf)
148 from vllm.model_executor.model_loader import get_model_cls
@@ -181,6 +181,7 @@ def get_score_prompt(
181 tokenization_kwargs: dict[str, Any],
182 data_1: str | ScoreContentPartParam,
183 data_2: str | ScoreContentPartParam,
184 + score_template: str | None = None,
185 ) -> tuple[str, TokensPrompt]:
186 prompt_1, prompt_2, mm_data = parse_score_data(
187 data_1,
@@ -190,19 +191,48 @@ def get_score_prompt(
191 from vllm.model_executor.model_loader import get_model_cls
192
193 model = get_model_cls(model_config)
193 - if supports_score_template(model):
194 - full_prompt = apply_score_template(model_config, prompt_1, prompt_2)
195 - prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
196 - elif model_config.use_pad_token:
197 - # cross_encoder models defaults to using pad_token.
198 - prompt_inputs = tokenizer(
199 - text=prompt_1, text_pair=prompt_2, **tokenization_kwargs
200 - )
201 - full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
194 +
195 + def default_tokenizer_encode():
196 + if supports_score_template(model):
197 + full_prompt = _apply_model_score_template(model_config, prompt_1, prompt_2)
198 + prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
199 + else:
200 + if model_config.use_pad_token:
201 + # cross_encoder models defaults to using pad_token.
202 + prompt_inputs = tokenizer(
203 + text=prompt_1, text_pair=prompt_2, **tokenization_kwargs
204 + )
205 + full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
206 + else:
207 + # `llm as reranker` models defaults to not using pad_token.
208 + full_prompt = prompt_1 + prompt_2
209 + prompt_inputs = tokenizer(text=full_prompt, **tokenization_kwargs)
210 + return full_prompt, prompt_inputs
211 +
212 + # FIXME: For now, we only apply a template when one is explicitly provided.
213 + # We cannot rely on the tokenizer's chat template because many models
214 + # inherit junk templates from their base LLM, which breaks both the models
215 + # and the tests that use them.
216 + if score_template is None:
217 + full_prompt, prompt_inputs = default_tokenizer_encode()
218 else:
203 - # `llm as reranker` models defaults to not using pad_token.
204 - full_prompt = prompt_1 + prompt_2
205 - prompt_inputs = tokenizer(text=full_prompt, **tokenization_kwargs)
219 + # FIXME: Try applying a score template from the CLI arg or tokenizer_config.json
220 + # If that fails because there is no such template,
221 + # fall back to the default implementation.
222 + try:
223 + full_prompt = apply_hf_chat_template(
224 + tokenizer,
225 + [
226 + {"role": "query", "content": prompt_1},
227 + {"role": "document", "content": prompt_2},
228 + ],
229 + score_template,
230 + tools=None,
231 + model_config=model_config,
232 + )
233 + prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
234 + except ChatTemplateResolutionError:
235 + full_prompt, prompt_inputs = default_tokenizer_encode()
236
237 engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"])
238
vllm/model_executor/models/config.py
+22
@@ -88,6 +88,26 @@ class JinaRobertaModelConfig(VerifyAndUpdateConfig):
88 }
89
90
91 +class LlamaBidirectionalConfig(VerifyAndUpdateConfig):
92 + @staticmethod
93 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
94 + from vllm.config.pooler import PoolingTypeStr
95 +
96 + hf_config = vllm_config.model_config.hf_config
97 + hf_config.is_causal = False
98 +
99 + pooling_type_map: dict[str, PoolingTypeStr] = {
100 + "avg": "MEAN",
101 + "cls": "CLS",
102 + "last": "LAST",
103 + }
104 +
105 + pooling_type = pooling_type_map.get(hf_config.pooling, None)
106 + if pooling_type is None:
107 + raise ValueError(f"pool_type {hf_config.pooling} not supported")
108 + vllm_config.model_config.pooler_config.pooling_type = pooling_type
109 +
110 +
111 class NomicBertModelConfig(VerifyAndUpdateConfig):
112 @staticmethod
113 def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -509,6 +529,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
529 "GteNewModel": GteNewModelConfig,
530 "GteNewForSequenceClassification": GteNewModelConfig,
531 "Gemma3TextModel": Gemma3TextModelConfig,
532 + "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
533 + "LlamaBidirectionalModel": LlamaBidirectionalConfig,
534 "NomicBertModel": NomicBertModelConfig,
535 "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
536 "Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
vllm/model_executor/models/llama.py
+22 -1
@@ -57,7 +57,14 @@ from vllm.model_executor.model_loader.weight_utils import (
57 )
58 from vllm.sequence import IntermediateTensors
59
60 -from .interfaces import SupportsEagle, SupportsEagle3, SupportsLoRA, SupportsPP
60 +from .adapters import as_embedding_model, as_seq_cls_model
61 +from .interfaces import (
62 + SupportsEagle,
63 + SupportsEagle3,
64 + SupportsLoRA,
65 + SupportsPP,
66 +)
67 +from .interfaces_base import attn_type
68 from .utils import (
69 AutoWeightsLoader,
70 PPMissingLayer,
@@ -698,3 +705,17 @@ class LlamaForCausalLM(
705 name = name.replace(item, mapping[item])
706
707 return name, loaded_weight
708 +
709 +
710 +@attn_type("encoder_only")
711 +class LlamaBidirectionalForSequenceClassification(as_seq_cls_model(LlamaForCausalLM)):
712 + # This class sets the correct attention type and pooling type
713 + # through LlamaBidirectionalConfig.
714 + pass
715 +
716 +
717 +@attn_type("encoder_only")
718 +class LlamaBidirectionalModel(as_embedding_model(LlamaForCausalLM)):
719 + # This class sets the correct attention type and pooling type
720 + # through LlamaBidirectionalConfig.
721 + pass
vllm/model_executor/models/registry.py
+6 -2
@@ -203,6 +203,7 @@ _EMBEDDING_MODELS = {
203 "GteNewModel": ("bert_with_rope", "GteNewModel"),
204 "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
205 "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), # noqa: E501
206 + "LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"),
207 "LlamaModel": ("llama", "LlamaForCausalLM"),
208 **{
209 # Multiple models share the same architecture, so we include them all
@@ -246,6 +247,11 @@ _CROSS_ENCODER_MODELS = {
247 "bert_with_rope",
248 "GteNewForSequenceClassification",
249 ),
250 + "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"),
251 + "LlamaBidirectionalForSequenceClassification": (
252 + "llama",
253 + "LlamaBidirectionalForSequenceClassification",
254 + ),
255 "ModernBertForSequenceClassification": (
256 "modernbert",
257 "ModernBertForSequenceClassification",
@@ -259,8 +265,6 @@ _CROSS_ENCODER_MODELS = {
265 "roberta",
266 "RobertaForSequenceClassification",
267 ),
262 - # [Auto-converted (see adapters.py)]
263 - "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"), # noqa: E501,
268 }
269
270 _MULTIMODAL_MODELS = {