@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cb5f7501cb
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 07:15 UTC

[New Model]: jinaai/jina-reranker-v3 (#38800)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

wang.yuqi committed Apr 10, 2026 at 23:20 UTC cb5f7501cbc80a394afbce0c877593be8cf6175c
14 files changed +660 -19
docs/models/pooling_models/token_embed.md
+8
@@ -71,6 +71,14 @@ Models of any architecture can be converted into embedding models using `--conve
71
72 If your model is not in the above list, we will try to automatically convert the model using [as_embedding_model][vllm.model_executor.models.adapters.as_embedding_model].
73
74 +### Special models
75 +
76 +| Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) |
77 +| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
78 +| `JinaForRanking` | Qwen3-based | `jinaai/jina-reranker-v3` | | |
79 +
80 +jina-reranker-v3 is a listwise document reranker model with a novel `last but not late interaction` architecture. More information can be found at: [examples/pooling/token_embed/jina_reranker_v3_offline.py](../../../examples/pooling/token_embed/jina_reranker_v3_offline.py)
81 +
82 --8<-- [end:supported-token-embed-models]
83
84 ## Offline Inference
examples/pooling/token_embed/jina_reranker_v3_offline.py new
+56
@@ -0,0 +1,56 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# ruff: noqa: E501
4 +
5 +import torch.nn.functional as F
6 +
7 +from vllm import LLM
8 +
9 +query = "What are the health benefits of green tea?"
10 +documents = [
11 + "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.",
12 + "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.",
13 + "Studies show that drinking green tea regularly can improve brain function and boost metabolism.",
14 + "Basketball is one of the most popular sports in the United States.",
15 + "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。",
16 + "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.",
17 +]
18 +
19 +
20 +def main():
21 + # Initialize model
22 + llm = LLM(
23 + model="jinaai/jina-reranker-v3",
24 + runner="pooling",
25 + )
26 +
27 + # Generate scores.
28 + outputs = llm.score(query, documents)
29 +
30 + # Print the outputs.
31 + print("\nGenerated Outputs:\n" + "-" * 60)
32 + for document, output in zip(documents, outputs):
33 + score = output.outputs.score
34 + print(f"Pair: {[query, document]!r} \nScore: {score}")
35 + print("-" * 60)
36 +
37 + # Generate embeddings.
38 + # The JinaForRanking model concatenates docs first, then query.
39 + # Let's stay consistent with this novel design.
40 + outputs = llm.encode(documents + [query], pooling_task="token_embed")
41 + embeds = outputs[0].outputs.data.float()
42 +
43 + doc_embeds = embeds[:-1]
44 + query_embeds = embeds[-1]
45 +
46 + scores = F.cosine_similarity(query_embeds, doc_embeds)
47 +
48 + # Print the outputs.
49 + print("\nGenerated Outputs:\n" + "-" * 60)
50 + for document, score in zip(documents, scores):
51 + print(f"Pair: {[query, document]!r} \nScore: {score}")
52 + print("-" * 60)
53 +
54 +
55 +if __name__ == "__main__":
56 + main()
pyproject.toml
+2 -1
@@ -120,7 +120,8 @@ python = "./.venv"
120 [tool.typos.files]
121 # these files may be written in non english words
122 extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*",
123 - "benchmarks/sonnet.txt", "tests/lora/data/*", "examples/pooling/token_embed/*", "build/*",
123 + "benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
124 + "examples/pooling/token_embed/*", "tests/models/language/pooling/*",
125 "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py",
126 "docs/governance/process.md", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"]
127 ignore-hidden = false
tests/models/language/pooling/test_jina_reranker_v3.py new
+275
@@ -0,0 +1,275 @@
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 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 +
43 +
44 +def test_offline(vllm_runner):
45 + with vllm_runner(model_name, runner="pooling") as llm_runner:
46 + llm = llm_runner.get_llm()
47 + _test_offline_1_v_1(llm)
48 + _test_offline_1_v_n(llm)
49 + _test_offline_n_v_n(llm)
50 + _test_offline_token_embed_illegal_inputs(llm)
51 + assert llm.model_config.embedding_size == EMBEDDING_SIZE
52 +
53 +
54 +def test_online():
55 + with RemoteOpenAIServer(model_name, ["--runner", "pooling"]) as server:
56 + _test_online_1_v_1(server)
57 + _test_online_1_v_n(server)
58 + _test_online_n_v_n(server)
59 + _test_online_token_embed_illegal_inputs(server)
60 +
61 +
62 +def _test_offline_1_v_1(llm):
63 + # test llm.score
64 + outputs = llm.score(query, documents[0])
65 + assert len(outputs) == 1
66 + assert outputs[0].outputs.score == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
67 +
68 + # test llm.encode
69 + outputs = llm.encode(documents[:1] + [query], pooling_task="token_embed")
70 + embeds = outputs[0].outputs.data.float()
71 + assert embeds.shape[0] == 2
72 + assert embeds.shape[-1] == EMBEDDING_SIZE
73 +
74 + doc_embeds = embeds[:-1]
75 + query_embeds = embeds[-1]
76 +
77 + scores = F.cosine_similarity(query_embeds, doc_embeds)
78 + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
79 +
80 +
81 +def _test_offline_1_v_n(llm):
82 + # test llm.score
83 + outputs = llm.score(query, documents)
84 + assert len(outputs) == len(documents)
85 +
86 + for expected, output in zip(REFERENCE_1_VS_N, outputs):
87 + actual = output.outputs.score
88 + assert actual == pytest.approx(expected, abs=TOL)
89 +
90 + # test llm.encode
91 + outputs = llm.encode(documents + [query], pooling_task="token_embed")
92 + embeds = outputs[0].outputs.data.float()
93 + assert embeds.shape[0] == len(documents) + 1
94 +
95 + doc_embeds = embeds[:-1]
96 + query_embeds = embeds[-1]
97 +
98 + scores = F.cosine_similarity(query_embeds, doc_embeds)
99 +
100 + assert len(scores) == len(documents)
101 + for expected, actual in zip(REFERENCE_1_VS_N, scores):
102 + assert actual == pytest.approx(expected, abs=TOL)
103 +
104 +
105 +def _test_offline_n_v_n(llm):
106 + # test llm.score
107 + outputs = llm.score([query] * len(documents), documents)
108 + assert len(outputs) == len(documents)
109 +
110 + for expected, output in zip(REFERENCE_1_VS_1, outputs):
111 + actual = output.outputs.score
112 + assert actual == pytest.approx(expected, abs=TOL)
113 +
114 + # test llm.encode
115 + for doc, expected in zip(documents, REFERENCE_1_VS_1):
116 + outputs = llm.encode([doc, query], pooling_task="token_embed")
117 + embeds = outputs[0].outputs.data.float()
118 + assert embeds.shape[0] == 2
119 +
120 + doc_embeds = embeds[:-1]
121 + query_embeds = embeds[-1]
122 +
123 + scores = F.cosine_similarity(query_embeds, doc_embeds)
124 + assert scores[0] == pytest.approx(expected, abs=TOL)
125 +
126 +
127 +def _test_offline_token_embed_illegal_inputs(llm):
128 + with pytest.raises(
129 + ValueError, match="The JinaForRanking model requires at least 2 inputs."
130 + ):
131 + llm.encode([query], pooling_task="token_embed")
132 +
133 + with pytest.raises(
134 + ValueError, match="The JinaForRanking model only supports text as input."
135 + ):
136 + llm.encode([1, 2, 3], pooling_task="token_embed")
137 +
138 +
139 +def _get_scores(server, query, document):
140 + score_response = requests.post(
141 + server.url_for("score"),
142 + json={
143 + "model": model_name,
144 + "queries": query,
145 + "documents": document,
146 + },
147 + )
148 +
149 + score_response.raise_for_status()
150 + score = ScoreResponse.model_validate(score_response.json())
151 +
152 + return [d.score for d in score.data]
153 +
154 +
155 +def _get_embeds(server, prompts: list[str]):
156 + response = requests.post(
157 + server.url_for("pooling"),
158 + json={
159 + "model": model_name,
160 + "task": "token_embed",
161 + "input": prompts,
162 + "encoding_format": "float",
163 + },
164 + )
165 + response.raise_for_status()
166 + poolings = PoolingResponse.model_validate(response.json())
167 +
168 + return torch.as_tensor([d.data for d in poolings.data][0]).float()
169 +
170 +
171 +def _test_online_1_v_1(server):
172 + # test scoring api
173 + scores = _get_scores(server, query, documents[0])
174 + assert len(scores) == 1
175 + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
176 +
177 + # test pooling api
178 + embeds = _get_embeds(server, [documents[0], query])
179 + assert embeds.shape[0] == 2
180 + assert embeds.shape[-1] == EMBEDDING_SIZE
181 +
182 + doc_embeds = embeds[:-1]
183 + query_embeds = embeds[-1]
184 +
185 + scores = F.cosine_similarity(query_embeds, doc_embeds)
186 + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL)
187 +
188 +
189 +def _test_online_1_v_n(server):
190 + # test scoring api
191 + scores = _get_scores(server, query, documents)
192 + assert len(scores) == len(documents)
193 +
194 + for expected, actual in zip(REFERENCE_1_VS_N, scores):
195 + assert actual == pytest.approx(expected, abs=TOL)
196 +
197 + # test pooling api
198 + embeds = _get_embeds(server, documents + [query])
199 + assert embeds.shape[0] == len(documents) + 1
200 +
201 + doc_embeds = embeds[:-1]
202 + query_embeds = embeds[-1]
203 +
204 + scores = F.cosine_similarity(query_embeds, doc_embeds)
205 +
206 + assert len(scores) == len(documents)
207 + for expected, actual in zip(REFERENCE_1_VS_N, scores):
208 + assert actual == pytest.approx(expected, abs=TOL)
209 +
210 +
211 +def _test_online_n_v_n(server):
212 + # test scoring api
213 + scores = _get_scores(server, [query] * len(documents), documents)
214 + assert len(scores) == len(documents)
215 +
216 + for expected, actual in zip(REFERENCE_1_VS_1, scores):
217 + assert actual == pytest.approx(expected, abs=TOL)
218 +
219 + # test pooling api
220 + for doc, expected in zip(documents, REFERENCE_1_VS_1):
221 + embeds = _get_embeds(server, [doc, query])
222 + assert embeds.shape[0] == 2
223 +
224 + doc_embeds = embeds[:-1]
225 + query_embeds = embeds[-1]
226 +
227 + scores = F.cosine_similarity(query_embeds, doc_embeds)
228 + assert len(scores) == 1
229 + assert scores[0] == pytest.approx(expected, abs=TOL)
230 +
231 +
232 +def _test_online_token_embed_illegal_inputs(server):
233 + response = requests.post(
234 + server.url_for("pooling"),
235 + json={
236 + "model": model_name,
237 + "task": "token_embed",
238 + "input": [query],
239 + "encoding_format": "float",
240 + },
241 + )
242 + assert response.json()["error"]["message"].startswith(
243 + "The JinaForRanking model requires at least 2 inputs."
244 + )
245 +
246 + response = requests.post(
247 + server.url_for("pooling"),
248 + json={
249 + "model": model_name,
250 + "task": "token_embed",
251 + "input": [1, 2, 3],
252 + "encoding_format": "float",
253 + },
254 + )
255 + assert response.json()["error"]["message"].startswith(
256 + "The JinaForRanking model only supports text as input."
257 + )
258 +
259 + response = requests.post(
260 + server.url_for("pooling"),
261 + json={
262 + "model": model_name,
263 + "task": "token_embed",
264 + "messages": [
265 + {
266 + "role": "user",
267 + "content": "The cat sat on the mat.",
268 + }
269 + ],
270 + "encoding_format": "float",
271 + },
272 + )
273 + assert response.json()["error"]["message"].startswith(
274 + "The JinaForRanking does not support chat Request."
275 + )
tests/models/registry.py
+1
@@ -645,6 +645,7 @@ _LATE_INTERACTION_EXAMPLE_MODELS = {
645 trust_remote_code=True,
646 hf_overrides={"architectures": ["ColBERTLfm2Model"]},
647 ),
648 + "JinaForRanking": _HfExamplesInfo("jinaai/jina-reranker-v3"),
649 # [Multimodal]
650 "ColModernVBertForRetrieval": _HfExamplesInfo(
651 "ModernVBERT/colmodernvbert-merged",
vllm/entrypoints/pooling/embed/io_processor.py
+52 -1
@@ -24,7 +24,13 @@ from vllm.entrypoints.pooling.embed.protocol import (
24 EmbeddingChatRequest,
25 EmbeddingCompletionRequest,
26 )
27 -from vllm.entrypoints.pooling.typing import PoolingServeContext
27 +from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessorMixin
28 +from vllm.entrypoints.pooling.typing import (
29 + OfflineInputsContext,
30 + PoolingChatLikeRequest,
31 + PoolingCompletionLikeRequest,
32 + PoolingServeContext,
33 +)
34 from vllm.inputs import EngineInput, tokens_input
35 from vllm.logger import init_logger
36 from vllm.outputs import PoolingOutput, PoolingRequestOutput
@@ -553,3 +559,48 @@ class EmbedIOProcessor(PoolingIOProcessor):
559
560 class TokenEmbedIOProcessor(PoolingIOProcessor):
561 name = "token_embed"
562 +
563 +
564 +class JinaRankingTokenEmbedIOProcessor(
565 + TokenEmbedIOProcessor, JinaRankingIOProcessorMixin
566 +):
567 + def pre_process_online(self, ctx: PoolingServeContext):
568 + request = ctx.request
569 + if isinstance(request, PoolingCompletionLikeRequest):
570 + prompts = request.input
571 + if not isinstance(prompts, Sequence) or len(prompts) < 2:
572 + raise ValueError("The JinaForRanking model requires at least 2 inputs.")
573 +
574 + text_prompts = self.ensure_str(prompts)
575 +
576 + # The JinaForRanking model concatenates docs first, then query.
577 + # Let's stay consistent with this novel design.
578 + prompt_input = self.format_docs_prompts_func(
579 + query=text_prompts[-1], docs=text_prompts[:-1]
580 + )
581 +
582 + engine_inputs = self._preprocess_completion_online(
583 + request,
584 + prompt_input=prompt_input,
585 + prompt_embeds=None,
586 + )
587 + elif isinstance(request, PoolingChatLikeRequest):
588 + raise ValueError("The JinaForRanking does not support chat Request.")
589 + else:
590 + raise ValueError(f"Invalid {self.name} request type")
591 +
592 + ctx.engine_inputs = engine_inputs
593 +
594 + def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
595 + if not isinstance(ctx.prompts, Sequence) or len(ctx.prompts) < 2:
596 + raise ValueError("The JinaForRanking model requires at least 2 inputs.")
597 +
598 + text_prompts = self.ensure_str(ctx.prompts)
599 +
600 + # The JinaForRanking model concatenates docs first, then query.
601 + # Let's stay consistent with this novel design.
602 + ctx.prompts = self.format_docs_prompts_func(
603 + query=text_prompts[-1], docs=text_prompts[:-1]
604 + )
605 +
606 + return super().pre_process_offline(ctx)
vllm/entrypoints/pooling/io_processor_factories.py
+7
@@ -59,6 +59,13 @@ def init_pooling_io_processors(
59 if score_type is not None and score_type in ScoringIOProcessors:
60 processors[score_type] = ScoringIOProcessors[score_type]
61
62 + if model_config.architecture == "JinaForRanking":
63 + from .embed.io_processor import JinaRankingTokenEmbedIOProcessor
64 + from .scoring.io_processor import ScoringIOProcessors
65 +
66 + processors["token_embed"] = JinaRankingTokenEmbedIOProcessor
67 + processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"]
68 +
69 return {
70 task: processor_cls(
71 vllm_config=vllm_config,
vllm/entrypoints/pooling/scoring/io_processor.py
+126
@@ -416,11 +416,137 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
416 return full_prompt, engine_prompt
417
418
419 +class JinaRankingIOProcessorMixin:
420 + @staticmethod
421 + def sanitize_input(text: str, special_tokens: dict[str, str]) -> str:
422 + for token in special_tokens.values():
423 + text = text.replace(token, "")
424 + return text
425 +
426 + @staticmethod
427 + def format_docs_prompts_func(
428 + query: str,
429 + docs: list[str],
430 + special_tokens: dict[str, str] | None = None,
431 + instruction: str | None = None,
432 + no_thinking: bool = True,
433 + ) -> str:
434 + # TODO: Try converting the code below into a chat template.
435 +
436 + default_special_tokens = {
437 + "query_embed_token": "<|rerank_token|>",
438 + "doc_embed_token": "<|embed_token|>",
439 + }
440 + if special_tokens is None:
441 + special_tokens = default_special_tokens
442 +
443 + query = JinaRankingIOProcessorMixin.sanitize_input(query, special_tokens)
444 + docs = [
445 + JinaRankingIOProcessorMixin.sanitize_input(doc, special_tokens)
446 + for doc in docs
447 + ]
448 +
449 + prefix = (
450 + "<|im_start|>system\n"
451 + "You are a search relevance expert who can determine a ranking of the passages based on how relevant they are to the query. " # noqa: E501
452 + "If the query is a question, how relevant a passage is depends on how well it answers the question. " # noqa: E501
453 + "If not, try to analyze the intent of the query and assess how well each passage satisfies the intent. " # noqa: E501
454 + "If an instruction is provided, you should follow the instruction when determining the ranking." # noqa: E501
455 + "<|im_end|>\n<|im_start|>user\n"
456 + )
457 + suffix = "<|im_end|>\n<|im_start|>assistant\n"
458 + if no_thinking:
459 + suffix += "<think>\n\n</think>\n\n"
460 +
461 + doc_emb_token = special_tokens["doc_embed_token"]
462 + query_emb_token = special_tokens["query_embed_token"]
463 +
464 + prompt = (
465 + f"I will provide you with {len(docs)} passages, each indicated by a numerical identifier. " # noqa: E501
466 + f"Rank the passages based on their relevance to query: {query}\n"
467 + )
468 +
469 + if instruction:
470 + prompt += f"<instruct>\n{instruction}\n</instruct>\n"
471 +
472 + doc_prompts = [
473 + f'<passage id="{i}">\n{doc}{doc_emb_token}\n</passage>'
474 + for i, doc in enumerate(docs)
475 + ]
476 + prompt += "\n".join(doc_prompts) + "\n"
477 + prompt += f"<query>\n{query}{query_emb_token}\n</query>"
478 +
479 + return prefix + prompt + suffix
480 +
481 + @staticmethod
482 + def ensure_str(data: Sequence[Any]) -> list[str]:
483 + text: list[str] = []
484 + for prompt in data:
485 + if not isinstance(prompt, str):
486 + raise ValueError(
487 + "The JinaForRanking model only supports text as input."
488 + )
489 + text.append(prompt)
490 + return text
491 +
492 +
493 +class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorMixin):
494 + name = "jina-reranking-scoring"
495 + pooling_task: PoolingTask = "token_embed"
496 +
497 + def _pre_process(
498 + self,
499 + scoring_data: ScoringData,
500 + tok_params: TokenizeParams,
501 + prompt_extras: dict[str, Any] | None = None,
502 + ) -> Sequence[EngineInput]:
503 + queries = self.ensure_str(scoring_data.data_1)
504 + docs = self.ensure_str(scoring_data.data_2)
505 +
506 + if len(queries) == 1:
507 + prompts = [self.format_docs_prompts_func(query=queries[0], docs=docs)]
508 + else:
509 + prompts = [
510 + self.format_docs_prompts_func(query=q, docs=[d])
511 + for q, d in zip(queries, docs)
512 + ]
513 +
514 + return self._preprocess_completion_offline(
515 + prompts=prompts, tok_params=tok_params, prompt_extras=prompt_extras
516 + )
517 +
518 + def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
519 + final_res_batch: list[PoolingRequestOutput] = []
520 +
521 + for i in range(len(outputs)):
522 + embeds = outputs[i].outputs.data.float()
523 +
524 + # The JinaForRanking model concatenates docs first, then query.
525 + # Let's stay consistent with this novel design.
526 + query_embeds = embeds[-1]
527 + doc_embeds = embeds[:-1]
528 +
529 + scores = F.cosine_similarity(query_embeds, doc_embeds)
530 +
531 + for score in scores:
532 + final_res_batch.append(
533 + PoolingRequestOutput(
534 + request_id=outputs[i].request_id,
535 + outputs=score,
536 + prompt_token_ids=outputs[i].prompt_token_ids,
537 + num_cached_tokens=outputs[i].num_cached_tokens,
538 + finished=True,
539 + )
540 + )
541 + return final_res_batch
542 +
543 +
544 ScoringIOProcessors: dict[str, type[ScoringIOProcessor]] = {
545 p.name: p
546 for p in [
547 BiEncoderIOProcessor,
548 LateInteractionIOProcessor,
549 + JinaRankingIOProcessor,
550 FlashLateInteractionIOProcessor,
551 CrossEncoderIOProcessor,
552 ]
vllm/entrypoints/pooling/scoring/serving.py
+11 -14
@@ -4,7 +4,6 @@
4 from fastapi.responses import JSONResponse, Response
5
6 from vllm import PoolingParams
7 -from vllm.config import VllmConfig
7 from vllm.engine.protocol import EngineClient
8 from vllm.entrypoints.openai.engine.protocol import UsageInfo
9 from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
@@ -42,25 +41,23 @@ class ServingScores(PoolingServing):
41 enable_flash_late_interaction: bool = True,
42 **kwargs,
43 ):
45 - self.score_type = engine_client.model_config.score_type
44 + self.io_processor_name: str = engine_client.model_config.score_type
45 self.enable_flash_late_interaction = (
47 - self.score_type == "late-interaction" and enable_flash_late_interaction
46 + self.io_processor_name == "late-interaction"
47 + and enable_flash_late_interaction
48 )
49
50 - super().__init__(engine_client, *args, **kwargs)
50 + if self.enable_flash_late_interaction:
51 + self.io_processor_name = "flash-late-interaction"
52
52 - def init_io_processor(
53 - self, vllm_config: VllmConfig, *args, **kwargs
54 - ) -> PoolingIOProcessor:
55 - model_config = vllm_config.model_config
53 + if engine_client.model_config.architecture == "JinaForRanking":
54 + self.io_processor_name = "jina-reranking-scoring"
55 + self.enable_flash_late_interaction = False
56
57 - score_type: str = model_config.score_type
58 - if self.enable_flash_late_interaction:
59 - score_type = "flash-late-interaction"
57 + super().__init__(engine_client, *args, **kwargs)
58
61 - assert score_type in ScoringIOProcessors
62 - processor_cls = ScoringIOProcessors[score_type]
63 - return processor_cls(vllm_config, *args, **kwargs)
59 + def init_io_processor(self, *args, **kwargs) -> PoolingIOProcessor:
60 + return ScoringIOProcessors[self.io_processor_name](*args, **kwargs)
61
62 async def __call__(self, *args, **kwargs) -> Response:
63 if not self.enable_flash_late_interaction:
vllm/model_executor/layers/pooler/tokwise/methods.py
+1 -1
@@ -100,7 +100,7 @@ class StepPool(AllPool):
100 ):
101 # for unfinished chunked prefill
102 if data is None:
103 - pass
103 + pooled_data.append(None)
104 else:
105 step_tag_id = pooling_param.step_tag_id
106 returned_token_ids = pooling_param.returned_token_ids
vllm/model_executor/layers/pooler/tokwise/poolers.py
+3 -2
@@ -58,7 +58,7 @@ class TokenPooler(Pooler):
58 def __init__(
59 self,
60 pooling: TokenPoolingMethod | TokenPoolingFn,
61 - head: TokenPoolerHead | TokenPoolingHeadFn,
61 + head: TokenPoolerHead | TokenPoolingHeadFn | None = None,
62 ) -> None:
63 super().__init__()
64
@@ -89,7 +89,8 @@ class TokenPooler(Pooler):
89 pooling_metadata: PoolingMetadata,
90 ) -> TokenPoolerOutput:
91 pooled_data = self.pooling(hidden_states, pooling_metadata)
92 - pooled_data = self.head(pooled_data, pooling_metadata)
92 + if self.head is not None:
93 + pooled_data = self.head(pooled_data, pooling_metadata)
94 return pooled_data
95
96
vllm/model_executor/models/config.py
+7
@@ -192,6 +192,12 @@ class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
192 pooler_config.use_activation = False
193
194
195 +class JinaForRankingConfig(VerifyAndUpdateConfig):
196 + @staticmethod
197 + def verify_and_update_model_config(model_config: "ModelConfig") -> None:
198 + model_config.hf_config.embedding_size = 512
199 +
200 +
201 class JinaRobertaModelConfig(VerifyAndUpdateConfig):
202 @staticmethod
203 def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -612,6 +618,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
618 "GteNewForSequenceClassification": GteNewModelConfig,
619 "GteNewModel": GteNewModelConfig,
620 "JambaForSequenceClassification": JambaForSequenceClassificationConfig,
621 + "JinaForRanking": JinaForRankingConfig,
622 "JinaVLForRanking": JinaVLForSequenceClassificationConfig,
623 "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
624 "LlamaBidirectionalModel": LlamaBidirectionalConfig,
vllm/model_executor/models/jina.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Adapted from https://huggingface.co/jinaai/jina-reranker-v3/blob/main/modeling.py
4 +from collections.abc import Iterable
5 +
6 +import torch
7 +from torch import nn
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.sequence import IntermediateTensors
11 +from vllm.tasks import PoolingTask
12 +from vllm.v1.pool.metadata import PoolingMetadata
13 +
14 +from ..layers.pooler import DispatchPooler
15 +from ..layers.pooler.tokwise import (
16 + StepPool,
17 + TokenPooler,
18 + TokenPoolingMethodOutputItem,
19 +)
20 +from .interfaces import SupportsLateInteraction
21 +from .qwen3 import Qwen3Model
22 +from .utils import AutoWeightsLoader, maybe_prefix
23 +
24 +
25 +class JinaForRanking(nn.Module, SupportsLateInteraction):
26 + is_pooling_model = True
27 +
28 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
29 + super().__init__()
30 + config = vllm_config.model_config.hf_config
31 + quant_config = vllm_config.quant_config
32 +
33 + self.config = config
34 + self.projector_dim: int = config.embedding_size
35 +
36 + self.vllm_config = vllm_config
37 + self.quant_config = quant_config
38 + self.model = Qwen3Model(
39 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
40 + )
41 +
42 + self.projector = nn.Sequential(
43 + nn.Linear(config.hidden_size, config.hidden_size // 2, bias=False),
44 + nn.ReLU(),
45 + nn.Linear(config.hidden_size // 2, self.projector_dim, bias=False),
46 + )
47 +
48 + self.pooler = DispatchPooler(
49 + {
50 + "token_embed": TokenPooler(
51 + pooling=JinaForRankingPool(self.projector),
52 + )
53 + }
54 + )
55 +
56 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
57 + return self.model.embed_input_ids(input_ids)
58 +
59 + def forward(
60 + self,
61 + input_ids: torch.Tensor | None,
62 + positions: torch.Tensor,
63 + intermediate_tensors: IntermediateTensors | None = None,
64 + inputs_embeds: torch.Tensor | None = None,
65 + ) -> torch.Tensor | IntermediateTensors:
66 + hidden_states = self.model(
67 + input_ids, positions, intermediate_tensors, inputs_embeds
68 + )
69 + return hidden_states
70 +
71 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
72 + loader = AutoWeightsLoader(self, skip_prefixes=(["lm_head."]))
73 + return loader.load_weights(weights)
74 +
75 +
76 +class JinaForRankingPool(StepPool):
77 + def __init__(self, projector: nn.Sequential):
78 + super().__init__()
79 +
80 + self.doc_token_id = 151670
81 + self.query_token_id = 151671
82 + self.projector = projector
83 +
84 + def get_supported_tasks(self) -> set[PoolingTask]:
85 + return {"token_embed"}
86 +
87 + def forward(
88 + self,
89 + hidden_states: torch.Tensor,
90 + pooling_metadata: PoolingMetadata,
91 + ) -> list[TokenPoolingMethodOutputItem]:
92 + pooled_data_lst = super().forward(hidden_states, pooling_metadata)
93 + prompt_token_ids = pooling_metadata.get_prompt_token_ids()
94 +
95 + embeds_list = list[torch.Tensor | None]()
96 + for data, token_ids in zip(pooled_data_lst, prompt_token_ids):
97 + # for unfinished chunked prefill
98 + if data is None:
99 + embeds_list.append(None)
100 + else:
101 + docs_indexes = torch.where(torch.eq(token_ids, self.doc_token_id))[0]
102 + query_indexes = torch.where(torch.eq(token_ids, self.query_token_id))[0]
103 +
104 + # The JinaForRanking model concatenates docs first, then query.
105 + # Let's stay consistent with this novel design.
106 + indexes = torch.cat([docs_indexes, query_indexes])
107 + embeds = self.projector(data[indexes])
108 + embeds_list.append(embeds)
109 +
110 + return embeds_list
vllm/model_executor/models/registry.py
+1
@@ -273,6 +273,7 @@ _LATE_INTERACTION_MODELS = {
273 "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
274 "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
275 "ColBERTLfm2Model": ("colbert", "ColBERTLfm2Model"),
276 + "JinaForRanking": ("jina", "JinaForRanking"),
277 # [Multimodal]
278 "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
279 "ColPaliForRetrieval": ("colpali", "ColPaliModel"),