From 3c13653738dad0acd877c23373816fefffbdd373 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Delhaye?= Date: Sun, 18 Oct 2026 13:30:12 +0200 Subject: [PATCH] [Frontend] Score with the chat template declared by Sentence Transformers CrossEncoders (#60877) Signed-off-by: rdlhio --- docs/models/pooling_models/scoring.md | 25 ++++++ docs/models/pooling_models/specific_models.md | 8 +- .../pooling/scoring/test_io_processor_unit.py | 2 + .../pooling_mteb_test/test_nemotron.py | 8 ++ .../pooling/test_llama_nemotron_vl.py | 26 ++---- .../pooling/scoring/io_processor.py | 64 +++++++++++++- vllm/transformers_utils/config.py | 86 +++++++++++++++++++ 7 files changed, 192 insertions(+), 27 deletions(-) diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index 7ff186b5cd..3bda12a09c 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -466,6 +466,31 @@ This approach is more robust than index-based access (`messages[0]`, `messages[1 Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../../examples/pooling/score/template/nemotron-rerank.jinja) +#### Templates declared by Sentence Transformers checkpoints + +vLLM does not apply a checkpoint's chat template by default, since many rerankers inherit an unrelated template from their base LLM. The exception is a Sentence Transformers (v5.4+) `CrossEncoder` with a single `sequence-classification` Transformer module whose `sentence_bert_config.json` declares the `message` modality. Sentence Transformers formats the pairs of such checkpoints as `query`/`document` messages, so vLLM does the same with the saved template: + +```json +{ + "transformer_task": "sequence-classification", + "modality_config": { + "text": {"method": "forward", "method_output_name": "logits"}, + "message": {"method": "forward", "method_output_name": "logits", "format": "flat"} + }, + "module_output_name": "scores", + "processing_kwargs": {"chat_template": {"add_generation_prompt": true}} +} +``` + +- The default chat template is used, unless `processing_kwargs.chat_template.chat_template` names another template saved with the tokenizer (e.g. `additional_chat_templates/score.jinja` for `"score"`). +- The other `processing_kwargs.chat_template` settings, such as `add_generation_prompt`, are passed to the template; the request's `chat_template_kwargs` override them. +- Message contents are passed as strings, as with `--chat-template`. +- vLLM does not restore the template suffix after truncation (`restore_suffix`). Use `max_tokens_per_doc` rather than `truncate_prompt_tokens` to truncate long documents. +- `--chat-template` (or the `chat_template` argument of `LLM.score`) takes precedence over the saved template and its settings. +- Models with a built-in score template keep using it. + +The template in use is logged at startup. + ### Enable/disable activation You can enable or disable activation via `use_activation` only works for cross-encoder models. diff --git a/docs/models/pooling_models/specific_models.md b/docs/models/pooling_models/specific_models.md index 0c7b0a1bf3..5b2871ca92 100644 --- a/docs/models/pooling_models/specific_models.md +++ b/docs/models/pooling_models/specific_models.md @@ -313,14 +313,12 @@ Start the server: ```shell vllm serve nvidia/llama-nemotron-rerank-vl-1b-v2 \ --runner pooling \ - --trust-remote-code \ - --chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja + --trust-remote-code ``` !!! note - The chat template bundled with this checkpoint's tokenizer is not suitable - for the Score/Rerank APIs. Use the provided override template when serving: - `examples/pooling/score/template/nemotron-vl-rerank.jinja`. + The checkpoint declares its score template in its Sentence Transformers config, so vLLM applies it automatically (see [Templates declared by Sentence Transformers checkpoints](scoring.md#templates-declared-by-sentence-transformers-checkpoints)). + Revisions older than October 2026 need `--chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja`. Score a text query against an image document: diff --git a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py index 7f2c157f25..d5a585f579 100644 --- a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py @@ -65,6 +65,8 @@ def llm_reranker_processor() -> CrossEncoderIOProcessor: processor.supports_score_template = False processor.use_sep_token = False processor.model = None + processor.saved_chat_template = None + processor.saved_chat_template_kwargs = {} return processor diff --git a/tests/models/language/pooling_mteb_test/test_nemotron.py b/tests/models/language/pooling_mteb_test/test_nemotron.py index 711224a1a2..3c3618de65 100644 --- a/tests/models/language/pooling_mteb_test/test_nemotron.py +++ b/tests/models/language/pooling_mteb_test/test_nemotron.py @@ -39,6 +39,14 @@ RERANK_MODELS = [ is_prefix_caching_supported=False, is_chunked_prefill_supported=False, ), + # Declares its score template in its Sentence Transformers config, so no + # chat template is passed. + RerankModelInfo( + "nvidia/llama-nemotron-rerank-vl-1b-v2", + architecture="LlamaNemotronVLForSequenceClassification", + mteb_score=0.33919, + mteb_tol=1e-2, + ), ] diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index 3cbeda2e49..227542023a 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -10,7 +10,6 @@ Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone. from collections.abc import Sequence from io import BytesIO -from pathlib import Path import pybase64 as base64 import pytest @@ -202,13 +201,6 @@ def test_models( RERANKER_MODELS = ["nvidia/llama-nemotron-rerank-vl-1b-v2"] -# The tokenizer's built-in chat template is not suitable for the Score/Rerank -# APIs (it's inherited from the base LLM). We must use the provided override. -_RERANKER_SCORE_TEMPLATE = ( - Path(__file__).parents[4] - / "examples/pooling/score/template/nemotron-vl-rerank.jinja" -).read_text() - RERANKER_TEXT_QUERY = "How is AI improving the intelligence and capabilities of robots?" RERANKER_TEXT_DOCS = [ "AI enables robots to perceive, plan, and act autonomously.", @@ -281,7 +273,11 @@ def _run_vllm_reranker( dtype: str, input_cases: Sequence[RerankerCase], ) -> list[list[float]]: - """Run all vLLM reranker cases in one model lifecycle.""" + """Run all vLLM reranker cases in one model lifecycle. + + The checkpoint declares its score template in its Sentence Transformers + config, so no chat template is passed. + """ with vllm_runner( model, runner="pooling", @@ -298,11 +294,7 @@ def _run_vllm_reranker( if not has_images: queries = [query] * len(docs) doc_texts = [doc_text for doc_text, _ in docs] - outputs = vllm_model.score( - queries, - doc_texts, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) + outputs = vllm_model.score(queries, doc_texts) else: query_params = [ ScoreMultiModalParam( @@ -334,11 +326,7 @@ def _run_vllm_reranker( ) doc_params.append(ScoreMultiModalParam(content=content)) - raw_outputs = vllm_model.llm.score( - query_params, - doc_params, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) + raw_outputs = vllm_model.llm.score(query_params, doc_params) outputs = [output.outputs.score for output in raw_outputs] scores_per_case.append(outputs) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 7c4a1fe917..2613a4073a 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -9,14 +9,17 @@ import torch.nn.functional as F from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt from vllm.exceptions import ChatTemplateResolutionError from vllm.logger import init_logger -from vllm.renderers import TokenizeParams -from vllm.renderers.hf import safe_apply_chat_template +from vllm.renderers import TokenizeParams, merge_kwargs +from vllm.renderers.hf import resolve_chat_template, safe_apply_chat_template from vllm.renderers.inputs.preprocess import ( extract_target_prompt, parse_model_prompt, prompt_to_seq, ) from vllm.tasks import PoolingTask +from vllm.transformers_utils.config import ( + get_sentence_transformers_chat_template_kwargs, +) from vllm.utils.mistral import is_mistral_tokenizer from ..base.io_processor import PoolingIOProcessor @@ -457,9 +460,22 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): self.model = model if self.supports_score_template else None self.use_sep_token = self.model_config.use_sep_token + self.saved_chat_template: str | None = None + self.saved_chat_template_kwargs: dict[str, Any] = {} + if not self.supports_score_template: + try: + self.saved_chat_template, self.saved_chat_template_kwargs = ( + self._get_saved_chat_template() + ) + except ValueError: + # An explicit chat template takes precedence anyway. + if self.chat_template is None: + raise + if ( getattr(self.model_config.hf_config, "is_original_qwen3_reranker", False) and self.chat_template is None + and self.saved_chat_template is None ): suggested_template = ( "examples/pooling/score/template/qwen3_vl_reranker.jinja" @@ -475,6 +491,41 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): suggested_template, ) + def _get_saved_chat_template(self) -> tuple[str | None, dict[str, Any]]: + """Get the chat template and kwargs declared by the Sentence + Transformers config, if any.""" + saved_kwargs = get_sentence_transformers_chat_template_kwargs( + self.model_config.model, self.model_config.revision + ) + if saved_kwargs is None: + return None, {} + + saved_kwargs = dict(saved_kwargs) + name = saved_kwargs.pop("chat_template", None) + if name == "default": + name = None + template = resolve_chat_template( + self.tokenizer, + chat_template=name, + tools=None, + model_config=self.model_config, + ) + # Hugging Face returns an unknown template name unchanged. + if template is None or template == name: + raise ValueError( + "The Sentence Transformers config declares message inputs, but " + f"the checkpoint has no {name or 'default'!r} chat template." + ) + + if self.chat_template is None: + logger.info_once( + "Scoring with the %r chat template declared by the Sentence " + "Transformers config. Its suffix is not restored after " + "truncation; use max_tokens_per_doc to truncate documents.", + name or "default", + ) + return template, saved_kwargs + ####################################### # online APIs @@ -648,6 +699,12 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): model_config = self.model_config tokenizer = self.tokenizer + if chat_template is None and self.saved_chat_template is not None: + chat_template = self.saved_chat_template + chat_template_kwargs = merge_kwargs( + self.saved_chat_template_kwargs, chat_template_kwargs + ) + prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data( data_1, data_2, @@ -703,7 +760,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): prompt_inputs = tokenizer(text=full_prompt, **local_kwargs) return full_prompt, prompt_inputs - # FIXME: For now, we only apply a template when one is explicitly provided. + # FIXME: For now, we only apply a template when one is explicitly provided, + # either directly or by the Sentence Transformers config. # We cannot rely on the tokenizer's chat template because many models # inherit junk templates from their base LLM, which breaks both the models # and the tests that use them. diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index c67d468a13..8587c02b68 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -89,6 +89,10 @@ CONFIG_FIELD_ALIASES: dict[str, tuple[str, ...]] = { `ModelArchConfigConvertorBase` reads a scalar from these; a caller that needs the raw value (a per-layer `list`, say) passes the aliases to `getattr_iter` itself.""" +_ST_TRANSFORMER_MODULE_TYPES = { + "sentence_transformers.models.Transformer", + "sentence_transformers.base.modules.transformer.Transformer", +} _ST_POOLING_MODULE_TYPES = { "sentence_transformers.models.Pooling", "sentence_transformers.sentence_transformer.modules.pooling.Pooling", @@ -1354,6 +1358,88 @@ def try_get_dense_modules( return None +@cache +def get_sentence_transformers_chat_template_kwargs( + model: str | Path, + revision: str | None = None, +) -> dict[str, Any] | None: + """Get the saved chat template settings of a Sentence Transformers + CrossEncoder whose single Transformer module declares the `message` + modality. + + Sentence Transformers formats the pairs of such checkpoints as + `query`/`document` messages with the saved chat template. + + Args: + model: Local checkpoint path or Hugging Face model ID. + revision: Model revision to inspect. + + Returns: + The saved `processing_kwargs.chat_template` (possibly empty), or + `None` if the checkpoint does not declare the `message` modality. + + Raises: + ValueError: If the saved chat template settings are unsupported. + + """ + try: + st_config = get_hf_file_to_dict( + "config_sentence_transformers.json", model, revision + ) + if not isinstance(st_config, dict) or ( + st_config.get("model_type") != "CrossEncoder" + ): + return None + modules = get_hf_file_to_dict("modules.json", model, revision) + transformer_config = get_hf_file_to_dict( + "sentence_bert_config.json", model, revision + ) + except Exception: + logger.debug("Failed to read the Sentence Transformers config", exc_info=True) + return None + + if isinstance(modules, dict): + modules = modules.get("modules") + if not ( + isinstance(modules, list) + and len(modules) == 1 + and isinstance(modules[0], dict) + and modules[0].get("type") in _ST_TRANSFORMER_MODULE_TYPES + and not modules[0].get("path") + and isinstance(transformer_config, dict) + and transformer_config.get("transformer_task") == "sequence-classification" + and isinstance( + modality_config := transformer_config.get("modality_config"), dict + ) + and "message" in modality_config + ): + return None + + processing_kwargs = transformer_config.get("processing_kwargs") or {} + chat_template_kwargs = ( + processing_kwargs.get("chat_template") or {} + if isinstance(processing_kwargs, dict) + else None + ) + if not isinstance(chat_template_kwargs, dict): + raise ValueError( + "The Sentence Transformers processing_kwargs.chat_template must be " + "a JSON object." + ) + chat_template_kwargs = dict(chat_template_kwargs) + # vLLM does not restore the template suffix after truncation. + chat_template_kwargs.pop("restore_suffix", None) + if not isinstance(chat_template_kwargs.get("chat_template", ""), str): + raise ValueError( + "The Sentence Transformers chat template name must be a string." + ) + if reserved := {"tokenize", "tools"} & chat_template_kwargs.keys(): + raise ValueError( + f"Unsupported Sentence Transformers chat template settings: {reserved}" + ) + return chat_template_kwargs + + def _read_safetensors_metadata_in_dir(local_dir: Path) -> dict[str, Any]: return { param_name: info