From 70c00163ffa80f5821bc8f87f80bde632bc29015 Mon Sep 17 00:00:00 2001 From: Krish Gupta Date: Thu, 14 May 2026 07:11:22 +0530 Subject: [PATCH] [Feature] Add instruction support for score/rerank chat templates (#42412) Signed-off-by: KrxGu --- .../score/template/qwen3_reranker.jinja | 2 +- .../score/template/qwen3_vl_reranker.jinja | 8 +- .../test_cross_encoder_online_vision.py | 132 ++++++++++++++++++ .../pooling/scoring/io_processor.py | 19 ++- vllm/entrypoints/pooling/scoring/protocol.py | 33 ++++- 5 files changed, 182 insertions(+), 12 deletions(-) diff --git a/examples/pooling/score/template/qwen3_reranker.jinja b/examples/pooling/score/template/qwen3_reranker.jinja index f33f526dc05..558e9b9566d 100644 --- a/examples/pooling/score/template/qwen3_reranker.jinja +++ b/examples/pooling/score/template/qwen3_reranker.jinja @@ -1,7 +1,7 @@ <|im_start|>system 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|> <|im_start|>user -: {{ messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query") }} +: {{ instruction | default(instruct | default(messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query", true), true), true) }} : {{ messages | selectattr("role", "eq", "query") | map(attribute="content") | first }} : {{ messages | selectattr("role", "eq", "document") | map(attribute="content") | first }}<|im_end|> <|im_start|>assistant diff --git a/examples/pooling/score/template/qwen3_vl_reranker.jinja b/examples/pooling/score/template/qwen3_vl_reranker.jinja index ed89f2a547a..90488d0f5d3 100644 --- a/examples/pooling/score/template/qwen3_vl_reranker.jinja +++ b/examples/pooling/score/template/qwen3_vl_reranker.jinja @@ -1,13 +1,7 @@ <|im_start|>system 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|> <|im_start|>user -: {{ - messages - | selectattr("role", "eq", "system") - | map(attribute="content") - | first - | default("Given a search query, retrieve relevant candidates that answer the query.") -}}:{{ +: {{ instruction | default(instruct | default(messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a search query, retrieve relevant candidates that answer the query.", true), true), true) }}:{{ messages | selectattr("role", "eq", "query") | map(attribute="content") diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py index a815e6e509c..e6b4d3f873e 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py @@ -377,3 +377,135 @@ async def test_score_api_queries_list_documents_list( backend, "paired[3]_text_vs_text_plus_image", ) + + +INSTRUCTION = ( + "Given a multimodal retrieval query, retrieve candidates that " + "visually or textually match the requested scene, object, or action." +) + + +@pytest.mark.asyncio +async def test_score_api_instruction_field( + server: tuple[RemoteOpenAIServer, str], +): + remote_server, _ = server + + default_response = requests.post( + remote_server.url_for("score"), + json={ + "model": MODEL_NAME, + "queries": query, + "documents": document, + }, + ) + default_response.raise_for_status() + default_score = ScoreResponse.model_validate(default_response.json()) + + instruction_response = requests.post( + remote_server.url_for("score"), + json={ + "model": MODEL_NAME, + "queries": query, + "documents": document, + "instruction": INSTRUCTION, + }, + ) + instruction_response.raise_for_status() + instruction_score = ScoreResponse.model_validate(instruction_response.json()) + + assert instruction_score.id is not None + assert instruction_score.data is not None + assert len(instruction_score.data) == 1 + assert instruction_score.usage.prompt_tokens > default_score.usage.prompt_tokens + + +@pytest.mark.asyncio +async def test_rerank_api_instruction_field( + server: tuple[RemoteOpenAIServer, str], +): + remote_server, _ = server + + doc_list = [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ] + + default_response = requests.post( + remote_server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": doc_list, + }, + ) + default_response.raise_for_status() + default_rerank = RerankResponse.model_validate(default_response.json()) + + instruction_response = requests.post( + remote_server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": doc_list, + "instruction": INSTRUCTION, + }, + ) + instruction_response.raise_for_status() + instruction_rerank = RerankResponse.model_validate(instruction_response.json()) + + assert instruction_rerank.id is not None + assert instruction_rerank.model is not None + assert instruction_rerank.usage is not None + assert len(instruction_rerank.results) == len(default_rerank.results) + assert instruction_rerank.usage.prompt_tokens > default_rerank.usage.prompt_tokens + + +@pytest.mark.asyncio +async def test_rerank_api_instruction_field_matches_chat_template_kwargs( + server: tuple[RemoteOpenAIServer, str], +): + remote_server, _ = server + + doc_list = [ + document, + {"content": [documents[0]]}, + {"content": [documents[1]]}, + {"content": [documents[0], documents[1]]}, + ] + + field_response = requests.post( + remote_server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": doc_list, + "instruction": INSTRUCTION, + }, + ) + field_response.raise_for_status() + field_rerank = RerankResponse.model_validate(field_response.json()) + + kwargs_response = requests.post( + remote_server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": doc_list, + "chat_template_kwargs": {"instruction": INSTRUCTION}, + }, + ) + kwargs_response.raise_for_status() + kwargs_rerank = RerankResponse.model_validate(kwargs_response.json()) + + assert kwargs_rerank.usage.prompt_tokens == field_rerank.usage.prompt_tokens + + field_scores = [ + r.relevance_score for r in sorted(field_rerank.results, key=lambda x: x.index) + ] + kwargs_scores = [ + r.relevance_score for r in sorted(kwargs_rerank.results, key=lambda x: x.index) + ] + assert field_scores == pytest.approx(kwargs_scores) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index e706601c568..7c58e7e9aba 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -157,7 +157,7 @@ class BiEncoderIOProcessor(ScoringIOProcessor): tok_params, prompt_extras={ k: v - for k in ("mm_processor_kwargs", "cache_salt") + for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs") if (v := getattr(request, k, None)) is not None }, ) @@ -384,7 +384,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): max_tokens_per_doc=max_tokens_per_doc, prompt_extras={ k: v - for k in ("mm_processor_kwargs", "cache_salt") + for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs") if (v := getattr(request, k, None)) is not None }, ) @@ -407,6 +407,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): pooling_params=ctx.pooling_params ) + prompt_extras = ctx.pooling_params.extra_kwargs if ctx.pooling_params else None engine_inputs, pooling_params_list = self._pre_process( ctx.prompts, tok_params, @@ -414,6 +415,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): ctx.chat_template, max_tokens_per_query=max_tokens_per_query, max_tokens_per_doc=max_tokens_per_doc, + prompt_extras=prompt_extras, ) ctx.pooling_params = pooling_params_list return engine_inputs @@ -453,6 +455,9 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): chat_template=chat_template, max_tokens_per_query=max_tokens_per_query, max_tokens_per_doc=max_tokens_per_doc, + chat_template_kwargs=prompt_extras.get("chat_template_kwargs") + if prompt_extras + else None, ) if token_type_ids := engine_prompt.pop("token_type_ids", None): @@ -477,6 +482,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): chat_template: str | None = None, max_tokens_per_query: int = 0, max_tokens_per_doc: int = 0, + chat_template_kwargs: dict[str, Any] | None = None, ): model_config = self.model_config tokenizer = self.tokenizer @@ -556,6 +562,14 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): # If that fails because there is no such template, # fall back to the default implementation. try: + _safe_kwargs = chat_template_kwargs or {} + _reserved = {"chat_template", "tools", "tokenize"} + _unexpected = _reserved & _safe_kwargs.keys() + if _unexpected: + raise ValueError( + "chat_template_kwargs contains reserved keys that " + f"conflict with fixed scorer arguments: {_unexpected}" + ) full_prompt = safe_apply_chat_template( model_config, tokenizer, @@ -566,6 +580,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): chat_template=chat_template, tools=None, tokenize=False, + **_safe_kwargs, ) prompt_inputs = tokenizer(full_prompt, **encode_kwargs) except ChatTemplateResolutionError: diff --git a/vllm/entrypoints/pooling/scoring/protocol.py b/vllm/entrypoints/pooling/scoring/protocol.py index b2fa0f2194a..49cce7a4ee2 100644 --- a/vllm/entrypoints/pooling/scoring/protocol.py +++ b/vllm/entrypoints/pooling/scoring/protocol.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time -from typing import TypeAlias +from typing import Any, TypeAlias -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator from vllm import PoolingParams from vllm.config import ModelConfig @@ -35,8 +35,37 @@ class ScoringRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): "applies to the combined query+document)." ), ) + instruction: str | None = Field( + default=None, + description=( + "Task instruction prepended to each scored pair via the chat " + "template. Equivalent to passing " + "chat_template_kwargs={'instruction': ...}." + ), + ) + chat_template_kwargs: dict[str, Any] | None = Field( + default=None, + description=( + "Additional keyword args to pass to the chat template renderer. " + "Will be accessible by the score/rerank chat template." + ), + ) # --8<-- [end:scoring-common-params] + @model_validator(mode="after") + def _merge_instruction_into_kwargs(self) -> "ScoringRequestMixin": + """Fold the top-level `instruction` field into `chat_template_kwargs`. + + This allows callers to use either the convenience field or the generic + dict. Explicit keys inside `chat_template_kwargs` take precedence over + the top-level `instruction` field. + """ + if self.instruction is not None: + merged = dict(self.chat_template_kwargs or {}) + merged.setdefault("instruction", self.instruction) + self.chat_template_kwargs = merged + return self + def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: return self._build_pooling_tok_params( model_config,