class CrossEncoderIOProcessor(ScoringIOProcessor):
name = "cross-encoder"
pooling_task: PoolingTask = "classify"
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if is_mistral_tokenizer(self.tokenizer):
raise ValueError("MistralTokenizer not supported for cross-encoding")
from vllm.model_executor.model_loader import get_model_cls
from vllm.model_executor.models.interfaces import supports_score_template
model = get_model_cls(self.model_config)
self.supports_score_template = supports_score_template(model)
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"
if self.model_config.is_multimodal_model
else "examples/pooling/score/template/qwen3_reranker.jinja"
)
logger.warning(
"Serving an original Qwen3 reranker (%s) without a "
"--chat-template. The model was trained with a specific prompt "
"template; running without it may produce inaccurate relevance "
"scores. Consider specifying `--chat-template %s`.",
self.model_config.model,
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
def get_request_factory_online(
self, ctx: PoolingServeContext
) -> Sequence[AnyRenderParam]:
request = ctx.request
scoring_data = self.valid_inputs_online(request)
data_1 = scoring_data.data_1
data_2 = scoring_data.data_2
num_requests = len(data_2)
if len(data_1) == 1:
data_1 = data_1 * num_requests
max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
request=request
)
tok_params = request.build_tok_params(self.model_config)
seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
requests = [
ScoringRenderParams(
data_1=data_1[i],
data_2=data_2[i],
chat_template=self.chat_template,
max_tokens_per_query=max_tokens_per_query,
max_tokens_per_doc=max_tokens_per_doc,
tok_params=tok_params,
prompt_extras=ctx.prompt_extras,
skip_mm_cache=False,
params=ctx.pooling_params,
lora_requests=seq_lora_requests[i],
priorities=seq_priority[i],
)
for i in range(num_requests)
]
return requests
#######################################
# offline APIs
def get_request_factory_offline(
self, ctx: AnyOfflineInputsContext
) -> tuple[RequestFactory, int]:
assert isinstance(ctx, OfflineScoringInputsContext)
data_1 = ctx.scoring_data.data_1
data_2 = ctx.scoring_data.data_2
num_requests = len(data_2)
if len(data_1) == 1:
data_1 = data_1 * num_requests
max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
pooling_params=ctx.pooling_params
)
tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
**(ctx.tokenization_kwargs or {})
)
prompt_extras = ctx.pooling_params.extra_kwargs
seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
def request_factory() -> RequestGenerator:
for i in range(num_requests):
yield ScoringRenderParams(
data_1=data_1[i],
data_2=data_2[i],
chat_template=(
self.chat_template
if ctx.chat_template is None
else ctx.chat_template
),
max_tokens_per_query=max_tokens_per_query,
max_tokens_per_doc=max_tokens_per_doc,
tok_params=tok_params,
prompt_extras=prompt_extras,
skip_mm_cache=False,
params=ctx.pooling_params,
lora_requests=seq_lora_requests[i],
priorities=seq_priority[i],
)
return request_factory, num_requests
#######################################
# helpers
@cached_property
def num_post_process_tokens(self) -> int:
if self.model is None:
return 0
probe = TokensPrompt(prompt_token_ids=[])
self.model.post_process_tokens(probe)
return len(probe["prompt_token_ids"])
def render(
self,
render_params: EncodeCMPLRenderParams
| EncodeChatRenderParams
| ScoringRenderParams,
) -> PoolingEngineInput:
if "data_1" not in render_params:
raise ValueError(
f"Unsupported render_params type {render_params.__class__.__name__}"
)
render_params = cast(ScoringRenderParams, render_params)
arrival_time = time.time()
tok_params = render_params["tok_params"]
params = render_params["params"]
prompt_extras = render_params["prompt_extras"]
# Leave room for the tokens that post_process_tokens appends after
# truncation (e.g. the JinaVL score token), so the limits count them.
if self.num_post_process_tokens:
tok_params = reserve_tokens(tok_params, self.num_post_process_tokens)
_, engine_prompt = self.get_score_prompt(
data_1=render_params["data_1"],
data_2=render_params["data_2"],
encode_kwargs=tok_params.get_encode_kwargs(),
chat_template=render_params["chat_template"],
max_tokens_per_query=render_params["max_tokens_per_query"],
max_tokens_per_doc=render_params["max_tokens_per_doc"],
chat_template_kwargs=prompt_extras.get("chat_template_kwargs")
if prompt_extras
else None,
)
tok_params.apply_post_tokenization(self.tokenizer, engine_prompt)
if self.model is not None:
self.model.post_process_tokens(engine_prompt)
if token_type_ids := engine_prompt.pop("token_type_ids", None):
params = params.clone()
compressed = compress_token_type_ids(
_apply_post_tokenization_to_token_type_ids(
self.tokenizer, tok_params, token_type_ids
)
)
params.extra_kwargs = {
**(params.extra_kwargs or {}),
"compressed_token_type_ids": compressed,
}
engine_prompt_extras = (
{
k: v
for k in ("mm_processor_kwargs", "cache_salt")
if (v := prompt_extras.get(k)) is not None
}
if prompt_extras
else None
)
if engine_prompt_extras:
target_prompt = extract_target_prompt(self.model_config, engine_prompt)
target_prompt.update(engine_prompt_extras)
engine_input = self.renderer.process_for_engine(engine_prompt, arrival_time)
return PoolingEngineInput(
prompts=engine_input,
params=params,
lora_requests=render_params["lora_requests"],
priorities=render_params["priorities"],
)
def get_score_prompt(
self,
data_1: ScoreData,
data_2: ScoreData,
encode_kwargs: dict[str, Any],
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
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,
model_config,
)
# Limit the query and document separately before composing them.
if max_tokens_per_query > 0 and isinstance(prompt_1, str):
prompt_1 = truncate_text_to_tokens(
prompt_1, tokenizer, max_tokens_per_query
)
if max_tokens_per_doc > 0 and isinstance(prompt_2, str):
prompt_2 = truncate_text_to_tokens(prompt_2, tokenizer, max_tokens_per_doc)
def default_tokenizer_encode():
local_kwargs = encode_kwargs.copy()
if self.supports_score_template:
assert self.model is not None
full_prompt = self.model.get_score_template(prompt_1, prompt_2)
if full_prompt is None:
raise ValueError("Get empty score template from model")
prompt_inputs = tokenizer(full_prompt, **local_kwargs)
else:
if self.use_sep_token:
# cross_encoder models defaults to using separating token.
if max_tokens_per_doc > 0 and isinstance(prompt_2, str):
query_tokens = tokenizer.encode(
prompt_1, add_special_tokens=False
)
num_special = (
get_num_special_tokens_for_pair(tokenizer)
if local_kwargs.get("add_special_tokens", True)
else 0
)
doc_limit_max_length = (
len(query_tokens) + max_tokens_per_doc + num_special
)
existing_max_length = local_kwargs.get("max_length")
if existing_max_length is not None:
effective_max_length = min(
doc_limit_max_length, existing_max_length
)
else:
effective_max_length = doc_limit_max_length
if len(query_tokens) + num_special < effective_max_length:
local_kwargs["truncation"] = "only_second"
local_kwargs["max_length"] = effective_max_length
prompt_inputs = tokenizer(
text=prompt_1, text_pair=prompt_2, **local_kwargs
)
full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
else:
# `llm as reranker` defaults to not using separating token.
full_prompt = prompt_1 + prompt_2
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,
# 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.
if chat_template is None:
full_prompt, prompt_inputs = default_tokenizer_encode()
else:
# FIXME:
# Try applying a score template from the CLI arg or tokenizer_config.json
# 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,
[
{"role": "query", "content": prompt_1},
{"role": "document", "content": prompt_2},
],
chat_template=chat_template,
tools=None,
tokenize=False,
**_safe_kwargs,
)
prompt_inputs = tokenizer(full_prompt, **encode_kwargs)
except ChatTemplateResolutionError:
full_prompt, prompt_inputs = default_tokenizer_encode()
engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"])
if (token_type_ids := prompt_inputs.get("token_type_ids")) is not None:
engine_prompt["token_type_ids"] = token_type_ids
if mm_data is not None:
engine_prompt["multi_modal_data"] = mm_data
if mm_uuids is not None:
engine_prompt["multi_modal_uuids"] = mm_uuids
return full_prompt, engine_prompt