From a2bc2ffb2c8bf46759c4bbeb88e2f4dc32968762 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Fri, 11 Sep 2026 03:26:44 -0700 Subject: [PATCH] [Bugfix][Pooling] Restore token limits for offline Jina scoring (#56415) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- .../test_jina_ranking_io_processor_unit.py | 87 ++++++++++++++++++- .../pooling/scoring/io_processor.py | 7 ++ 2 files changed, 93 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py index c2fe62f609..cb64524699 100644 --- a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py @@ -1,15 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for JinaRankingIOProcessor online request building.""" +"""Unit tests for JinaRankingIOProcessor request building.""" +from types import SimpleNamespace from unittest.mock import MagicMock import pytest +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast +from vllm import PoolingParams from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor from vllm.entrypoints.pooling.scoring.protocol import RerankRequest from vllm.entrypoints.pooling.scoring.typing import ScoringData +from vllm.entrypoints.pooling.typing import OfflineScoringInputsContext +from vllm.renderers import TokenizeParams pytestmark = pytest.mark.skip_global_cleanup @@ -57,3 +63,82 @@ def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch): assert captured["truncation_side"] == "left" # The real request is restored after delegating. assert ctx.request is request + + +@pytest.fixture +def offline_processor_and_context(): + backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + backend.pre_tokenizer = pre_tokenizers.WhitespaceSplit() + proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor) + proc.tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend) + proc.model_config = SimpleNamespace(max_model_len=1024, is_encoder_decoder=False) + proc.renderer = SimpleNamespace( + default_cmpl_tok_params=TokenizeParams(max_total_tokens=1024) + ) + ctx = OfflineScoringInputsContext( + pooling_task="token_embed", + scoring_data=ScoringData( + data_1=["query alpha beta", "query gamma delta"], + data_2=["document one two", "document three four"], + ), + pooling_params=PoolingParams( + extra_kwargs={ + "chat_template_kwargs": {"instruction": "Keep this instruction"} + } + ), + tokenization_kwargs=None, + chat_template=None, + lora_request=None, + priorities=None, + ) + return proc, ctx + + +@pytest.mark.parametrize( + ("query_limit", "doc_limit", "n_queries"), + [(None, None, 1), (0, 0, 2), (1, None, 1), (None, 2, 1), (1, 2, 1), (1, 2, 2)], + ids=["unset", "zero", "query", "doc", "both-1-to-n", "both-n-to-n"], +) +def test_offline_truncates_text_before_formatting( + offline_processor_and_context, query_limit, doc_limit, n_queries +): + proc, ctx = offline_processor_and_context + queries = ctx.scoring_data.data_1[:n_queries] + docs = ctx.scoring_data.data_2 + ctx.scoring_data = ScoringData(data_1=queries, data_2=docs) + ctx.pooling_params.extra_kwargs.update( + { + name: limit + for name, limit in ( + ("max_tokens_per_query", query_limit), + ("max_tokens_per_doc", doc_limit), + ) + if limit is not None + } + ) + + factory, num_requests = proc.get_request_factory_offline(ctx) + requests = list(factory()) + assert len(requests) == num_requests == n_queries + for i, request in enumerate(requests): + prompt = request["prompts"]["prompt"] + query = " ".join(queries[i].split()[: query_limit or None]) + assert f"\n{query}<|rerank_token|>\n" in prompt + prompt_docs = docs if n_queries == 1 else [docs[i]] + for j, doc in enumerate(prompt_docs): + doc = " ".join(doc.split()[: doc_limit or None]) + assert f'\n{doc}<|embed_token|>\n' in prompt + assert prompt.count("<|embed_token|>") == len(prompt_docs) + assert prompt.count("<|rerank_token|>") == 1 + assert "\nKeep this instruction\n" in prompt + + +@pytest.mark.parametrize("name", ["max_tokens_per_query", "max_tokens_per_doc"]) +@pytest.mark.parametrize("limit", [-1, 1024]) +def test_offline_rejects_invalid_token_limits( + offline_processor_and_context, name, limit +): + proc, ctx = offline_processor_and_context + ctx.pooling_params.extra_kwargs[name] = limit + with pytest.raises(ValueError, match=name): + proc.get_request_factory_offline(ctx) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 4feb40eadb..f8cdda20d5 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -845,7 +845,14 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM ) -> tuple[RequestFactory, int]: assert isinstance(ctx, OfflineScoringInputsContext) + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + pooling_params=ctx.pooling_params + ) scoring_data = ctx.scoring_data + if max_tokens_per_query > 0 or max_tokens_per_doc > 0: + scoring_data = self._truncate_scoring_data( + scoring_data, max_tokens_per_query, max_tokens_per_doc + ) prompt_extras = ctx.pooling_params.extra_kwargs queries = self.ensure_str(scoring_data.data_1)