[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>

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.
 

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

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.

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: