| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for max_tokens_per_doc and max_tokens_per_query.""" |
| 4 | |
| 5 | import json |
| 6 | import os |
| 7 | from dataclasses import dataclass |
| 8 | |
| 9 | import pytest |
| 10 | import requests |
| 11 | |
| 12 | from tests.utils import VLLM_PATH, RemoteOpenAIServer |
| 13 | from vllm.entrypoints.pooling.scoring.protocol import RerankResponse |
| 14 | |
| 15 | TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template") |
| 16 | ExpectedPromptTokens = int | tuple[int, ...] |
| 17 | |
| 18 | long_query = "What is the capital of France?" * 20 |
| 19 | long_doc = "The capital of France is Paris. " * 20 |
| 20 | |
| 21 | |
| 22 | @dataclass |
| 23 | class TestConfig: |
| 24 | model: str |
| 25 | args: list[str] |
| 26 | without_truncated_prompt_tokens: ExpectedPromptTokens |
| 27 | with_max_tokens_per_query_prompt_tokens: ExpectedPromptTokens |
| 28 | with_max_tokens_per_doc_prompt_tokens: ExpectedPromptTokens |
| 29 | with_max_tokens_per_query_and_doc_prompt_tokens: ExpectedPromptTokens |
| 30 | |
| 31 | |
| 32 | RERANK_CONFIGS = [ |
| 33 | # 1. cross-encoder |
| 34 | TestConfig( |
| 35 | model="jinaai/jina-reranker-v2-base-multilingual", |
| 36 | args=[ |
| 37 | "--enforce-eager", |
| 38 | "--max-model-len", |
| 39 | "1024", |
| 40 | "--trust-remote-code", |
| 41 | ], |
| 42 | without_truncated_prompt_tokens=284, |
| 43 | with_max_tokens_per_query_prompt_tokens=154, |
| 44 | with_max_tokens_per_doc_prompt_tokens=154, |
| 45 | with_max_tokens_per_query_and_doc_prompt_tokens=24, |
| 46 | ), |
| 47 | # 2. cross-encoder + score template |
| 48 | TestConfig( |
| 49 | model="Qwen/Qwen3-Reranker-0.6B", |
| 50 | args=[ |
| 51 | "--enforce-eager", |
| 52 | "--max-model-len", |
| 53 | "1024", |
| 54 | "--hf-overrides", |
| 55 | json.dumps( |
| 56 | { |
| 57 | "architectures": ["Qwen3ForSequenceClassification"], |
| 58 | "classifier_from_token": ["no", "yes"], |
| 59 | "is_original_qwen3_reranker": True, |
| 60 | } |
| 61 | ), |
| 62 | "--chat-template", |
| 63 | os.path.join(TEMPLATE_DIR, "qwen3_reranker.jinja"), |
| 64 | ], |
| 65 | without_truncated_prompt_tokens=352, |
| 66 | with_max_tokens_per_query_prompt_tokens=223, |
| 67 | with_max_tokens_per_doc_prompt_tokens=221, |
| 68 | with_max_tokens_per_query_and_doc_prompt_tokens=92, |
| 69 | ), |
| 70 | # 3. bi-encoder |
| 71 | TestConfig( |
| 72 | model="intfloat/multilingual-e5-small", |
| 73 | args=[ |
| 74 | "--enforce-eager", |
| 75 | "--max-model-len", |
| 76 | "512", |
| 77 | "--trust-remote-code", |
| 78 | ], |
| 79 | # This model has produced both prompt-token totals in CI/local cache; |
| 80 | # keep truncation checks exact while tolerating the boundary delta. |
| 81 | without_truncated_prompt_tokens=(284, 285), |
| 82 | with_max_tokens_per_query_prompt_tokens=(154, 155), |
| 83 | with_max_tokens_per_doc_prompt_tokens=154, |
| 84 | with_max_tokens_per_query_and_doc_prompt_tokens=24, |
| 85 | ), |
| 86 | # 4. late-interaction |
| 87 | TestConfig( |
| 88 | model="answerdotai/answerai-colbert-small-v1", |
| 89 | args=[ |
| 90 | "--enforce-eager", |
| 91 | "--max-model-len", |
| 92 | "512", |
| 93 | "--trust-remote-code", |
| 94 | ], |
| 95 | without_truncated_prompt_tokens=284, |
| 96 | with_max_tokens_per_query_prompt_tokens=154, |
| 97 | with_max_tokens_per_doc_prompt_tokens=154, |
| 98 | with_max_tokens_per_query_and_doc_prompt_tokens=24, |
| 99 | ), |
| 100 | # 5. jinaai/jina-reranker-v3 |
| 101 | TestConfig( |
| 102 | model="jinaai/jina-reranker-v3", |
| 103 | args=[ |
| 104 | "--enforce-eager", |
| 105 | "--max-model-len", |
| 106 | "1024", |
| 107 | "--trust-remote-code", |
| 108 | ], |
| 109 | without_truncated_prompt_tokens=567, |
| 110 | with_max_tokens_per_query_prompt_tokens=308, |
| 111 | with_max_tokens_per_doc_prompt_tokens=436, |
| 112 | with_max_tokens_per_query_and_doc_prompt_tokens=177, |
| 113 | ), |
| 114 | ] |
| 115 | |
| 116 | |
| 117 | def assert_prompt_tokens(actual: int, expected: ExpectedPromptTokens) -> None: |
| 118 | if isinstance(expected, int): |
| 119 | assert actual == expected |
| 120 | else: |
| 121 | assert actual in expected |
| 122 | |
| 123 | |
| 124 | @pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model) |
| 125 | def server(request): |
| 126 | config: TestConfig = request.param |
| 127 | env_dict = {"VLLM_LOGGING_LEVEL": "WARNING"} |
| 128 | with RemoteOpenAIServer( |
| 129 | config.model, config.args, env_dict=env_dict |
| 130 | ) as remote_server: |
| 131 | yield config, remote_server |
| 132 | |
| 133 | |
| 134 | def test_without_truncated(server): |
| 135 | """Test that max_tokens_per_doc truncates documents correctly.""" |
| 136 | config, remote_server = server |
| 137 | |
| 138 | response = requests.post( |
| 139 | remote_server.url_for("rerank"), |
| 140 | json={"model": config.model, "query": long_query, "documents": [long_doc]}, |
| 141 | ) |
| 142 | response.raise_for_status() |
| 143 | rerank = RerankResponse.model_validate(response.json()) |
| 144 | |
| 145 | assert rerank.id is not None |
| 146 | assert rerank.results is not None |
| 147 | assert len(rerank.results) == 1 |
| 148 | assert_prompt_tokens( |
| 149 | rerank.usage.prompt_tokens, |
| 150 | config.without_truncated_prompt_tokens, |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | def test_max_tokens_per_query(server): |
| 155 | """Test that max_tokens_per_doc truncates documents correctly.""" |
| 156 | config, remote_server = server |
| 157 | |
| 158 | response = requests.post( |
| 159 | remote_server.url_for("rerank"), |
| 160 | json={ |
| 161 | "model": config.model, |
| 162 | "query": long_query, |
| 163 | "documents": [long_doc], |
| 164 | "max_tokens_per_query": 10, |
| 165 | }, |
| 166 | ) |
| 167 | response.raise_for_status() |
| 168 | rerank = RerankResponse.model_validate(response.json()) |
| 169 | |
| 170 | assert rerank.id is not None |
| 171 | assert rerank.results is not None |
| 172 | assert len(rerank.results) == 1 |
| 173 | assert_prompt_tokens( |
| 174 | rerank.usage.prompt_tokens, |
| 175 | config.with_max_tokens_per_query_prompt_tokens, |
| 176 | ) |
| 177 | |
| 178 | |
| 179 | def test_max_tokens_per_doc(server): |
| 180 | """Test that max_tokens_per_doc truncates documents correctly.""" |
| 181 | config, remote_server = server |
| 182 | |
| 183 | response = requests.post( |
| 184 | remote_server.url_for("rerank"), |
| 185 | json={ |
| 186 | "model": config.model, |
| 187 | "query": long_query, |
| 188 | "documents": [long_doc], |
| 189 | "max_tokens_per_doc": 10, |
| 190 | }, |
| 191 | ) |
| 192 | response.raise_for_status() |
| 193 | rerank = RerankResponse.model_validate(response.json()) |
| 194 | |
| 195 | assert rerank.id is not None |
| 196 | assert rerank.results is not None |
| 197 | assert len(rerank.results) == 1 |
| 198 | assert_prompt_tokens( |
| 199 | rerank.usage.prompt_tokens, |
| 200 | config.with_max_tokens_per_doc_prompt_tokens, |
| 201 | ) |
| 202 | |
| 203 | |
| 204 | def test_max_tokens_per_query_and_doc(server): |
| 205 | """Test that max_tokens_per_doc truncates documents correctly.""" |
| 206 | config, remote_server = server |
| 207 | |
| 208 | response = requests.post( |
| 209 | remote_server.url_for("rerank"), |
| 210 | json={ |
| 211 | "model": config.model, |
| 212 | "query": long_query, |
| 213 | "documents": [long_doc], |
| 214 | "max_tokens_per_query": 10, |
| 215 | "max_tokens_per_doc": 10, |
| 216 | }, |
| 217 | ) |
| 218 | response.raise_for_status() |
| 219 | rerank = RerankResponse.model_validate(response.json()) |
| 220 | |
| 221 | assert rerank.id is not None |
| 222 | assert rerank.results is not None |
| 223 | assert len(rerank.results) == 1 |
| 224 | assert_prompt_tokens( |
| 225 | rerank.usage.prompt_tokens, |
| 226 | config.with_max_tokens_per_query_and_doc_prompt_tokens, |
| 227 | ) |