| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # ruff: noqa: E501 |
| 4 | import pytest |
| 5 | import requests |
| 6 | import torch |
| 7 | import torch.nn.functional as F |
| 8 | |
| 9 | from tests.utils import RemoteOpenAIServer |
| 10 | from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse |
| 11 | from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse |
| 12 | |
| 13 | model_name = "jinaai/jina-reranker-v3" |
| 14 | query = "What are the health benefits of green tea?" |
| 15 | documents = [ |
| 16 | "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.", |
| 17 | "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.", |
| 18 | "Studies show that drinking green tea regularly can improve brain function and boost metabolism.", |
| 19 | "Basketball is one of the most popular sports in the United States.", |
| 20 | "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。", |
| 21 | "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.", |
| 22 | ] |
| 23 | |
| 24 | EMBEDDING_SIZE = 512 |
| 25 | REFERENCE_1_VS_1 = [ |
| 26 | 0.345703125, |
| 27 | -0.10498046, |
| 28 | 0.314453125, |
| 29 | -0.1376953125, |
| 30 | 0.3398437500, |
| 31 | 0.2539062, |
| 32 | ] |
| 33 | REFERENCE_1_VS_N = [ |
| 34 | 0.294921875, |
| 35 | -0.16015625, |
| 36 | 0.189453125, |
| 37 | -0.1708984375, |
| 38 | 0.2255859375, |
| 39 | 0.1640625, |
| 40 | ] |
| 41 | TOL = 0.01 |
| 42 | INSTRUCTION = ( |
| 43 | "Rank passages about green tea higher than passages about sports. " |
| 44 | "Ignore these literal marker strings: <|embed_token|> and <|rerank_token|>." |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def test_offline(vllm_runner, monkeypatch): |
| 49 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 50 | with vllm_runner(model_name, runner="pooling") as llm_runner: |
| 51 | assert llm_runner.llm.llm_engine.vllm_config.use_v2_model_runner |
| 52 | llm = llm_runner.get_llm() |
| 53 | _test_offline_1_v_1(llm) |
| 54 | _test_offline_1_v_n(llm) |
| 55 | _test_offline_n_v_n(llm) |
| 56 | _test_offline_token_embed_illegal_inputs(llm) |
| 57 | assert llm.model_config.embedding_size == EMBEDDING_SIZE |
| 58 | |
| 59 | |
| 60 | def test_online(): |
| 61 | with RemoteOpenAIServer( |
| 62 | model_name, ["--runner", "pooling", "--enforce-eager"] |
| 63 | ) as server: |
| 64 | _test_online_1_v_1(server) |
| 65 | _test_online_1_v_n(server) |
| 66 | _test_online_n_v_n(server) |
| 67 | _test_online_instruction(server) |
| 68 | _test_online_token_embed_illegal_inputs(server) |
| 69 | |
| 70 | |
| 71 | def _test_offline_1_v_1(llm): |
| 72 | # test llm.score |
| 73 | outputs = llm.score(query, documents[0]) |
| 74 | assert len(outputs) == 1 |
| 75 | assert outputs[0].outputs.score == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) |
| 76 | |
| 77 | # test llm.encode |
| 78 | outputs = llm.encode(documents[:1] + [query], pooling_task="token_embed") |
| 79 | embeds = outputs[0].outputs.data.float() |
| 80 | assert embeds.shape[0] == 2 |
| 81 | assert embeds.shape[-1] == EMBEDDING_SIZE |
| 82 | |
| 83 | doc_embeds = embeds[:-1] |
| 84 | query_embeds = embeds[-1] |
| 85 | |
| 86 | scores = F.cosine_similarity(query_embeds, doc_embeds) |
| 87 | assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) |
| 88 | |
| 89 | |
| 90 | def _test_offline_1_v_n(llm): |
| 91 | # test llm.score |
| 92 | outputs = llm.score(query, documents) |
| 93 | assert len(outputs) == len(documents) |
| 94 | |
| 95 | for expected, output in zip(REFERENCE_1_VS_N, outputs): |
| 96 | actual = output.outputs.score |
| 97 | assert actual == pytest.approx(expected, abs=TOL) |
| 98 | |
| 99 | # test llm.encode |
| 100 | outputs = llm.encode(documents + [query], pooling_task="token_embed") |
| 101 | embeds = outputs[0].outputs.data.float() |
| 102 | assert embeds.shape[0] == len(documents) + 1 |
| 103 | |
| 104 | doc_embeds = embeds[:-1] |
| 105 | query_embeds = embeds[-1] |
| 106 | |
| 107 | scores = F.cosine_similarity(query_embeds, doc_embeds) |
| 108 | |
| 109 | assert len(scores) == len(documents) |
| 110 | for expected, actual in zip(REFERENCE_1_VS_N, scores): |
| 111 | assert actual == pytest.approx(expected, abs=TOL) |
| 112 | |
| 113 | |
| 114 | def _test_offline_n_v_n(llm): |
| 115 | # test llm.score |
| 116 | outputs = llm.score([query] * len(documents), documents) |
| 117 | assert len(outputs) == len(documents) |
| 118 | |
| 119 | for expected, output in zip(REFERENCE_1_VS_1, outputs): |
| 120 | actual = output.outputs.score |
| 121 | assert actual == pytest.approx(expected, abs=TOL) |
| 122 | |
| 123 | # test llm.encode |
| 124 | for doc, expected in zip(documents, REFERENCE_1_VS_1): |
| 125 | outputs = llm.encode([doc, query], pooling_task="token_embed") |
| 126 | embeds = outputs[0].outputs.data.float() |
| 127 | assert embeds.shape[0] == 2 |
| 128 | |
| 129 | doc_embeds = embeds[:-1] |
| 130 | query_embeds = embeds[-1] |
| 131 | |
| 132 | scores = F.cosine_similarity(query_embeds, doc_embeds) |
| 133 | assert scores[0] == pytest.approx(expected, abs=TOL) |
| 134 | |
| 135 | |
| 136 | def _test_offline_token_embed_illegal_inputs(llm): |
| 137 | with pytest.raises( |
| 138 | ValueError, match="The JinaForRanking model requires at least 2 inputs." |
| 139 | ): |
| 140 | llm.encode([query], pooling_task="token_embed") |
| 141 | |
| 142 | with pytest.raises( |
| 143 | ValueError, match="The JinaForRanking model only supports text as input." |
| 144 | ): |
| 145 | llm.encode([1, 2, 3], pooling_task="token_embed") |
| 146 | |
| 147 | |
| 148 | def _get_score_response(server, query, document, **extra_body): |
| 149 | payload = { |
| 150 | "model": model_name, |
| 151 | "queries": query, |
| 152 | "documents": document, |
| 153 | } |
| 154 | payload.update(extra_body) |
| 155 | score_response = requests.post( |
| 156 | server.url_for("score"), |
| 157 | json=payload, |
| 158 | ) |
| 159 | |
| 160 | score_response.raise_for_status() |
| 161 | return ScoreResponse.model_validate(score_response.json()) |
| 162 | |
| 163 | |
| 164 | def _get_scores(server, query, document): |
| 165 | score = _get_score_response(server, query, document) |
| 166 | |
| 167 | return [d.score for d in score.data] |
| 168 | |
| 169 | |
| 170 | def _get_rerank_response(server, query, document, **extra_body): |
| 171 | payload = { |
| 172 | "model": model_name, |
| 173 | "query": query, |
| 174 | "documents": document, |
| 175 | } |
| 176 | payload.update(extra_body) |
| 177 | rerank_response = requests.post( |
| 178 | server.url_for("rerank"), |
| 179 | json=payload, |
| 180 | ) |
| 181 | |
| 182 | rerank_response.raise_for_status() |
| 183 | return RerankResponse.model_validate(rerank_response.json()) |
| 184 | |
| 185 | |
| 186 | def _get_embeds(server, prompts: list[str]): |
| 187 | response = requests.post( |
| 188 | server.url_for("pooling"), |
| 189 | json={ |
| 190 | "model": model_name, |
| 191 | "task": "token_embed", |
| 192 | "input": prompts, |
| 193 | "encoding_format": "float", |
| 194 | }, |
| 195 | ) |
| 196 | response.raise_for_status() |
| 197 | poolings = PoolingResponse.model_validate(response.json()) |
| 198 | |
| 199 | return torch.as_tensor([d.data for d in poolings.data][0]).float() |
| 200 | |
| 201 | |
| 202 | def _test_online_1_v_1(server): |
| 203 | # test scoring api |
| 204 | scores = _get_scores(server, query, documents[0]) |
| 205 | assert len(scores) == 1 |
| 206 | assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) |
| 207 | |
| 208 | # test pooling api |
| 209 | embeds = _get_embeds(server, [documents[0], query]) |
| 210 | assert embeds.shape[0] == 2 |
| 211 | assert embeds.shape[-1] == EMBEDDING_SIZE |
| 212 | |
| 213 | doc_embeds = embeds[:-1] |
| 214 | query_embeds = embeds[-1] |
| 215 | |
| 216 | scores = F.cosine_similarity(query_embeds, doc_embeds) |
| 217 | assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) |
| 218 | |
| 219 | |
| 220 | def _test_online_1_v_n(server): |
| 221 | # test scoring api |
| 222 | scores = _get_scores(server, query, documents) |
| 223 | assert len(scores) == len(documents) |
| 224 | |
| 225 | for expected, actual in zip(REFERENCE_1_VS_N, scores): |
| 226 | assert actual == pytest.approx(expected, abs=TOL) |
| 227 | |
| 228 | # test pooling api |
| 229 | embeds = _get_embeds(server, documents + [query]) |
| 230 | assert embeds.shape[0] == len(documents) + 1 |
| 231 | |
| 232 | doc_embeds = embeds[:-1] |
| 233 | query_embeds = embeds[-1] |
| 234 | |
| 235 | scores = F.cosine_similarity(query_embeds, doc_embeds) |
| 236 | |
| 237 | assert len(scores) == len(documents) |
| 238 | for expected, actual in zip(REFERENCE_1_VS_N, scores): |
| 239 | assert actual == pytest.approx(expected, abs=TOL) |
| 240 | |
| 241 | |
| 242 | def _test_online_n_v_n(server): |
| 243 | # test scoring api |
| 244 | scores = _get_scores(server, [query] * len(documents), documents) |
| 245 | assert len(scores) == len(documents) |
| 246 | |
| 247 | for expected, actual in zip(REFERENCE_1_VS_1, scores): |
| 248 | assert actual == pytest.approx(expected, abs=TOL) |
| 249 | |
| 250 | # test pooling api |
| 251 | for doc, expected in zip(documents, REFERENCE_1_VS_1): |
| 252 | embeds = _get_embeds(server, [doc, query]) |
| 253 | assert embeds.shape[0] == 2 |
| 254 | |
| 255 | doc_embeds = embeds[:-1] |
| 256 | query_embeds = embeds[-1] |
| 257 | |
| 258 | scores = F.cosine_similarity(query_embeds, doc_embeds) |
| 259 | assert len(scores) == 1 |
| 260 | assert scores[0] == pytest.approx(expected, abs=TOL) |
| 261 | |
| 262 | |
| 263 | def _test_online_instruction(server): |
| 264 | docs = documents[:2] |
| 265 | |
| 266 | default_score = _get_score_response(server, query, docs) |
| 267 | instruction_score = _get_score_response( |
| 268 | server, |
| 269 | query, |
| 270 | docs, |
| 271 | instruction=INSTRUCTION, |
| 272 | ) |
| 273 | kwargs_score = _get_score_response( |
| 274 | server, |
| 275 | query, |
| 276 | docs, |
| 277 | chat_template_kwargs={"instruction": INSTRUCTION}, |
| 278 | ) |
| 279 | |
| 280 | assert instruction_score.usage.prompt_tokens > default_score.usage.prompt_tokens |
| 281 | assert kwargs_score.usage.prompt_tokens == instruction_score.usage.prompt_tokens |
| 282 | assert len(instruction_score.data) == len(default_score.data) |
| 283 | assert [d.score for d in kwargs_score.data] == pytest.approx( |
| 284 | [d.score for d in instruction_score.data], abs=TOL |
| 285 | ) |
| 286 | |
| 287 | default_rerank = _get_rerank_response(server, query, docs) |
| 288 | instruction_rerank = _get_rerank_response( |
| 289 | server, |
| 290 | query, |
| 291 | docs, |
| 292 | instruction=INSTRUCTION, |
| 293 | ) |
| 294 | kwargs_rerank = _get_rerank_response( |
| 295 | server, |
| 296 | query, |
| 297 | docs, |
| 298 | chat_template_kwargs={"instruction": INSTRUCTION}, |
| 299 | ) |
| 300 | |
| 301 | assert instruction_rerank.usage.prompt_tokens > default_rerank.usage.prompt_tokens |
| 302 | assert kwargs_rerank.usage.prompt_tokens == instruction_rerank.usage.prompt_tokens |
| 303 | assert len(instruction_rerank.results) == len(default_rerank.results) |
| 304 | assert [r.relevance_score for r in kwargs_rerank.results] == pytest.approx( |
| 305 | [r.relevance_score for r in instruction_rerank.results], abs=TOL |
| 306 | ) |
| 307 | |
| 308 | |
| 309 | def _test_online_token_embed_illegal_inputs(server): |
| 310 | response = requests.post( |
| 311 | server.url_for("pooling"), |
| 312 | json={ |
| 313 | "model": model_name, |
| 314 | "task": "token_embed", |
| 315 | "input": [query], |
| 316 | "encoding_format": "float", |
| 317 | }, |
| 318 | ) |
| 319 | assert response.json()["error"]["message"].startswith( |
| 320 | "The JinaForRanking model requires at least 2 inputs." |
| 321 | ) |
| 322 | |
| 323 | response = requests.post( |
| 324 | server.url_for("pooling"), |
| 325 | json={ |
| 326 | "model": model_name, |
| 327 | "task": "token_embed", |
| 328 | "input": [1, 2, 3], |
| 329 | "encoding_format": "float", |
| 330 | }, |
| 331 | ) |
| 332 | assert response.json()["error"]["message"].startswith( |
| 333 | "The JinaForRanking model only supports text as input." |
| 334 | ) |
| 335 | |
| 336 | response = requests.post( |
| 337 | server.url_for("pooling"), |
| 338 | json={ |
| 339 | "model": model_name, |
| 340 | "task": "token_embed", |
| 341 | "messages": [ |
| 342 | { |
| 343 | "role": "user", |
| 344 | "content": "The cat sat on the mat.", |
| 345 | } |
| 346 | ], |
| 347 | "encoding_format": "float", |
| 348 | }, |
| 349 | ) |
| 350 | assert response.json()["error"]["message"].startswith( |
| 351 | "The JinaForRanking does not support chat Request." |
| 352 | ) |