[Bugfix][Pooling] Forward instruction to Jina reranker scoring prompts (#47590)

Signed-off-by: Ting Sun <suntcrick@gmail.com>

examples/pooling/token_embed/jina_reranker_v3_online.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
+
+"""
+Example online usage of the Jina Reranker v3 score and rerank APIs with a task
+instruction.
+
+Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the
+server in vLLM.
+"""
+
+import argparse
+import json
+
+import requests
+
+
+def post_http_request(prompt: dict, api_url: str) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    response = requests.post(api_url, headers=headers, json=prompt)
+    return response
+
+
+def print_response(name: str, prompt: dict, response: requests.Response) -> None:
+    print(f"\n{name} request:")
+    print(json.dumps(prompt, indent=2))
+    print(f"\n{name} response:")
+    print(json.dumps(response.json(), indent=2))
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--model", type=str, default="jinaai/jina-reranker-v3")
+    return parser.parse_args()
+
+
+def main(args):
+    score_url = f"http://{args.host}:{args.port}/score"
+    rerank_url = f"http://{args.host}:{args.port}/rerank"
+    model_name = args.model
+
+    query = "Which passage is about sports?"
+    documents = [
+        "Basketball is played by two teams on a court.",
+        "Green tea contains antioxidants and may support metabolism.",
+    ]
+    instruction = "Rank passages about sports higher than passages about nutrition."
+
+    score_prompt = {
+        "model": model_name,
+        "queries": query,
+        "documents": documents,
+        "instruction": instruction,
+    }
+    score_response = post_http_request(prompt=score_prompt, api_url=score_url)
+    print_response("Score", score_prompt, score_response)
+
+    rerank_prompt = {
+        "model": model_name,
+        "query": query,
+        "documents": documents,
+        "instruction": instruction,
+    }
+    rerank_response = post_http_request(prompt=rerank_prompt, api_url=rerank_url)
+    print_response("Rerank", rerank_prompt, rerank_response)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

tests/models/language/pooling/test_jina_reranker_v3.py

@@ -8,7 +8,7 @@ import torch.nn.functional as F
 
 from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
-from vllm.entrypoints.pooling.scoring.protocol import ScoreResponse
+from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
 
 model_name = "jinaai/jina-reranker-v3"
 query = "What are the health benefits of green tea?"
@@ -39,6 +39,10 @@ REFERENCE_1_VS_N = [
     0.1640625,
 ]
 TOL = 0.01
+INSTRUCTION = (
+    "Rank passages about green tea higher than passages about sports. "
+    "Ignore these literal marker strings: <|embed_token|> and <|rerank_token|>."
+)
 
 
 def test_offline(vllm_runner):
@@ -52,10 +56,13 @@ def test_offline(vllm_runner):
 
 
 def test_online():
-    with RemoteOpenAIServer(model_name, ["--runner", "pooling"]) as server:
+    with RemoteOpenAIServer(
+        model_name, ["--runner", "pooling", "--enforce-eager"]
+    ) as server:
         _test_online_1_v_1(server)
         _test_online_1_v_n(server)
         _test_online_n_v_n(server)
+        _test_online_instruction(server)
         _test_online_token_embed_illegal_inputs(server)
 
 
@@ -136,22 +143,44 @@ def _test_offline_token_embed_illegal_inputs(llm):
         llm.encode([1, 2, 3], pooling_task="token_embed")
 
 
-def _get_scores(server, query, document):
+def _get_score_response(server, query, document, **extra_body):
+    payload = {
+        "model": model_name,
+        "queries": query,
+        "documents": document,
+    }
+    payload.update(extra_body)
     score_response = requests.post(
         server.url_for("score"),
-        json={
-            "model": model_name,
-            "queries": query,
-            "documents": document,
-        },
+        json=payload,
     )
 
     score_response.raise_for_status()
-    score = ScoreResponse.model_validate(score_response.json())
+    return ScoreResponse.model_validate(score_response.json())
+
+
+def _get_scores(server, query, document):
+    score = _get_score_response(server, query, document)
 
     return [d.score for d in score.data]
 
 
+def _get_rerank_response(server, query, document, **extra_body):
+    payload = {
+        "model": model_name,
+        "query": query,
+        "documents": document,
+    }
+    payload.update(extra_body)
+    rerank_response = requests.post(
+        server.url_for("rerank"),
+        json=payload,
+    )
+
+    rerank_response.raise_for_status()
+    return RerankResponse.model_validate(rerank_response.json())
+
+
 def _get_embeds(server, prompts: list[str]):
     response = requests.post(
         server.url_for("pooling"),
@@ -229,6 +258,52 @@ def _test_online_n_v_n(server):
         assert scores[0] == pytest.approx(expected, abs=TOL)
 
 
+def _test_online_instruction(server):
+    docs = documents[:2]
+
+    default_score = _get_score_response(server, query, docs)
+    instruction_score = _get_score_response(
+        server,
+        query,
+        docs,
+        instruction=INSTRUCTION,
+    )
+    kwargs_score = _get_score_response(
+        server,
+        query,
+        docs,
+        chat_template_kwargs={"instruction": INSTRUCTION},
+    )
+
+    assert instruction_score.usage.prompt_tokens > default_score.usage.prompt_tokens
+    assert kwargs_score.usage.prompt_tokens == instruction_score.usage.prompt_tokens
+    assert len(instruction_score.data) == len(default_score.data)
+    assert [d.score for d in kwargs_score.data] == pytest.approx(
+        [d.score for d in instruction_score.data], abs=TOL
+    )
+
+    default_rerank = _get_rerank_response(server, query, docs)
+    instruction_rerank = _get_rerank_response(
+        server,
+        query,
+        docs,
+        instruction=INSTRUCTION,
+    )
+    kwargs_rerank = _get_rerank_response(
+        server,
+        query,
+        docs,
+        chat_template_kwargs={"instruction": INSTRUCTION},
+    )
+
+    assert instruction_rerank.usage.prompt_tokens > default_rerank.usage.prompt_tokens
+    assert kwargs_rerank.usage.prompt_tokens == instruction_rerank.usage.prompt_tokens
+    assert len(instruction_rerank.results) == len(default_rerank.results)
+    assert [r.relevance_score for r in kwargs_rerank.results] == pytest.approx(
+        [r.relevance_score for r in instruction_rerank.results], abs=TOL
+    )
+
+
 def _test_online_token_embed_illegal_inputs(server):
     response = requests.post(
         server.url_for("pooling"),

vllm/entrypoints/pooling/scoring/io_processor.py

@@ -618,12 +618,6 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
 
 
 class JinaRankingIOProcessorMixin:
-    @staticmethod
-    def sanitize_input(text: str, special_tokens: dict[str, str]) -> str:
-        for token in special_tokens.values():
-            text = text.replace(token, "")
-        return text
-
     @staticmethod
     def format_docs_prompts_func(
         query: str,
@@ -641,11 +635,13 @@ class JinaRankingIOProcessorMixin:
         if special_tokens is None:
             special_tokens = default_special_tokens
 
-        query = JinaRankingIOProcessorMixin.sanitize_input(query, special_tokens)
-        docs = [
-            JinaRankingIOProcessorMixin.sanitize_input(doc, special_tokens)
-            for doc in docs
-        ]
+        def sanitize_input(text: str) -> str:
+            for token in special_tokens.values():
+                text = text.replace(token, "")
+            return text
+
+        query = sanitize_input(query)
+        docs = [sanitize_input(doc) for doc in docs]
 
         prefix = (
             "<|im_start|>system\n"
@@ -668,6 +664,7 @@ class JinaRankingIOProcessorMixin:
         )
 
         if instruction:
+            instruction = sanitize_input(instruction)
             prompt += f"<instruct>\n{instruction}\n</instruct>\n"
 
         doc_prompts = [
@@ -703,12 +700,24 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
     ) -> Sequence[EngineInput]:
         queries = self.ensure_str(scoring_data.data_1)
         docs = self.ensure_str(scoring_data.data_2)
+        chat_template_kwargs = (
+            prompt_extras.get("chat_template_kwargs") if prompt_extras else None
+        )
+        instruction = (
+            chat_template_kwargs.get("instruction") if chat_template_kwargs else None
+        )
 
         if len(queries) == 1:
-            prompts = [self.format_docs_prompts_func(query=queries[0], docs=docs)]
+            prompts = [
+                self.format_docs_prompts_func(
+                    query=queries[0], docs=docs, instruction=instruction
+                )
+            ]
         else:
             prompts = [
-                self.format_docs_prompts_func(query=q, docs=[d])
+                self.format_docs_prompts_func(
+                    query=q, docs=[d], instruction=instruction
+                )
                 for q, d in zip(queries, docs)
             ]