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)