| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import time |
| 4 | from collections.abc import Sequence |
| 5 | from typing import Any, TypeAlias, cast |
| 6 | |
| 7 | import torch.nn.functional as F |
| 8 | |
| 9 | from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.renderers import TokenizeParams |
| 12 | from vllm.renderers.hf import safe_apply_chat_template |
| 13 | from vllm.renderers.inputs.preprocess import ( |
| 14 | extract_target_prompt, |
| 15 | parse_model_prompt, |
| 16 | prompt_to_seq, |
| 17 | ) |
| 18 | from vllm.tasks import PoolingTask |
| 19 | from vllm.utils.mistral import is_mistral_tokenizer |
| 20 | |
| 21 | from ...chat_utils import ChatTemplateResolutionError |
| 22 | from ..base.io_processor import PoolingIOProcessor |
| 23 | from ..pooling.protocol import PoolingCompletionRequest |
| 24 | from ..typing import ( |
| 25 | AnyOfflineInputsContext, |
| 26 | AnyPoolingRequest, |
| 27 | AnyRenderParam, |
| 28 | EncodeChatRenderParams, |
| 29 | EncodeCMPLRenderParams, |
| 30 | OfflineEncodeInputsContext, |
| 31 | OfflineOutputsContext, |
| 32 | OfflineScoringInputsContext, |
| 33 | PoolingEngineInput, |
| 34 | PoolingServeContext, |
| 35 | RequestFactory, |
| 36 | RequestGenerator, |
| 37 | ScoringRenderParams, |
| 38 | ) |
| 39 | from .protocol import RerankRequest, ScoreRequest, ScoringRequest |
| 40 | from .typing import ScoreData, ScoreInput, ScoringData |
| 41 | from .utils import ( |
| 42 | compress_token_type_ids, |
| 43 | compute_maxsim_score, |
| 44 | get_num_special_tokens_for_pair, |
| 45 | parse_score_data, |
| 46 | score_data_to_prompts, |
| 47 | truncate_text_to_tokens, |
| 48 | validate_score_input, |
| 49 | ) |
| 50 | |
| 51 | logger = init_logger(__name__) |
| 52 | |
| 53 | ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest] |
| 54 | |
| 55 | |
| 56 | def _apply_post_tokenization_to_token_type_ids( |
| 57 | tokenizer: Any, |
| 58 | tok_params: TokenizeParams, |
| 59 | token_type_ids: list[int], |
| 60 | ) -> list[int]: |
| 61 | # Must stay in the same order as `TokenizeParams._validate_tokens`, which |
| 62 | # truncates before padding. These ids are parallel to the prompt tokens |
| 63 | # and locate the query/document boundary for the cross-encoder; applying |
| 64 | # the two steps in a different order silently misplaces that boundary. |
| 65 | max_length = tok_params.truncate_prompt_tokens |
| 66 | if max_length is not None and max_length < 0: |
| 67 | max_length = tok_params.max_input_tokens |
| 68 | |
| 69 | if max_length is not None and max_length < len(token_type_ids): |
| 70 | if max_length == 0: |
| 71 | token_type_ids = token_type_ids[:0] |
| 72 | else: |
| 73 | side = tok_params.truncation_side or ( |
| 74 | tokenizer.truncation_side if tokenizer is not None else None |
| 75 | ) |
| 76 | token_type_ids = ( |
| 77 | token_type_ids[-max_length:] |
| 78 | if side == "left" |
| 79 | else token_type_ids[:max_length] |
| 80 | ) |
| 81 | |
| 82 | pad_length = tok_params.pad_prompt_tokens |
| 83 | if pad_length is not None and pad_length < 0: |
| 84 | pad_length = tok_params.max_input_tokens |
| 85 | |
| 86 | if pad_length is not None and pad_length > len(token_type_ids): |
| 87 | pad_token_type_id = token_type_ids[-1] if token_type_ids else 0 |
| 88 | token_type_ids = token_type_ids + [pad_token_type_id] * ( |
| 89 | pad_length - len(token_type_ids) |
| 90 | ) |
| 91 | |
| 92 | return token_type_ids |
| 93 | |
| 94 | |
| 95 | class ScoringIOProcessor(PoolingIOProcessor): |
| 96 | name: str |
| 97 | pooling_task: PoolingTask |
| 98 | |
| 99 | def __init__(self, *args, **kwargs): |
| 100 | super().__init__(*args, **kwargs) |
| 101 | |
| 102 | self.tokenizer = self.renderer.get_tokenizer() |
| 103 | self.architecture = self.model_config.architecture |
| 104 | self.is_multimodal_model = self.model_config.is_multimodal_model |
| 105 | self.pad_token_id = self.tokenizer.pad_token_id |
| 106 | |
| 107 | def create_pooling_params(self, request): |
| 108 | return request.to_pooling_params(self.pooling_task) |
| 109 | |
| 110 | def _validate_token_limit(self, value: int, name: str) -> None: |
| 111 | if value < 0: |
| 112 | raise ValueError(f"{name} must be a non-negative integer") |
| 113 | if value >= self.model_config.max_model_len: |
| 114 | raise ValueError( |
| 115 | f"{name} ({value}) must be less " |
| 116 | f"than max_model_len ({self.model_config.max_model_len})." |
| 117 | ) |
| 118 | |
| 119 | def _get_token_limits( |
| 120 | self, |
| 121 | request: ScoringRequest | None = None, |
| 122 | pooling_params: PoolingParams | None = None, |
| 123 | ) -> tuple[int, int]: |
| 124 | """Extract and validate token limits from request or pooling_params.""" |
| 125 | if request is not None: |
| 126 | max_tokens_per_query = getattr(request, "max_tokens_per_query", 0) |
| 127 | max_tokens_per_doc = getattr(request, "max_tokens_per_doc", 0) |
| 128 | else: |
| 129 | extra = ( |
| 130 | (pooling_params.extra_kwargs or {}) |
| 131 | if pooling_params is not None |
| 132 | else {} |
| 133 | ) |
| 134 | max_tokens_per_query = extra.get("max_tokens_per_query", 0) |
| 135 | max_tokens_per_doc = extra.get("max_tokens_per_doc", 0) |
| 136 | |
| 137 | if max_tokens_per_query != 0: |
| 138 | self._validate_token_limit(max_tokens_per_query, "max_tokens_per_query") |
| 139 | if max_tokens_per_doc != 0: |
| 140 | self._validate_token_limit(max_tokens_per_doc, "max_tokens_per_doc") |
| 141 | return max_tokens_per_query, max_tokens_per_doc |
| 142 | |
| 143 | def _truncate_scoring_data( |
| 144 | self, |
| 145 | scoring_data: ScoringData, |
| 146 | max_tokens_per_query: int = 0, |
| 147 | max_tokens_per_doc: int = 0, |
| 148 | ) -> ScoringData: |
| 149 | """Truncate query/document texts to token limits.""" |
| 150 | data_1 = scoring_data.data_1 |
| 151 | data_2 = scoring_data.data_2 |
| 152 | if max_tokens_per_query > 0: |
| 153 | data_1 = [ |
| 154 | truncate_text_to_tokens(d, self.tokenizer, max_tokens_per_query) |
| 155 | if isinstance(d, str) |
| 156 | else d |
| 157 | for d in data_1 |
| 158 | ] |
| 159 | if max_tokens_per_doc > 0: |
| 160 | data_2 = [ |
| 161 | truncate_text_to_tokens(d, self.tokenizer, max_tokens_per_doc) |
| 162 | if isinstance(d, str) |
| 163 | else d |
| 164 | for d in data_2 |
| 165 | ] |
| 166 | return ScoringData(data_1=data_1, data_2=data_2) |
| 167 | |
| 168 | def valid_inputs( |
| 169 | self, |
| 170 | data_1: ScoreInput | list[ScoreInput], |
| 171 | data_2: ScoreInput | list[ScoreInput], |
| 172 | ) -> ScoringData: |
| 173 | scoring_data = validate_score_input( |
| 174 | data_1, |
| 175 | data_2, |
| 176 | is_multimodal_model=self.is_multimodal_model, |
| 177 | architecture=self.architecture, |
| 178 | ) |
| 179 | return scoring_data |
| 180 | |
| 181 | def valid_inputs_online(self, request: AnyPoolingRequest): |
| 182 | if isinstance(request, ScoreRequest): |
| 183 | data_1 = request.data_1 |
| 184 | data_2 = request.data_2 |
| 185 | elif isinstance(request, RerankRequest): |
| 186 | data_1 = request.query |
| 187 | data_2 = request.documents |
| 188 | else: |
| 189 | raise ValueError(f"Invalid {request.__class__.__name__} request type") |
| 190 | |
| 191 | scoring_data = self.valid_inputs(data_1, data_2) |
| 192 | return scoring_data |
| 193 | |
| 194 | def _pair_error_output( |
| 195 | self, |
| 196 | output_1: PoolingRequestOutput, |
| 197 | output_2: PoolingRequestOutput, |
| 198 | ) -> PoolingRequestOutput | None: |
| 199 | failed = output_1 if output_1.error is not None else output_2 |
| 200 | if failed.error is None: |
| 201 | return None |
| 202 | |
| 203 | padding = [self.pad_token_id] if self.pad_token_id is not None else [] |
| 204 | return PoolingRequestOutput( |
| 205 | request_id=f"{output_1.request_id}_{output_2.request_id}", |
| 206 | outputs=failed.outputs, |
| 207 | prompt_token_ids=( |
| 208 | output_1.prompt_token_ids + padding + output_2.prompt_token_ids |
| 209 | ), |
| 210 | num_cached_tokens=(output_1.num_cached_tokens + output_2.num_cached_tokens), |
| 211 | finished=failed.finished, |
| 212 | error=failed.error, |
| 213 | ) |
| 214 | |
| 215 | |
| 216 | class BiEncoderIOProcessor(ScoringIOProcessor): |
| 217 | name = "bi-encoder" |
| 218 | pooling_task: PoolingTask = "embed" |
| 219 | |
| 220 | ####################################### |
| 221 | # online APIs |
| 222 | |
| 223 | def get_request_factory_online( |
| 224 | self, ctx: PoolingServeContext |
| 225 | ) -> Sequence[AnyRenderParam]: |
| 226 | request = ctx.request |
| 227 | scoring_data = self.valid_inputs_online(request) |
| 228 | |
| 229 | max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( |
| 230 | request=request |
| 231 | ) |
| 232 | if max_tokens_per_query > 0 or max_tokens_per_doc > 0: |
| 233 | scoring_data = self._truncate_scoring_data( |
| 234 | scoring_data, max_tokens_per_query, max_tokens_per_doc |
| 235 | ) |
| 236 | |
| 237 | data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config) |
| 238 | data_2 = score_data_to_prompts( |
| 239 | scoring_data.data_2, "document", self.model_config |
| 240 | ) |
| 241 | prompts = data_1 + data_2 |
| 242 | ctx.n_queries = len(data_1) |
| 243 | |
| 244 | prompts_seq = prompt_to_seq(prompts) |
| 245 | parsed_prompts = [ |
| 246 | parse_model_prompt(self.model_config, prompt) for prompt in prompts_seq |
| 247 | ] |
| 248 | num_requests = len(parsed_prompts) |
| 249 | |
| 250 | tok_params = request.build_tok_params(self.model_config) |
| 251 | params_seq = self._params_to_seq(ctx.pooling_params, num_requests) |
| 252 | seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) |
| 253 | seq_priority = self._priority_to_seq(ctx.priorities, num_requests) |
| 254 | |
| 255 | requests = [ |
| 256 | EncodeCMPLRenderParams( |
| 257 | prompts=parsed_prompts[i], |
| 258 | tok_params=tok_params, |
| 259 | prompt_extras=ctx.prompt_extras, |
| 260 | skip_mm_cache=False, |
| 261 | params=params_seq[i], |
| 262 | lora_requests=seq_lora_requests[i], |
| 263 | priorities=seq_priority[i], |
| 264 | ) |
| 265 | for i in range(num_requests) |
| 266 | ] |
| 267 | return requests |
| 268 | |
| 269 | def post_process_online( |
| 270 | self, |
| 271 | ctx: ScoringServeContext, |
| 272 | ): |
| 273 | assert ctx.final_res_batch is not None |
| 274 | assert isinstance(ctx.n_queries, int) |
| 275 | |
| 276 | ctx.final_res_batch = self._post_process( |
| 277 | outputs=ctx.final_res_batch, n_queries=ctx.n_queries |
| 278 | ) |
| 279 | |
| 280 | ####################################### |
| 281 | # offline APIs |
| 282 | |
| 283 | def get_request_factory_offline( |
| 284 | self, ctx: AnyOfflineInputsContext |
| 285 | ) -> tuple[RequestFactory, int]: |
| 286 | assert isinstance(ctx, OfflineScoringInputsContext) |
| 287 | |
| 288 | max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( |
| 289 | pooling_params=ctx.pooling_params |
| 290 | ) |
| 291 | |
| 292 | scoring_data = ctx.scoring_data |
| 293 | if max_tokens_per_query > 0 or max_tokens_per_doc > 0: |
| 294 | scoring_data = self._truncate_scoring_data( |
| 295 | scoring_data, max_tokens_per_query, max_tokens_per_doc |
| 296 | ) |
| 297 | |
| 298 | data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config) |
| 299 | data_2 = score_data_to_prompts( |
| 300 | scoring_data.data_2, "document", self.model_config |
| 301 | ) |
| 302 | prompts = data_1 + data_2 |
| 303 | |
| 304 | return super().get_request_factory_offline( |
| 305 | OfflineEncodeInputsContext( |
| 306 | pooling_task=self.pooling_task, |
| 307 | prompts=prompts, |
| 308 | tokenization_kwargs=ctx.tokenization_kwargs, |
| 309 | pooling_params=ctx.pooling_params, |
| 310 | lora_request=ctx.lora_request, |
| 311 | priorities=ctx.priorities, |
| 312 | ) |
| 313 | ) |
| 314 | |
| 315 | def post_process_offline( |
| 316 | self, |
| 317 | ctx: OfflineOutputsContext, |
| 318 | ) -> list[PoolingRequestOutput]: |
| 319 | assert ctx.n_queries is not None |
| 320 | return self._post_process(outputs=ctx.outputs, n_queries=ctx.n_queries) |
| 321 | |
| 322 | ####################################### |
| 323 | # helpers |
| 324 | |
| 325 | def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): |
| 326 | emb_data_1 = outputs[:n_queries] |
| 327 | emb_data_2 = outputs[n_queries:] |
| 328 | |
| 329 | if len(emb_data_1) == 1: |
| 330 | emb_data_1 = emb_data_1 * len(emb_data_2) |
| 331 | |
| 332 | final_res_batch: list[PoolingRequestOutput] = [] |
| 333 | for emb_1, emb_2 in zip(emb_data_1, emb_data_2): |
| 334 | if error_output := self._pair_error_output(emb_1, emb_2): |
| 335 | final_res_batch.append(error_output) |
| 336 | continue |
| 337 | |
| 338 | pair_score = F.cosine_similarity( |
| 339 | emb_1.outputs.data.float(), emb_2.outputs.data.float(), dim=0 |
| 340 | ) |
| 341 | |
| 342 | padding: list[int] = [] |
| 343 | if self.pad_token_id is not None: |
| 344 | padding = [self.pad_token_id] |
| 345 | |
| 346 | tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids |
| 347 | |
| 348 | final_res_batch.append( |
| 349 | PoolingRequestOutput( |
| 350 | request_id=f"{emb_1.request_id}_{emb_2.request_id}", |
| 351 | outputs=pair_score, |
| 352 | prompt_token_ids=tokens, |
| 353 | num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens, |
| 354 | finished=True, |
| 355 | ) |
| 356 | ) |
| 357 | return final_res_batch |
| 358 | |
| 359 | |
| 360 | class LateInteractionIOProcessor(BiEncoderIOProcessor): |
| 361 | name = "late-interaction" |
| 362 | pooling_task: PoolingTask = "token_embed" |
| 363 | |
| 364 | def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): |
| 365 | # Split into query and document embeddings |
| 366 | emb_data_1 = outputs[:n_queries] |
| 367 | emb_data_2 = outputs[n_queries:] |
| 368 | |
| 369 | # Expand queries if 1:N scoring |
| 370 | if len(emb_data_1) == 1: |
| 371 | emb_data_1 = emb_data_1 * len(emb_data_2) |
| 372 | |
| 373 | final_res_batch: list[PoolingRequestOutput] = [] |
| 374 | padding: list[int] = [] |
| 375 | if (pad_token_id := self.pad_token_id) is not None: |
| 376 | padding = [pad_token_id] |
| 377 | |
| 378 | # Compute MaxSim scores |
| 379 | for emb_1, emb_2 in zip(emb_data_1, emb_data_2): |
| 380 | if error_output := self._pair_error_output(emb_1, emb_2): |
| 381 | final_res_batch.append(error_output) |
| 382 | continue |
| 383 | |
| 384 | # emb_1.outputs.data: [query_len, dim] |
| 385 | # emb_2.outputs.data: [doc_len, dim] |
| 386 | q_emb = emb_1.outputs.data |
| 387 | d_emb = emb_2.outputs.data |
| 388 | |
| 389 | maxsim_score = compute_maxsim_score(q_emb, d_emb) |
| 390 | |
| 391 | tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids |
| 392 | |
| 393 | final_res_batch.append( |
| 394 | PoolingRequestOutput( |
| 395 | request_id=f"{emb_1.request_id}_{emb_2.request_id}", |
| 396 | outputs=maxsim_score, |
| 397 | prompt_token_ids=tokens, |
| 398 | num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens, |
| 399 | finished=True, |
| 400 | ) |
| 401 | ) |
| 402 | return final_res_batch |
| 403 | |
| 404 | |
| 405 | class FlashLateInteractionIOProcessor(LateInteractionIOProcessor): |
| 406 | name = "flash-late-interaction" |
| 407 | |
| 408 | def post_process_online( |
| 409 | self, |
| 410 | ctx: ScoringServeContext, |
| 411 | ): |
| 412 | assert ctx.query_final_res_batch is not None |
| 413 | assert ctx.final_res_batch is not None |
| 414 | assert isinstance(ctx.n_queries, int) |
| 415 | |
| 416 | # Expand queries if 1:N scoring |
| 417 | if len(ctx.query_final_res_batch) == 1: |
| 418 | ctx.query_final_res_batch = ctx.query_final_res_batch * len( |
| 419 | ctx.final_res_batch |
| 420 | ) |
| 421 | |
| 422 | final_res_batch: list[PoolingRequestOutput] = [] |
| 423 | for d1, d2 in zip(ctx.query_final_res_batch, ctx.final_res_batch): |
| 424 | padding: list[int] = [] |
| 425 | if (pad_token_id := self.pad_token_id) is not None: |
| 426 | padding = [pad_token_id] |
| 427 | |
| 428 | tokens = d1.prompt_token_ids + padding + d2.prompt_token_ids |
| 429 | |
| 430 | final_res_batch.append( |
| 431 | PoolingRequestOutput( |
| 432 | request_id=f"{d1.request_id}_{d2.request_id}", |
| 433 | outputs=d2.outputs, |
| 434 | prompt_token_ids=tokens, |
| 435 | num_cached_tokens=d1.num_cached_tokens + d2.num_cached_tokens, |
| 436 | finished=True, |
| 437 | ) |
| 438 | ) |
| 439 | ctx.final_res_batch = final_res_batch |
| 440 | |
| 441 | |
| 442 | class CrossEncoderIOProcessor(ScoringIOProcessor): |
| 443 | name = "cross-encoder" |
| 444 | pooling_task: PoolingTask = "classify" |
| 445 | |
| 446 | def __init__(self, *args, **kwargs): |
| 447 | super().__init__(*args, **kwargs) |
| 448 | |
| 449 | if is_mistral_tokenizer(self.tokenizer): |
| 450 | raise ValueError("MistralTokenizer not supported for cross-encoding") |
| 451 | |
| 452 | from vllm.model_executor.model_loader import get_model_cls |
| 453 | from vllm.model_executor.models.interfaces import supports_score_template |
| 454 | |
| 455 | model = get_model_cls(self.model_config) |
| 456 | self.supports_score_template = supports_score_template(model) |
| 457 | self.model = model if self.supports_score_template else None |
| 458 | self.use_sep_token = self.model_config.use_sep_token |
| 459 | |
| 460 | if ( |
| 461 | getattr(self.model_config.hf_config, "is_original_qwen3_reranker", False) |
| 462 | and self.chat_template is None |
| 463 | ): |
| 464 | suggested_template = ( |
| 465 | "examples/pooling/score/template/qwen3_vl_reranker.jinja" |
| 466 | if self.model_config.is_multimodal_model |
| 467 | else "examples/pooling/score/template/qwen3_reranker.jinja" |
| 468 | ) |
| 469 | logger.warning( |
| 470 | "Serving an original Qwen3 reranker (%s) without a " |
| 471 | "--chat-template. The model was trained with a specific prompt " |
| 472 | "template; running without it may produce inaccurate relevance " |
| 473 | "scores. Consider specifying `--chat-template %s`.", |
| 474 | self.model_config.model, |
| 475 | suggested_template, |
| 476 | ) |
| 477 | |
| 478 | ####################################### |
| 479 | # online APIs |
| 480 | |
| 481 | def get_request_factory_online( |
| 482 | self, ctx: PoolingServeContext |
| 483 | ) -> Sequence[AnyRenderParam]: |
| 484 | request = ctx.request |
| 485 | scoring_data = self.valid_inputs_online(request) |
| 486 | data_1 = scoring_data.data_1 |
| 487 | data_2 = scoring_data.data_2 |
| 488 | num_requests = len(data_2) |
| 489 | |
| 490 | if len(data_1) == 1: |
| 491 | data_1 = data_1 * num_requests |
| 492 | |
| 493 | max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( |
| 494 | request=request |
| 495 | ) |
| 496 | |
| 497 | tok_params = request.build_tok_params(self.model_config) |
| 498 | seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) |
| 499 | seq_priority = self._priority_to_seq(ctx.priorities, num_requests) |
| 500 | |
| 501 | requests = [ |
| 502 | ScoringRenderParams( |
| 503 | data_1=data_1[i], |
| 504 | data_2=data_2[i], |
| 505 | chat_template=self.chat_template, |
| 506 | max_tokens_per_query=max_tokens_per_query, |
| 507 | max_tokens_per_doc=max_tokens_per_doc, |
| 508 | tok_params=tok_params, |
| 509 | prompt_extras=ctx.prompt_extras, |
| 510 | skip_mm_cache=False, |
| 511 | params=ctx.pooling_params, |
| 512 | lora_requests=seq_lora_requests[i], |
| 513 | priorities=seq_priority[i], |
| 514 | ) |
| 515 | for i in range(num_requests) |
| 516 | ] |
| 517 | |
| 518 | return requests |
| 519 | |
| 520 | ####################################### |
| 521 | # offline APIs |
| 522 | |
| 523 | def get_request_factory_offline( |
| 524 | self, ctx: AnyOfflineInputsContext |
| 525 | ) -> tuple[RequestFactory, int]: |
| 526 | assert isinstance(ctx, OfflineScoringInputsContext) |
| 527 | |
| 528 | data_1 = ctx.scoring_data.data_1 |
| 529 | data_2 = ctx.scoring_data.data_2 |
| 530 | num_requests = len(data_2) |
| 531 | |
| 532 | if len(data_1) == 1: |
| 533 | data_1 = data_1 * num_requests |
| 534 | |
| 535 | max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( |
| 536 | pooling_params=ctx.pooling_params |
| 537 | ) |
| 538 | tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( |
| 539 | **(ctx.tokenization_kwargs or {}) |
| 540 | ) |
| 541 | prompt_extras = ctx.pooling_params.extra_kwargs |
| 542 | |
| 543 | seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) |
| 544 | seq_priority = self._priority_to_seq(ctx.priorities, num_requests) |
| 545 | |
| 546 | def request_factory() -> RequestGenerator: |
| 547 | for i in range(num_requests): |
| 548 | yield ScoringRenderParams( |
| 549 | data_1=data_1[i], |
| 550 | data_2=data_2[i], |
| 551 | chat_template=ctx.chat_template, |
| 552 | max_tokens_per_query=max_tokens_per_query, |
| 553 | max_tokens_per_doc=max_tokens_per_doc, |
| 554 | tok_params=tok_params, |
| 555 | prompt_extras=prompt_extras, |
| 556 | skip_mm_cache=False, |
| 557 | params=ctx.pooling_params, |
| 558 | lora_requests=seq_lora_requests[i], |
| 559 | priorities=seq_priority[i], |
| 560 | ) |
| 561 | |
| 562 | return request_factory, num_requests |
| 563 | |
| 564 | ####################################### |
| 565 | # helpers |
| 566 | |
| 567 | def render( |
| 568 | self, |
| 569 | render_params: EncodeCMPLRenderParams |
| 570 | | EncodeChatRenderParams |
| 571 | | ScoringRenderParams, |
| 572 | ) -> PoolingEngineInput: |
| 573 | if "data_1" not in render_params: |
| 574 | raise ValueError( |
| 575 | f"Unsupported render_params type {render_params.__class__.__name__}" |
| 576 | ) |
| 577 | render_params = cast(ScoringRenderParams, render_params) |
| 578 | |
| 579 | arrival_time = time.time() |
| 580 | |
| 581 | tok_params = render_params["tok_params"] |
| 582 | params = render_params["params"] |
| 583 | prompt_extras = render_params["prompt_extras"] |
| 584 | |
| 585 | _, engine_prompt = self.get_score_prompt( |
| 586 | data_1=render_params["data_1"], |
| 587 | data_2=render_params["data_2"], |
| 588 | encode_kwargs=tok_params.get_encode_kwargs(), |
| 589 | chat_template=render_params["chat_template"], |
| 590 | max_tokens_per_query=render_params["max_tokens_per_query"], |
| 591 | max_tokens_per_doc=render_params["max_tokens_per_doc"], |
| 592 | chat_template_kwargs=prompt_extras.get("chat_template_kwargs") |
| 593 | if prompt_extras |
| 594 | else None, |
| 595 | ) |
| 596 | |
| 597 | tok_params.apply_post_tokenization(self.tokenizer, engine_prompt) |
| 598 | |
| 599 | if token_type_ids := engine_prompt.pop("token_type_ids", None): |
| 600 | params = params.clone() |
| 601 | compressed = compress_token_type_ids( |
| 602 | _apply_post_tokenization_to_token_type_ids( |
| 603 | self.tokenizer, tok_params, token_type_ids |
| 604 | ) |
| 605 | ) |
| 606 | params.extra_kwargs = { |
| 607 | **(params.extra_kwargs or {}), |
| 608 | "compressed_token_type_ids": compressed, |
| 609 | } |
| 610 | |
| 611 | engine_prompt_extras = ( |
| 612 | { |
| 613 | k: v |
| 614 | for k in ("mm_processor_kwargs", "cache_salt") |
| 615 | if (v := prompt_extras.get(k)) is not None |
| 616 | } |
| 617 | if prompt_extras |
| 618 | else None |
| 619 | ) |
| 620 | |
| 621 | if engine_prompt_extras: |
| 622 | target_prompt = extract_target_prompt(self.model_config, engine_prompt) |
| 623 | target_prompt.update(engine_prompt_extras) |
| 624 | |
| 625 | engine_input = self.renderer.process_for_engine(engine_prompt, arrival_time) |
| 626 | |
| 627 | return PoolingEngineInput( |
| 628 | prompts=engine_input, |
| 629 | params=params, |
| 630 | lora_requests=render_params["lora_requests"], |
| 631 | priorities=render_params["priorities"], |
| 632 | ) |
| 633 | |
| 634 | def get_score_prompt( |
| 635 | self, |
| 636 | data_1: ScoreData, |
| 637 | data_2: ScoreData, |
| 638 | encode_kwargs: dict[str, Any], |
| 639 | chat_template: str | None = None, |
| 640 | max_tokens_per_query: int = 0, |
| 641 | max_tokens_per_doc: int = 0, |
| 642 | chat_template_kwargs: dict[str, Any] | None = None, |
| 643 | ): |
| 644 | model_config = self.model_config |
| 645 | tokenizer = self.tokenizer |
| 646 | |
| 647 | prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data( |
| 648 | data_1, |
| 649 | data_2, |
| 650 | model_config, |
| 651 | ) |
| 652 | |
| 653 | # Limit the query and document separately before composing them. |
| 654 | if max_tokens_per_query > 0 and isinstance(prompt_1, str): |
| 655 | prompt_1 = truncate_text_to_tokens( |
| 656 | prompt_1, tokenizer, max_tokens_per_query |
| 657 | ) |
| 658 | if max_tokens_per_doc > 0 and isinstance(prompt_2, str): |
| 659 | prompt_2 = truncate_text_to_tokens(prompt_2, tokenizer, max_tokens_per_doc) |
| 660 | |
| 661 | def default_tokenizer_encode(): |
| 662 | local_kwargs = encode_kwargs.copy() |
| 663 | |
| 664 | if self.supports_score_template: |
| 665 | assert self.model is not None |
| 666 | full_prompt = self.model.get_score_template(prompt_1, prompt_2) |
| 667 | if full_prompt is None: |
| 668 | raise ValueError("Get empty score template from model") |
| 669 | |
| 670 | prompt_inputs = tokenizer(full_prompt, **local_kwargs) |
| 671 | else: |
| 672 | if self.use_sep_token: |
| 673 | # cross_encoder models defaults to using separating token. |
| 674 | if max_tokens_per_doc > 0 and isinstance(prompt_2, str): |
| 675 | query_tokens = tokenizer.encode( |
| 676 | prompt_1, add_special_tokens=False |
| 677 | ) |
| 678 | num_special = get_num_special_tokens_for_pair(tokenizer) |
| 679 | doc_limit_max_length = ( |
| 680 | len(query_tokens) + max_tokens_per_doc + num_special |
| 681 | ) |
| 682 | existing_max_length = local_kwargs.get("max_length") |
| 683 | if existing_max_length is not None: |
| 684 | effective_max_length = min( |
| 685 | doc_limit_max_length, existing_max_length |
| 686 | ) |
| 687 | else: |
| 688 | effective_max_length = doc_limit_max_length |
| 689 | local_kwargs["truncation"] = "only_second" |
| 690 | local_kwargs["max_length"] = effective_max_length |
| 691 | |
| 692 | prompt_inputs = tokenizer( |
| 693 | text=prompt_1, text_pair=prompt_2, **local_kwargs |
| 694 | ) |
| 695 | full_prompt = tokenizer.decode(prompt_inputs["input_ids"]) |
| 696 | else: |
| 697 | # `llm as reranker` defaults to not using separating token. |
| 698 | full_prompt = prompt_1 + prompt_2 |
| 699 | prompt_inputs = tokenizer(text=full_prompt, **local_kwargs) |
| 700 | return full_prompt, prompt_inputs |
| 701 | |
| 702 | # FIXME: For now, we only apply a template when one is explicitly provided. |
| 703 | # We cannot rely on the tokenizer's chat template because many models |
| 704 | # inherit junk templates from their base LLM, which breaks both the models |
| 705 | # and the tests that use them. |
| 706 | if chat_template is None: |
| 707 | full_prompt, prompt_inputs = default_tokenizer_encode() |
| 708 | else: |
| 709 | # FIXME: |
| 710 | # Try applying a score template from the CLI arg or tokenizer_config.json |
| 711 | # If that fails because there is no such template, |
| 712 | # fall back to the default implementation. |
| 713 | try: |
| 714 | _safe_kwargs = chat_template_kwargs or {} |
| 715 | _reserved = {"chat_template", "tools", "tokenize"} |
| 716 | _unexpected = _reserved & _safe_kwargs.keys() |
| 717 | if _unexpected: |
| 718 | raise ValueError( |
| 719 | "chat_template_kwargs contains reserved keys that " |
| 720 | f"conflict with fixed scorer arguments: {_unexpected}" |
| 721 | ) |
| 722 | full_prompt = safe_apply_chat_template( |
| 723 | model_config, |
| 724 | tokenizer, |
| 725 | [ |
| 726 | {"role": "query", "content": prompt_1}, |
| 727 | {"role": "document", "content": prompt_2}, |
| 728 | ], |
| 729 | chat_template=chat_template, |
| 730 | tools=None, |
| 731 | tokenize=False, |
| 732 | **_safe_kwargs, |
| 733 | ) |
| 734 | prompt_inputs = tokenizer(full_prompt, **encode_kwargs) |
| 735 | except ChatTemplateResolutionError: |
| 736 | full_prompt, prompt_inputs = default_tokenizer_encode() |
| 737 | |
| 738 | engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"]) |
| 739 | |
| 740 | if (token_type_ids := prompt_inputs.get("token_type_ids")) is not None: |
| 741 | engine_prompt["token_type_ids"] = token_type_ids |
| 742 | |
| 743 | if self.model is not None: |
| 744 | self.model.post_process_tokens(engine_prompt) |
| 745 | |
| 746 | if mm_data is not None: |
| 747 | engine_prompt["multi_modal_data"] = mm_data |
| 748 | if mm_uuids is not None: |
| 749 | engine_prompt["multi_modal_uuids"] = mm_uuids |
| 750 | |
| 751 | return full_prompt, engine_prompt |
| 752 | |
| 753 | |
| 754 | class JinaRankingIOProcessorMixin: |
| 755 | @staticmethod |
| 756 | def format_docs_prompts_func( |
| 757 | query: str, |
| 758 | docs: list[str], |
| 759 | special_tokens: dict[str, str] | None = None, |
| 760 | instruction: str | None = None, |
| 761 | no_thinking: bool = True, |
| 762 | ) -> str: |
| 763 | # TODO: Try converting the code below into a chat template. |
| 764 | |
| 765 | default_special_tokens = { |
| 766 | "query_embed_token": "<|rerank_token|>", |
| 767 | "doc_embed_token": "<|embed_token|>", |
| 768 | } |
| 769 | if special_tokens is None: |
| 770 | special_tokens = default_special_tokens |
| 771 | |
| 772 | def sanitize_input(text: str) -> str: |
| 773 | for token in special_tokens.values(): |
| 774 | text = text.replace(token, "") |
| 775 | return text |
| 776 | |
| 777 | query = sanitize_input(query) |
| 778 | docs = [sanitize_input(doc) for doc in docs] |
| 779 | |
| 780 | prefix = ( |
| 781 | "<|im_start|>system\n" |
| 782 | "You are a search relevance expert who can determine a ranking of the passages based on how relevant they are to the query. " # noqa: E501 |
| 783 | "If the query is a question, how relevant a passage is depends on how well it answers the question. " # noqa: E501 |
| 784 | "If not, try to analyze the intent of the query and assess how well each passage satisfies the intent. " # noqa: E501 |
| 785 | "If an instruction is provided, you should follow the instruction when determining the ranking." # noqa: E501 |
| 786 | "<|im_end|>\n<|im_start|>user\n" |
| 787 | ) |
| 788 | suffix = "<|im_end|>\n<|im_start|>assistant\n" |
| 789 | if no_thinking: |
| 790 | suffix += "<think>\n\n</think>\n\n" |
| 791 | |
| 792 | doc_emb_token = special_tokens["doc_embed_token"] |
| 793 | query_emb_token = special_tokens["query_embed_token"] |
| 794 | |
| 795 | prompt = ( |
| 796 | f"I will provide you with {len(docs)} passages, each indicated by a numerical identifier. " # noqa: E501 |
| 797 | f"Rank the passages based on their relevance to query: {query}\n" |
| 798 | ) |
| 799 | |
| 800 | if instruction: |
| 801 | instruction = sanitize_input(instruction) |
| 802 | prompt += f"<instruct>\n{instruction}\n</instruct>\n" |
| 803 | |
| 804 | doc_prompts = [ |
| 805 | f'<passage id="{i}">\n{doc}{doc_emb_token}\n</passage>' |
| 806 | for i, doc in enumerate(docs) |
| 807 | ] |
| 808 | prompt += "\n".join(doc_prompts) + "\n" |
| 809 | prompt += f"<query>\n{query}{query_emb_token}\n</query>" |
| 810 | |
| 811 | return prefix + prompt + suffix |
| 812 | |
| 813 | @staticmethod |
| 814 | def ensure_str(data: Sequence[Any]) -> list[str]: |
| 815 | text: list[str] = [] |
| 816 | for prompt in data: |
| 817 | if not isinstance(prompt, str): |
| 818 | raise ValueError( |
| 819 | "The JinaForRanking model only supports text as input." |
| 820 | ) |
| 821 | text.append(prompt) |
| 822 | return text |
| 823 | |
| 824 | |
| 825 | class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorMixin): |
| 826 | name = "jina-reranking-scoring" |
| 827 | pooling_task: PoolingTask = "token_embed" |
| 828 | |
| 829 | def get_request_factory_online( |
| 830 | self, ctx: PoolingServeContext |
| 831 | ) -> Sequence[AnyRenderParam]: |
| 832 | request = ctx.request |
| 833 | ctx.n_queries = 1 |
| 834 | |
| 835 | prompt_extras = ctx.prompt_extras |
| 836 | scoring_data = self.valid_inputs_online(request) |
| 837 | |
| 838 | max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( |
| 839 | request=request |
| 840 | ) |
| 841 | |
| 842 | if max_tokens_per_query > 0 or max_tokens_per_doc > 0: |
| 843 | scoring_data = self._truncate_scoring_data( |
| 844 | scoring_data, max_tokens_per_query, max_tokens_per_doc |
| 845 | ) |
| 846 | |
| 847 | queries = self.ensure_str(scoring_data.data_1) |
| 848 | docs = self.ensure_str(scoring_data.data_2) |
| 849 | |
| 850 | chat_template_kwargs = ( |
| 851 | prompt_extras.get("chat_template_kwargs") if prompt_extras else None |
| 852 | ) |
| 853 | instruction = ( |
| 854 | chat_template_kwargs.get("instruction") if chat_template_kwargs else None |
| 855 | ) |
| 856 | |
| 857 | if len(queries) == 1: |
| 858 | prompts = [ |
| 859 | self.format_docs_prompts_func( |
| 860 | query=queries[0], docs=docs, instruction=instruction |
| 861 | ) |
| 862 | ] |
| 863 | else: |
| 864 | prompts = [ |
| 865 | self.format_docs_prompts_func( |
| 866 | query=q, docs=[d], instruction=instruction |
| 867 | ) |
| 868 | for q, d in zip(queries, docs) |
| 869 | ] |
| 870 | |
| 871 | # Forward truncation from the real request: the base factory reads |
| 872 | # these off ctx.request, so omitting them here silently drops |
| 873 | # truncate_prompt_tokens for Jina rerank/score (unlike the embed and |
| 874 | # bi/cross-encoder paths, which read them from the real request). |
| 875 | ctx.request = PoolingCompletionRequest( |
| 876 | task="token_embed", |
| 877 | input=prompts, |
| 878 | truncate_prompt_tokens=request.truncate_prompt_tokens, |
| 879 | truncation_side=request.truncation_side, |
| 880 | ) |
| 881 | requests = PoolingIOProcessor.get_request_factory_online(self, ctx) |
| 882 | ctx.request = request |
| 883 | return requests |
| 884 | |
| 885 | def get_request_factory_offline( |
| 886 | self, ctx: AnyOfflineInputsContext |
| 887 | ) -> tuple[RequestFactory, int]: |
| 888 | assert isinstance(ctx, OfflineScoringInputsContext) |
| 889 | |
| 890 | max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( |
| 891 | pooling_params=ctx.pooling_params |
| 892 | ) |
| 893 | scoring_data = ctx.scoring_data |
| 894 | if max_tokens_per_query > 0 or max_tokens_per_doc > 0: |
| 895 | scoring_data = self._truncate_scoring_data( |
| 896 | scoring_data, max_tokens_per_query, max_tokens_per_doc |
| 897 | ) |
| 898 | prompt_extras = ctx.pooling_params.extra_kwargs |
| 899 | |
| 900 | queries = self.ensure_str(scoring_data.data_1) |
| 901 | docs = self.ensure_str(scoring_data.data_2) |
| 902 | chat_template_kwargs = ( |
| 903 | prompt_extras.get("chat_template_kwargs") if prompt_extras else None |
| 904 | ) |
| 905 | instruction = ( |
| 906 | chat_template_kwargs.get("instruction") if chat_template_kwargs else None |
| 907 | ) |
| 908 | |
| 909 | if len(queries) == 1: |
| 910 | prompts = [ |
| 911 | self.format_docs_prompts_func( |
| 912 | query=queries[0], docs=docs, instruction=instruction |
| 913 | ) |
| 914 | ] |
| 915 | else: |
| 916 | prompts = [ |
| 917 | self.format_docs_prompts_func( |
| 918 | query=q, docs=[d], instruction=instruction |
| 919 | ) |
| 920 | for q, d in zip(queries, docs) |
| 921 | ] |
| 922 | |
| 923 | return PoolingIOProcessor.get_request_factory_offline( |
| 924 | self, |
| 925 | OfflineEncodeInputsContext( |
| 926 | pooling_task=self.pooling_task, |
| 927 | prompts=prompts, |
| 928 | tokenization_kwargs=ctx.tokenization_kwargs, |
| 929 | pooling_params=ctx.pooling_params, |
| 930 | lora_request=ctx.lora_request, |
| 931 | priorities=ctx.priorities, |
| 932 | ), |
| 933 | ) |
| 934 | |
| 935 | def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): |
| 936 | final_res_batch: list[PoolingRequestOutput] = [] |
| 937 | |
| 938 | for i in range(len(outputs)): |
| 939 | if outputs[i].error is not None: |
| 940 | final_res_batch.append(outputs[i]) |
| 941 | continue |
| 942 | |
| 943 | embeds = outputs[i].outputs.data.float() |
| 944 | |
| 945 | # The JinaForRanking model concatenates docs first, then query. |
| 946 | # Let's stay consistent with this novel design. |
| 947 | query_embeds = embeds[-1] |
| 948 | doc_embeds = embeds[:-1] |
| 949 | |
| 950 | scores = F.cosine_similarity(query_embeds, doc_embeds) |
| 951 | |
| 952 | for score in scores: |
| 953 | final_res_batch.append( |
| 954 | PoolingRequestOutput( |
| 955 | request_id=outputs[i].request_id, |
| 956 | outputs=score, |
| 957 | prompt_token_ids=outputs[i].prompt_token_ids, |
| 958 | num_cached_tokens=outputs[i].num_cached_tokens, |
| 959 | finished=True, |
| 960 | ) |
| 961 | ) |
| 962 | return final_res_batch |
| 963 | |
| 964 | |
| 965 | ScoringIOProcessors: dict[str, type[ScoringIOProcessor]] = { |
| 966 | p.name: p |
| 967 | for p in [ |
| 968 | BiEncoderIOProcessor, |
| 969 | LateInteractionIOProcessor, |
| 970 | JinaRankingIOProcessor, |
| 971 | FlashLateInteractionIOProcessor, |
| 972 | CrossEncoderIOProcessor, |
| 973 | ] |
| 974 | } |