[Bugfix][Pooling] Restore token limits for offline Jina scoring (#56415)
Signed-off-by: 子华 <huaxi.shx@alibaba-inc.com> Co-authored-by: Codex <noreply@openai.com>
shaohuaxi committed
Sep 11, 2026 at 03:26 UTC
a2bc2ffb2c8bf46759c4bbeb88e2f4dc32968762
2 files changed
+93
-1
tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py
+86
-1
@@ -1,15 +1,21 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Unit tests for JinaRankingIOProcessor online request building."""
3
+"""Unit tests for JinaRankingIOProcessor request building."""
4
5
+from types import SimpleNamespace
6
from unittest.mock import MagicMock
7
8
import pytest
9
+from tokenizers import Tokenizer, models, pre_tokenizers
10
+from transformers import PreTrainedTokenizerFast
11
12
+from vllm import PoolingParams
13
from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
14
from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor
15
from vllm.entrypoints.pooling.scoring.protocol import RerankRequest
16
from vllm.entrypoints.pooling.scoring.typing import ScoringData
17
+from vllm.entrypoints.pooling.typing import OfflineScoringInputsContext
18
+from vllm.renderers import TokenizeParams
19
20
pytestmark = pytest.mark.skip_global_cleanup
21
@@ -57,3 +63,82 @@ def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch):
63
assert captured["truncation_side"] == "left"
64
# The real request is restored after delegating.
65
assert ctx.request is request
66
+
67
+
68
+@pytest.fixture
69
+def offline_processor_and_context():
70
+ backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]"))
71
+ backend.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
72
+ proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)
73
+ proc.tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend)
74
+ proc.model_config = SimpleNamespace(max_model_len=1024, is_encoder_decoder=False)
75
+ proc.renderer = SimpleNamespace(
76
+ default_cmpl_tok_params=TokenizeParams(max_total_tokens=1024)
77
+ )
78
+ ctx = OfflineScoringInputsContext(
79
+ pooling_task="token_embed",
80
+ scoring_data=ScoringData(
81
+ data_1=["query alpha beta", "query gamma delta"],
82
+ data_2=["document one two", "document three four"],
83
+ ),
84
+ pooling_params=PoolingParams(
85
+ extra_kwargs={
86
+ "chat_template_kwargs": {"instruction": "Keep this instruction"}
87
+ }
88
+ ),
89
+ tokenization_kwargs=None,
90
+ chat_template=None,
91
+ lora_request=None,
92
+ priorities=None,
93
+ )
94
+ return proc, ctx
95
+
96
+
97
+@pytest.mark.parametrize(
98
+ ("query_limit", "doc_limit", "n_queries"),
99
+ [(None, None, 1), (0, 0, 2), (1, None, 1), (None, 2, 1), (1, 2, 1), (1, 2, 2)],
100
+ ids=["unset", "zero", "query", "doc", "both-1-to-n", "both-n-to-n"],
101
+)
102
+def test_offline_truncates_text_before_formatting(
103
+ offline_processor_and_context, query_limit, doc_limit, n_queries
104
+):
105
+ proc, ctx = offline_processor_and_context
106
+ queries = ctx.scoring_data.data_1[:n_queries]
107
+ docs = ctx.scoring_data.data_2
108
+ ctx.scoring_data = ScoringData(data_1=queries, data_2=docs)
109
+ ctx.pooling_params.extra_kwargs.update(
110
+ {
111
+ name: limit
112
+ for name, limit in (
113
+ ("max_tokens_per_query", query_limit),
114
+ ("max_tokens_per_doc", doc_limit),
115
+ )
116
+ if limit is not None
117
+ }
118
+ )
119
+
120
+ factory, num_requests = proc.get_request_factory_offline(ctx)
121
+ requests = list(factory())
122
+ assert len(requests) == num_requests == n_queries
123
+ for i, request in enumerate(requests):
124
+ prompt = request["prompts"]["prompt"]
125
+ query = " ".join(queries[i].split()[: query_limit or None])
126
+ assert f"<query>\n{query}<|rerank_token|>\n</query>" in prompt
127
+ prompt_docs = docs if n_queries == 1 else [docs[i]]
128
+ for j, doc in enumerate(prompt_docs):
129
+ doc = " ".join(doc.split()[: doc_limit or None])
130
+ assert f'<passage id="{j}">\n{doc}<|embed_token|>\n</passage>' in prompt
131
+ assert prompt.count("<|embed_token|>") == len(prompt_docs)
132
+ assert prompt.count("<|rerank_token|>") == 1
133
+ assert "<instruct>\nKeep this instruction\n</instruct>" in prompt
134
+
135
+
136
+@pytest.mark.parametrize("name", ["max_tokens_per_query", "max_tokens_per_doc"])
137
+@pytest.mark.parametrize("limit", [-1, 1024])
138
+def test_offline_rejects_invalid_token_limits(
139
+ offline_processor_and_context, name, limit
140
+):
141
+ proc, ctx = offline_processor_and_context
142
+ ctx.pooling_params.extra_kwargs[name] = limit
143
+ with pytest.raises(ValueError, match=name):
144
+ proc.get_request_factory_offline(ctx)
vllm/entrypoints/pooling/scoring/io_processor.py
+7
@@ -845,7 +845,14 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
845
) -> tuple[RequestFactory, int]:
846
assert isinstance(ctx, OfflineScoringInputsContext)
847
848
+ max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
849
+ pooling_params=ctx.pooling_params
850
+ )
851
scoring_data = ctx.scoring_data
852
+ if max_tokens_per_query > 0 or max_tokens_per_doc > 0:
853
+ scoring_data = self._truncate_scoring_data(
854
+ scoring_data, max_tokens_per_query, max_tokens_per_doc
855
+ )
856
prompt_extras = ctx.pooling_params.extra_kwargs
857
858
queries = self.ensure_str(scoring_data.data_1)