[Feature] Add instruction support for score/rerank chat templates (#42412)

Signed-off-by: KrxGu <[email protected]>
This commit is contained in:
Krish Gupta
2026-05-14 09:41:22 +08:00
committed by GitHub
parent f51f6844f9
commit 70c00163ff
5 changed files with 182 additions and 12 deletions
@@ -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
<Instruct>: {{ messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query") }}
<Instruct>: {{ 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) }}
<Query>: {{ messages | selectattr("role", "eq", "query") | map(attribute="content") | first }}
<Document>: {{ messages | selectattr("role", "eq", "document") | map(attribute="content") | first }}<|im_end|>
<|im_start|>assistant
@@ -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
<Instruct>: {{
messages
| selectattr("role", "eq", "system")
| map(attribute="content")
| first
| default("Given a search query, retrieve relevant candidates that answer the query.")
}}<Query>:{{
<Instruct>: {{ 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) }}<Query>:{{
messages
| selectattr("role", "eq", "query")
| map(attribute="content")
@@ -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)
@@ -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:
+31 -2
View File
@@ -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,