From bb35e2587bb5ca3f149b718744949fd10210e124 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 25 Sep 2026 10:04:41 -0500 Subject: [PATCH] [Pooling] Preserve reranker tokenization with document limits (#57666) Signed-off-by: Taneem Ibrahim Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- .../pooling/scoring/test_io_processor_unit.py | 51 +++++++++++++++++++ .../entrypoints/pooling/scoring/test_utils.py | 21 +++++++- .../pooling/scoring/io_processor.py | 14 ++--- vllm/entrypoints/pooling/scoring/utils.py | 26 +++++++--- 4 files changed, 94 insertions(+), 18 deletions(-) diff --git a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py index 8c2cf40f20..c83eaa544a 100644 --- a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py @@ -3,11 +3,15 @@ """Unit tests for ScoringIOProcessor post-tokenization helpers.""" from dataclasses import dataclass +from types import SimpleNamespace import pytest +from tokenizers import Tokenizer, models, pre_tokenizers, processors +from transformers import TokenizersBackend from vllm import TokensPrompt from vllm.entrypoints.pooling.scoring.io_processor import ( + CrossEncoderIOProcessor, _apply_post_tokenization_to_token_type_ids, ) from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids @@ -24,6 +28,53 @@ class _DummyTokenizer: pad_token_id: int = 99999 +@pytest.fixture +def llm_reranker_processor() -> CrossEncoderIOProcessor: + vocab = { + "[UNK]": 0, + "[CLS]": 1, + "[SEP]": 2, + "a": 3, + "bc": 4, + "##b": 5, + "##c": 6, + } + backend = Tokenizer(models.WordPiece(vocab, unk_token="[UNK]")) + backend.pre_tokenizer = pre_tokenizers.WhitespaceSplit() + backend.post_processor = processors.TemplateProcessing( + single="[CLS] $A [SEP]", + special_tokens=[("[CLS]", 1), ("[SEP]", 2)], + ) + + processor = CrossEncoderIOProcessor.__new__(CrossEncoderIOProcessor) + processor.model_config = SimpleNamespace(enable_prompt_embeds=False) + processor.tokenizer = TokenizersBackend( + tokenizer_object=backend, + unk_token="[UNK]", + cls_token="[CLS]", + sep_token="[SEP]", + ) + processor.supports_score_template = False + processor.use_sep_token = False + processor.model = None + return processor + + +def test_llm_reranker_tokenization_is_independent_of_nonbinding_doc_limit( + llm_reranker_processor: CrossEncoderIOProcessor, +): + encode_kwargs = {"add_special_tokens": True} + + _, uncapped = llm_reranker_processor.get_score_prompt("a", "bc", encode_kwargs) + _, nonbinding = llm_reranker_processor.get_score_prompt( + "a", "bc", encode_kwargs, max_tokens_per_doc=10 + ) + + expected = [1, 3, 5, 6, 2] + assert uncapped["prompt_token_ids"] == expected + assert nonbinding["prompt_token_ids"] == expected + + def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt(): """The cross-encoder segment boundary must survive truncate + pad. diff --git a/tests/entrypoints/pooling/scoring/test_utils.py b/tests/entrypoints/pooling/scoring/test_utils.py index 35b205eebf..ddd3c082be 100644 --- a/tests/entrypoints/pooling/scoring/test_utils.py +++ b/tests/entrypoints/pooling/scoring/test_utils.py @@ -2,9 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import TokenizersBackend from vllm.entrypoints.pooling.scoring.typing import ScoreInput -from vllm.entrypoints.pooling.scoring.utils import validate_score_input +from vllm.entrypoints.pooling.scoring.utils import ( + truncate_text_to_tokens, + validate_score_input, +) from vllm.exceptions import VLLMValidationError @@ -63,3 +68,17 @@ def test_validate_score_input_rejects_invalid_inputs( assert str(exc_info.value) == message assert exc_info.value.parameter is None assert exc_info.value.value is None + + +def test_truncate_text_to_tokens_handles_shared_character_offsets(): + vocab = { + token: i for i, token in enumerate(sorted(pre_tokenizers.ByteLevel.alphabet())) + } + backend = Tokenizer(models.BPE(vocab=vocab, merges=[])) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + tokenizer = TokenizersBackend(tokenizer_object=backend) + + truncated = truncate_text_to_tokens("aéx", tokenizer, max_tokens=2) + + assert truncated == "a" + assert len(tokenizer(truncated, add_special_tokens=False)["input_ids"]) == 1 diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 5e6104eedd..8ffa2e4ddd 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -621,7 +621,7 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): model_config, ) - # Apply truncation before defining closures + # Limit the query and document separately before composing them. if max_tokens_per_query > 0 and isinstance(prompt_1, str): prompt_1 = truncate_text_to_tokens( prompt_1, tokenizer, max_tokens_per_query @@ -666,16 +666,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): full_prompt = tokenizer.decode(prompt_inputs["input_ids"]) else: # `llm as reranker` defaults to not using separating token. - if max_tokens_per_doc > 0 and isinstance(prompt_2, str): - query_ids = tokenizer.encode(prompt_1, add_special_tokens=False) - doc_ids = tokenizer.encode(prompt_2, add_special_tokens=False) - doc_ids = doc_ids[:max_tokens_per_doc] - input_ids = query_ids + doc_ids - full_prompt = tokenizer.decode(input_ids) - prompt_inputs = {"input_ids": input_ids} - else: - full_prompt = prompt_1 + prompt_2 - prompt_inputs = tokenizer(text=full_prompt, **local_kwargs) + full_prompt = prompt_1 + prompt_2 + prompt_inputs = tokenizer(text=full_prompt, **local_kwargs) return full_prompt, prompt_inputs # FIXME: For now, we only apply a template when one is explicitly provided. diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index cff8ff61d9..d7cc888a15 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -48,13 +48,27 @@ def truncate_text_to_tokens( Uses offset_mapping to slice the original text at the exact character boundary, avoiding lossy encode→decode round-trips that can shift - the token count by 1-3 tokens due to BPE merge boundary changes. + the token count due to BPE merge boundary changes. """ - encoding = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) - if len(encoding["input_ids"]) <= max_tokens: - return text - char_end = encoding["offset_mapping"][max_tokens - 1][1] - return text[:char_end] + if max_tokens <= 0: + return "" + + while text: + encoding = tokenizer( + text, + add_special_tokens=False, + return_offsets_mapping=True, + ) + offsets = encoding["offset_mapping"] + if len(encoding["input_ids"]) <= max_tokens: + break + + char_end = min(offsets[max_tokens - 1][1], offsets[max_tokens][0]) + if char_end >= len(text): + char_end = len(text) - 1 + text = text[:char_end] + + return text def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor: