[Bugfix][Pooling] Forward instruction to Jina reranker scoring prompts (#47590)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Ting SUN committed
Jul 5, 2026 at 12:39 UTC
922661304357313700340e787cfe19d58b8c0a2f
3 files changed
+179
-22
examples/pooling/token_embed/jina_reranker_v3_online.py
new
+73
@@ -0,0 +1,73 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+
5
+"""
6
+Example online usage of the Jina Reranker v3 score and rerank APIs with a task
7
+instruction.
8
+
9
+Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the
10
+server in vLLM.
11
+"""
12
+
13
+import argparse
14
+import json
15
+
16
+import requests
17
+
18
+
19
+def post_http_request(prompt: dict, api_url: str) -> requests.Response:
20
+ headers = {"User-Agent": "Test Client"}
21
+ response = requests.post(api_url, headers=headers, json=prompt)
22
+ return response
23
+
24
+
25
+def print_response(name: str, prompt: dict, response: requests.Response) -> None:
26
+ print(f"\n{name} request:")
27
+ print(json.dumps(prompt, indent=2))
28
+ print(f"\n{name} response:")
29
+ print(json.dumps(response.json(), indent=2))
30
+
31
+
32
+def parse_args():
33
+ parser = argparse.ArgumentParser()
34
+ parser.add_argument("--host", type=str, default="localhost")
35
+ parser.add_argument("--port", type=int, default=8000)
36
+ parser.add_argument("--model", type=str, default="jinaai/jina-reranker-v3")
37
+ return parser.parse_args()
38
+
39
+
40
+def main(args):
41
+ score_url = f"http://{args.host}:{args.port}/score"
42
+ rerank_url = f"http://{args.host}:{args.port}/rerank"
43
+ model_name = args.model
44
+
45
+ query = "Which passage is about sports?"
46
+ documents = [
47
+ "Basketball is played by two teams on a court.",
48
+ "Green tea contains antioxidants and may support metabolism.",
49
+ ]
50
+ instruction = "Rank passages about sports higher than passages about nutrition."
51
+
52
+ score_prompt = {
53
+ "model": model_name,
54
+ "queries": query,
55
+ "documents": documents,
56
+ "instruction": instruction,
57
+ }
58
+ score_response = post_http_request(prompt=score_prompt, api_url=score_url)
59
+ print_response("Score", score_prompt, score_response)
60
+
61
+ rerank_prompt = {
62
+ "model": model_name,
63
+ "query": query,
64
+ "documents": documents,
65
+ "instruction": instruction,
66
+ }
67
+ rerank_response = post_http_request(prompt=rerank_prompt, api_url=rerank_url)
68
+ print_response("Rerank", rerank_prompt, rerank_response)
69
+
70
+
71
+if __name__ == "__main__":
72
+ args = parse_args()
73
+ main(args)
tests/models/language/pooling/test_jina_reranker_v3.py
+84
-9
@@ -8,7 +8,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
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?"
@@ -39,6 +39,10 @@ REFERENCE_1_VS_N = [
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):
@@ -52,10 +56,13 @@ def test_offline(vllm_runner):
56
57
58
def test_online():
55
- with RemoteOpenAIServer(model_name, ["--runner", "pooling"]) as server:
59
+ with RemoteOpenAIServer(
60
+ model_name, ["--runner", "pooling", "--enforce-eager"]
61
+ ) as server:
62
_test_online_1_v_1(server)
63
_test_online_1_v_n(server)
64
_test_online_n_v_n(server)
65
+ _test_online_instruction(server)
66
_test_online_token_embed_illegal_inputs(server)
67
68
@@ -136,22 +143,44 @@ def _test_offline_token_embed_illegal_inputs(llm):
143
llm.encode([1, 2, 3], pooling_task="token_embed")
144
145
139
-def _get_scores(server, query, document):
146
+def _get_score_response(server, query, document, **extra_body):
147
+ payload = {
148
+ "model": model_name,
149
+ "queries": query,
150
+ "documents": document,
151
+ }
152
+ payload.update(extra_body)
153
score_response = requests.post(
154
server.url_for("score"),
142
- json={
143
- "model": model_name,
144
- "queries": query,
145
- "documents": document,
146
- },
155
+ json=payload,
156
)
157
158
score_response.raise_for_status()
150
- score = ScoreResponse.model_validate(score_response.json())
159
+ return ScoreResponse.model_validate(score_response.json())
160
+
161
+
162
+def _get_scores(server, query, document):
163
+ score = _get_score_response(server, query, document)
164
165
return [d.score for d in score.data]
166
167
168
+def _get_rerank_response(server, query, document, **extra_body):
169
+ payload = {
170
+ "model": model_name,
171
+ "query": query,
172
+ "documents": document,
173
+ }
174
+ payload.update(extra_body)
175
+ rerank_response = requests.post(
176
+ server.url_for("rerank"),
177
+ json=payload,
178
+ )
179
+
180
+ rerank_response.raise_for_status()
181
+ return RerankResponse.model_validate(rerank_response.json())
182
+
183
+
184
def _get_embeds(server, prompts: list[str]):
185
response = requests.post(
186
server.url_for("pooling"),
@@ -229,6 +258,52 @@ def _test_online_n_v_n(server):
258
assert scores[0] == pytest.approx(expected, abs=TOL)
259
260
261
+def _test_online_instruction(server):
262
+ docs = documents[:2]
263
+
264
+ default_score = _get_score_response(server, query, docs)
265
+ instruction_score = _get_score_response(
266
+ server,
267
+ query,
268
+ docs,
269
+ instruction=INSTRUCTION,
270
+ )
271
+ kwargs_score = _get_score_response(
272
+ server,
273
+ query,
274
+ docs,
275
+ chat_template_kwargs={"instruction": INSTRUCTION},
276
+ )
277
+
278
+ assert instruction_score.usage.prompt_tokens > default_score.usage.prompt_tokens
279
+ assert kwargs_score.usage.prompt_tokens == instruction_score.usage.prompt_tokens
280
+ assert len(instruction_score.data) == len(default_score.data)
281
+ assert [d.score for d in kwargs_score.data] == pytest.approx(
282
+ [d.score for d in instruction_score.data], abs=TOL
283
+ )
284
+
285
+ default_rerank = _get_rerank_response(server, query, docs)
286
+ instruction_rerank = _get_rerank_response(
287
+ server,
288
+ query,
289
+ docs,
290
+ instruction=INSTRUCTION,
291
+ )
292
+ kwargs_rerank = _get_rerank_response(
293
+ server,
294
+ query,
295
+ docs,
296
+ chat_template_kwargs={"instruction": INSTRUCTION},
297
+ )
298
+
299
+ assert instruction_rerank.usage.prompt_tokens > default_rerank.usage.prompt_tokens
300
+ assert kwargs_rerank.usage.prompt_tokens == instruction_rerank.usage.prompt_tokens
301
+ assert len(instruction_rerank.results) == len(default_rerank.results)
302
+ assert [r.relevance_score for r in kwargs_rerank.results] == pytest.approx(
303
+ [r.relevance_score for r in instruction_rerank.results], abs=TOL
304
+ )
305
+
306
+
307
def _test_online_token_embed_illegal_inputs(server):
308
response = requests.post(
309
server.url_for("pooling"),
vllm/entrypoints/pooling/scoring/io_processor.py
+22
-13
@@ -618,12 +618,6 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
618
619
620
class JinaRankingIOProcessorMixin:
621
- @staticmethod
622
- def sanitize_input(text: str, special_tokens: dict[str, str]) -> str:
623
- for token in special_tokens.values():
624
- text = text.replace(token, "")
625
- return text
626
-
621
@staticmethod
622
def format_docs_prompts_func(
623
query: str,
@@ -641,11 +635,13 @@ class JinaRankingIOProcessorMixin:
635
if special_tokens is None:
636
special_tokens = default_special_tokens
637
644
- query = JinaRankingIOProcessorMixin.sanitize_input(query, special_tokens)
645
- docs = [
646
- JinaRankingIOProcessorMixin.sanitize_input(doc, special_tokens)
647
- for doc in docs
648
- ]
638
+ def sanitize_input(text: str) -> str:
639
+ for token in special_tokens.values():
640
+ text = text.replace(token, "")
641
+ return text
642
+
643
+ query = sanitize_input(query)
644
+ docs = [sanitize_input(doc) for doc in docs]
645
646
prefix = (
647
"<|im_start|>system\n"
@@ -668,6 +664,7 @@ class JinaRankingIOProcessorMixin:
664
)
665
666
if instruction:
667
+ instruction = sanitize_input(instruction)
668
prompt += f"<instruct>\n{instruction}\n</instruct>\n"
669
670
doc_prompts = [
@@ -703,12 +700,24 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
700
) -> Sequence[EngineInput]:
701
queries = self.ensure_str(scoring_data.data_1)
702
docs = self.ensure_str(scoring_data.data_2)
703
+ chat_template_kwargs = (
704
+ prompt_extras.get("chat_template_kwargs") if prompt_extras else None
705
+ )
706
+ instruction = (
707
+ chat_template_kwargs.get("instruction") if chat_template_kwargs else None
708
+ )
709
710
if len(queries) == 1:
708
- prompts = [self.format_docs_prompts_func(query=queries[0], docs=docs)]
711
+ prompts = [
712
+ self.format_docs_prompts_func(
713
+ query=queries[0], docs=docs, instruction=instruction
714
+ )
715
+ ]
716
else:
717
prompts = [
711
- self.format_docs_prompts_func(query=q, docs=[d])
718
+ self.format_docs_prompts_func(
719
+ query=q, docs=[d], instruction=instruction
720
+ )
721
for q, d in zip(queries, docs)
722
]
723