[Pooling] Preserve reranker tokenization with document limits (#57666)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Taneem Ibrahim committed
Sep 25, 2026 at 10:04 UTC
bb35e2587bb5ca3f149b718744949fd10210e124
4 files changed
+94
-18
tests/entrypoints/pooling/scoring/test_io_processor_unit.py
+51
@@ -3,11 +3,15 @@
3
"""Unit tests for ScoringIOProcessor post-tokenization helpers."""
4
5
from dataclasses import dataclass
6
+from types import SimpleNamespace
7
8
import pytest
9
+from tokenizers import Tokenizer, models, pre_tokenizers, processors
10
+from transformers import TokenizersBackend
11
12
from vllm import TokensPrompt
13
from vllm.entrypoints.pooling.scoring.io_processor import (
14
+ CrossEncoderIOProcessor,
15
_apply_post_tokenization_to_token_type_ids,
16
)
17
from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids
@@ -24,6 +28,53 @@ class _DummyTokenizer:
28
pad_token_id: int = 99999
29
30
31
+@pytest.fixture
32
+def llm_reranker_processor() -> CrossEncoderIOProcessor:
33
+ vocab = {
34
+ "[UNK]": 0,
35
+ "[CLS]": 1,
36
+ "[SEP]": 2,
37
+ "a": 3,
38
+ "bc": 4,
39
+ "##b": 5,
40
+ "##c": 6,
41
+ }
42
+ backend = Tokenizer(models.WordPiece(vocab, unk_token="[UNK]"))
43
+ backend.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
44
+ backend.post_processor = processors.TemplateProcessing(
45
+ single="[CLS] $A [SEP]",
46
+ special_tokens=[("[CLS]", 1), ("[SEP]", 2)],
47
+ )
48
+
49
+ processor = CrossEncoderIOProcessor.__new__(CrossEncoderIOProcessor)
50
+ processor.model_config = SimpleNamespace(enable_prompt_embeds=False)
51
+ processor.tokenizer = TokenizersBackend(
52
+ tokenizer_object=backend,
53
+ unk_token="[UNK]",
54
+ cls_token="[CLS]",
55
+ sep_token="[SEP]",
56
+ )
57
+ processor.supports_score_template = False
58
+ processor.use_sep_token = False
59
+ processor.model = None
60
+ return processor
61
+
62
+
63
+def test_llm_reranker_tokenization_is_independent_of_nonbinding_doc_limit(
64
+ llm_reranker_processor: CrossEncoderIOProcessor,
65
+):
66
+ encode_kwargs = {"add_special_tokens": True}
67
+
68
+ _, uncapped = llm_reranker_processor.get_score_prompt("a", "bc", encode_kwargs)
69
+ _, nonbinding = llm_reranker_processor.get_score_prompt(
70
+ "a", "bc", encode_kwargs, max_tokens_per_doc=10
71
+ )
72
+
73
+ expected = [1, 3, 5, 6, 2]
74
+ assert uncapped["prompt_token_ids"] == expected
75
+ assert nonbinding["prompt_token_ids"] == expected
76
+
77
+
78
def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt():
79
"""The cross-encoder segment boundary must survive truncate + pad.
80
tests/entrypoints/pooling/scoring/test_utils.py
+20
-1
@@ -2,9 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import pytest
5
+from tokenizers import Tokenizer, models, pre_tokenizers
6
+from transformers import TokenizersBackend
7
8
from vllm.entrypoints.pooling.scoring.typing import ScoreInput
7
-from vllm.entrypoints.pooling.scoring.utils import validate_score_input
9
+from vllm.entrypoints.pooling.scoring.utils import (
10
+ truncate_text_to_tokens,
11
+ validate_score_input,
12
+)
13
from vllm.exceptions import VLLMValidationError
14
15
@@ -63,3 +68,17 @@ def test_validate_score_input_rejects_invalid_inputs(
68
assert str(exc_info.value) == message
69
assert exc_info.value.parameter is None
70
assert exc_info.value.value is None
71
+
72
+
73
+def test_truncate_text_to_tokens_handles_shared_character_offsets():
74
+ vocab = {
75
+ token: i for i, token in enumerate(sorted(pre_tokenizers.ByteLevel.alphabet()))
76
+ }
77
+ backend = Tokenizer(models.BPE(vocab=vocab, merges=[]))
78
+ backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
79
+ tokenizer = TokenizersBackend(tokenizer_object=backend)
80
+
81
+ truncated = truncate_text_to_tokens("aéx", tokenizer, max_tokens=2)
82
+
83
+ assert truncated == "a"
84
+ assert len(tokenizer(truncated, add_special_tokens=False)["input_ids"]) == 1
vllm/entrypoints/pooling/scoring/io_processor.py
+3
-11
@@ -621,7 +621,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
621
model_config,
622
)
623
624
- # Apply truncation before defining closures
624
+ # Limit the query and document separately before composing them.
625
if max_tokens_per_query > 0 and isinstance(prompt_1, str):
626
prompt_1 = truncate_text_to_tokens(
627
prompt_1, tokenizer, max_tokens_per_query
@@ -666,16 +666,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
666
full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
667
else:
668
# `llm as reranker` defaults to not using separating token.
669
- if max_tokens_per_doc > 0 and isinstance(prompt_2, str):
670
- query_ids = tokenizer.encode(prompt_1, add_special_tokens=False)
671
- doc_ids = tokenizer.encode(prompt_2, add_special_tokens=False)
672
- doc_ids = doc_ids[:max_tokens_per_doc]
673
- input_ids = query_ids + doc_ids
674
- full_prompt = tokenizer.decode(input_ids)
675
- prompt_inputs = {"input_ids": input_ids}
676
- else:
677
- full_prompt = prompt_1 + prompt_2
678
- prompt_inputs = tokenizer(text=full_prompt, **local_kwargs)
669
+ full_prompt = prompt_1 + prompt_2
670
+ prompt_inputs = tokenizer(text=full_prompt, **local_kwargs)
671
return full_prompt, prompt_inputs
672
673
# FIXME: For now, we only apply a template when one is explicitly provided.
vllm/entrypoints/pooling/scoring/utils.py
+20
-6
@@ -48,13 +48,27 @@ def truncate_text_to_tokens(
48
49
Uses offset_mapping to slice the original text at the exact character
50
boundary, avoiding lossy encode→decode round-trips that can shift
51
- the token count by 1-3 tokens due to BPE merge boundary changes.
51
+ the token count due to BPE merge boundary changes.
52
"""
53
- encoding = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
54
- if len(encoding["input_ids"]) <= max_tokens:
55
- return text
56
- char_end = encoding["offset_mapping"][max_tokens - 1][1]
57
- return text[:char_end]
53
+ if max_tokens <= 0:
54
+ return ""
55
+
56
+ while text:
57
+ encoding = tokenizer(
58
+ text,
59
+ add_special_tokens=False,
60
+ return_offsets_mapping=True,
61
+ )
62
+ offsets = encoding["offset_mapping"]
63
+ if len(encoding["input_ids"]) <= max_tokens:
64
+ break
65
+
66
+ char_end = min(offsets[max_tokens - 1][1], offsets[max_tokens][0])
67
+ if char_end >= len(text):
68
+ char_end = len(text) - 1
69
+ text = text[:char_end]
70
+
71
+ return text
72
73
74
def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor: