@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4bed167768

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 15, 02:15 UTC

[Model][VLM] Support JinaVL Reranker (#20260)

Signed-off-by: shineran96 <shinewang96@gmail.com>

shineran96 committed Jul 11, 2025 at 01:43 UTC 4bed167768bd228b9251d77f7b338fb997d3aefe

15 files changed +993 -133

.buildkite/test-pipeline.yaml

+1 -1

@@ -282,7 +282,7 @@ steps:

282 - python3 offline_inference/llm_engine_example.py

283 - python3 offline_inference/audio_language.py --seed 0

284 - python3 offline_inference/vision_language.py --seed 0

285 - - python3 offline_inference/vision_language_embedding.py --seed 0

285 + - python3 offline_inference/vision_language_pooling.py --seed 0

286 - python3 offline_inference/vision_language_multi_image.py --seed 0

287 - VLLM_USE_V1=0 python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors

288 - python3 offline_inference/encoder_decoder.py

docs/models/supported_models.md

@@ -712,6 +712,14 @@ The following table lists those that are tested in vLLM.

712

713 ---

714

715 +#### Scoring

716 +

717 +Specified using `--task score`.

718 +

720 +|-------------------------------------|--------------------|----------|--------------------------|------------------------|-----------------------------|-----------------------|

722 +

723 ## Model Support Policy

724

725 At vLLM, we are committed to facilitating the integration and support of third-party models within our ecosystem. Our approach is designed to balance the need for robustness and the practical limitations of supporting a wide range of models. Here’s how we manage third-party model support:

docs/serving/openai_compatible_server.md

+51 -3

@@ -537,7 +537,7 @@ The following extra parameters are supported:

537

538 ### Score API

539

540 -Our Score API can apply a cross-encoder model or an embedding model to predict scores for sentence pairs. When using an embedding model the score corresponds to the cosine similarity between each embedding pair.

540 +Our Score API can apply a cross-encoder model or an embedding model to predict scores for sentence or multimodal pairs. When using an embedding model the score corresponds to the cosine similarity between each embedding pair.

541 Usually, the score for a sentence pair refers to the similarity between two sentences, on a scale of 0 to 1.

542

543 You can find the documentation for cross encoder models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).

@@ -676,6 +676,55 @@ The total number of pairs is `len(text_2)`.

676 }

677 ```

678

679 +#### Multi-modal inputs

680 +

681 +You can pass multi-modal inputs to scoring models by passing `content` including a list of multi-modal input (image, etc.) in the request. Refer to the examples below for illustration.

682 +

683 +=== "JinaVL-Reranker"

684 +

685 + To serve the model:

686 +

687 + ```bash

688 + vllm serve jinaai/jina-reranker-m0

689 + ```

690 +

691 + Since the request schema is not defined by OpenAI client, we post a request to the server using the lower-level `requests` library:

692 +

693 + ??? Code

694 +

695 + ```python

696 + import requests

697 +

698 + response = requests.post(

699 + "http://localhost:8000/v1/score",

700 + json={

701 + "model": "jinaai/jina-reranker-m0",

702 + "text_1": "slm markdown",

703 + "text_2": {

704 + "content": [

705 + {

706 + "type": "image_url",

707 + "image_url": {

708 + "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"

709 + },

710 + },

711 + {

712 + "type": "image_url",

713 + "image_url": {

714 + "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"

715 + },

716 + },

717 + ]

718 + }

719 + },

720 + )

721 + response.raise_for_status()

722 + response_json = response.json()

723 + print("Scoring output:", response_json["data"][0]["score"])

724 + print("Scoring output:", response_json["data"][1]["score"])

725 + ```

726 +Full example: <gh-file:examples/online_serving/openai_cross_encoder_score_for_multimodal.py>

727 +

728 #### Extra parameters

729

730 The following [pooling parameters][pooling-params] are supported.

@@ -695,8 +744,7 @@ The following extra parameters are supported:

744 ### Re-rank API

745

746 Our Re-rank API can apply an embedding model or a cross-encoder model to predict relevant scores between a single query, and

698 -each of a list of documents. Usually, the score for a sentence pair refers to the similarity between two sentences, on

699 -a scale of 0 to 1.

747 +each of a list of documents. Usually, the score for a sentence pair refers to the similarity between two sentences or multi-modal inputs (image, etc.), on a scale of 0 to 1.

748

749 You can find the documentation for cross encoder models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).

750

examples/offline_inference/vision_language_pooling.py renamed

+89 -7

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """

4 This example shows how to use vLLM for running offline inference with

5 -the correct prompt format on vision language models for multimodal embedding.

5 +the correct prompt format on vision language models for multimodal pooling.

7 For most models, the prompt format should follow corresponding examples

8 on HuggingFace model repository.

@@ -15,6 +15,7 @@ from typing import Literal, NamedTuple, Optional, TypedDict, Union, get_args

15 from PIL.Image import Image

17 from vllm import LLM, EngineArgs

18 +from vllm.entrypoints.score_utils import ScoreMultiModalParam

19 from vllm.multimodal.utils import fetch_image

20 from vllm.utils import FlexibleArgumentParser

@@ -35,14 +36,22 @@ class TextImageQuery(TypedDict):

36 image: Image

38 -QueryModality = Literal["text", "image", "text+image"]

39 -Query = Union[TextQuery, ImageQuery, TextImageQuery]

39 +class TextImagesQuery(TypedDict):

40 + modality: Literal["text+images"]

41 + text: str

42 + image: ScoreMultiModalParam

43 +

44 +

45 +QueryModality = Literal["text", "image", "text+image", "text+images"]

46 +Query = Union[TextQuery, ImageQuery, TextImageQuery, TextImagesQuery]

49 class ModelRequestData(NamedTuple):

50 engine_args: EngineArgs

44 - prompt: str

45 - image: Optional[Image]

51 + prompt: Optional[str] = None

52 + image: Optional[Image] = None

53 + query: Optional[str] = None

54 + documents: Optional[ScoreMultiModalParam] = None

57 def run_e5_v(query: Query) -> ModelRequestData:

@@ -107,6 +116,29 @@ def run_vlm2vec(query: Query) -> ModelRequestData:

116 )

117

118

119 +def run_jinavl_reranker(query: Query) -> ModelRequestData:

120 + if query["modality"] != "text+images":

121 + raise ValueError(f"Unsupported query modality: '{query['modality']}'")

122 +

123 + engine_args = EngineArgs(

124 + model="jinaai/jina-reranker-m0",

125 + task="score",

126 + max_model_len=32768,

127 + trust_remote_code=True,

128 + mm_processor_kwargs={

129 + "min_pixels": 3136,

130 + "max_pixels": 602112,

131 + },

132 + limit_mm_per_prompt={"image": 1},

133 + )

134 +

135 + return ModelRequestData(

136 + engine_args=engine_args,

137 + query=query["text"],

138 + documents=query["image"],

139 + )

140 +

141 +

142 def get_query(modality: QueryModality):

143 if modality == "text":

144 return TextQuery(modality="text", text="A dog sitting in the grass")

@@ -128,6 +160,28 @@ def get_query(modality: QueryModality):

160 ),

161 )

162

163 + if modality == "text+images":

164 + return TextImagesQuery(

165 + modality="text+images",

166 + text="slm markdown",

167 + image={

168 + "content": [

169 + {

170 + "type": "image_url",

171 + "image_url": {

172 + "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"

173 + },

174 + },

175 + {

176 + "type": "image_url",

177 + "image_url": {

178 + "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"

179 + },

180 + },

181 + ]

182 + },

183 + )

184 +

185 msg = f"Modality {modality} is not supported."

186 raise ValueError(msg)

187

@@ -162,16 +216,31 @@ def run_encode(model: str, modality: QueryModality, seed: Optional[int]):

216 print("-" * 50)

217

218

219 +def run_score(model: str, modality: QueryModality, seed: Optional[int]):

220 + query = get_query(modality)

221 + req_data = model_example_map[model](query)

222 +

223 + engine_args = asdict(req_data.engine_args) | {"seed": seed}

224 + llm = LLM(**engine_args)

225 +

226 + outputs = llm.score(req_data.query, req_data.documents)

227 +

228 + print("-" * 30)

229 + print([output.outputs.score for output in outputs])

230 + print("-" * 30)

231 +

232 +

233 model_example_map = {

234 "e5_v": run_e5_v,

235 "vlm2vec": run_vlm2vec,

236 + "jinavl_reranker": run_jinavl_reranker,

237 }

238

239

240 def parse_args():

241 parser = FlexibleArgumentParser(

242 description="Demo on using vLLM for offline inference with "

174 - "vision language models for multimodal embedding"

243 + "vision language models for multimodal pooling tasks."

244 )

245 parser.add_argument(

246 "--model-name",

@@ -181,6 +250,14 @@ def parse_args():

250 choices=model_example_map.keys(),

251 help="The name of the embedding model.",

252 )

253 + parser.add_argument(

254 + "--task",

255 + "-t",

256 + type=str,

257 + default="embedding",

258 + choices=["embedding", "scoring"],

259 + help="The task type.",

260 + )

261 parser.add_argument(

262 "--modality",

263 type=str,

@@ -198,7 +275,12 @@ def parse_args():

275

276

277 def main(args: Namespace):

201 - run_encode(args.model_name, args.modality, args.seed)

278 + if args.task == "embedding":

279 + run_encode(args.model_name, args.modality, args.seed)

280 + elif args.task == "scoring":

281 + run_score(args.model_name, args.modality, args.seed)

282 + else:

283 + raise ValueError(f"Unsupported task: {args.task}")

284

285

286 if __name__ == "__main__":

examples/online_serving/openai_cross_encoder_score_for_multimodal.py new

+60

@@ -0,0 +1,60 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""

4 +Example online usage of Score API.

5 +

6 +Run `vllm serve <model> --task score` to start up the server in vLLM.

7 +"""

8 +

9 +import argparse

10 +import pprint

11 +

12 +import requests

13 +

14 +

15 +def post_http_request(prompt: dict, api_url: str) -> requests.Response:

16 + headers = {"User-Agent": "Test Client"}

17 + response = requests.post(api_url, headers=headers, json=prompt)

18 + return response

19 +

20 +

21 +def parse_args():

22 + parser = argparse.ArgumentParser()

23 + parser.add_argument("--host", type=str, default="localhost")

24 + parser.add_argument("--port", type=int, default=8000)

25 + parser.add_argument("--model", type=str, default="jinaai/jina-reranker-m0")

26 + return parser.parse_args()

27 +

28 +

29 +def main(args):

30 + api_url = f"http://{args.host}:{args.port}/score"

31 + model_name = args.model

32 +

33 + text_1 = "slm markdown"

34 + text_2 = {

35 + "content": [

36 + {

37 + "type": "image_url",

38 + "image_url": {

39 + "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"

40 + },

41 + },

42 + {

43 + "type": "image_url",

44 + "image_url": {

45 + "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"

46 + },

47 + },

48 + ]

49 + }

50 + prompt = {"model": model_name, "text_1": text_1, "text_2": text_2}

51 + score_response = post_http_request(prompt=prompt, api_url=api_url)

52 + print("\nPrompt when text_1 is string and text_2 is a image list:")

53 + pprint.pprint(prompt)

54 + print("\nScore Response:")

55 + pprint.pprint(score_response.json())

56 +

57 +

58 +if __name__ == "__main__":

59 + args = parse_args()

60 + main(args)

tests/models/multimodal/pooling/test_jinavl_reranker.py new

+160

@@ -0,0 +1,160 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import pytest

5 +from transformers import AutoModel

6 +

7 +model_name = "jinaai/jina-reranker-m0"

8 +

9 +mm_processor_kwargs = {

10 + "min_pixels": 3136,

11 + "max_pixels": 602112,

12 +}

13 +

14 +limit_mm_per_prompt = {"image": 2}

15 +

16 +

17 +def vllm_reranker(model_name,

18 + query,

19 + documents,

20 + query_type="text",

21 + doc_type="text"):

22 + from vllm import LLM

23 +

24 + model = LLM(

25 + model=model_name,

26 + task="score",

27 + max_model_len=32768,

28 + mm_processor_kwargs=mm_processor_kwargs,

29 + limit_mm_per_prompt=limit_mm_per_prompt,

30 + )

31 +

32 + def create_image_param(url: str):

33 + return {"type": "image_url", "image_url": {"url": f"{url}"}}

34 +

35 + if query_type == "image":

36 + query = {"content": [create_image_param(url) for url in query]}

37 +

38 + if doc_type == "image":

39 + documents = {"content": [create_image_param(url) for url in documents]}

40 +

41 + outputs = model.score(query, documents)

42 +

43 + return [output.outputs.score for output in outputs]

44 +

45 +

46 +def hf_reranker(model_name,

47 + query,

48 + documents,

49 + query_type="text",

50 + doc_type="text"):

51 +

52 + checkpoint_to_hf_mapper = {

53 + "visual.": "model.visual.",

54 + "model.": "model.language_model.",

55 + }

56 +

57 + model = AutoModel.from_pretrained(

58 + model_name,

59 + torch_dtype="auto",

60 + trust_remote_code=True,

61 + key_mapping=checkpoint_to_hf_mapper).to("cuda").eval()

62 +

63 + data_pairs = [[query[0], d] for d in documents]

64 +

65 + scores = model.compute_score(data_pairs,

66 + max_length=2048,

67 + query_type=query_type,

68 + doc_type=doc_type)

69 + return scores

70 +

71 +

72 +# Visual Documents Reranking

73 +@pytest.mark.parametrize("model_name", [model_name])

74 +def test_model_text_image(model_name):

75 +

76 + query = ["slm markdown"]

77 + documents = [

78 + "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png",

79 + "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png",

80 + ]

81 +

82 + hf_outputs = hf_reranker(model_name, query, documents, "text", "image")

83 + vllm_outputs = vllm_reranker(model_name, query, documents, "text", "image")

84 +

85 + assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)

86 + assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)

87 +

88 +

89 +# Textual Documents Reranking

90 +@pytest.mark.parametrize("model_name", [model_name])

91 +def test_model_text_text(model_name):

92 +

93 + query = ["slm markdown"]

94 + documents = [

95 + """We present ReaderLM-v2, a compact 1.5 billion parameter language model designed for efficient

96 + web content extraction. Our model processes documents up to 512K tokens, transforming messy HTML

97 + into clean Markdown or JSON formats with high accuracy -- making it an ideal tool for grounding

98 + large language models. The models effectiveness results from two key innovations: (1) a three-stage

99 + data synthesis pipeline that generates high quality, diverse training data by iteratively drafting,

100 + refining, and critiquing web content extraction; and (2) a unified training framework combining

101 + continuous pre-training with multi-objective optimization. Intensive evaluation demonstrates that

102 + ReaderLM-v2 outperforms GPT-4o-2024-08-06 and other larger models by 15-20% on carefully curated

103 + benchmarks, particularly excelling at documents exceeding 100K tokens, while maintaining significantly

104 + lower computational requirements.""", # noqa: E501

105 + "数据提取么？为什么不用正则啊，你用正则不就全解决了么？",

106 + ]

107 +

108 + hf_outputs = hf_reranker(model_name, query, documents, "text", "text")

109 + vllm_outputs = vllm_reranker(model_name, query, documents, "text", "text")

110 +

111 + assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)

112 + assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)

113 +

114 +

115 +# Image Querying for Textual Documents

116 +@pytest.mark.parametrize("model_name", [model_name])

117 +def test_model_image_text(model_name):

118 +

119 + query = [

120 + "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"

121 + ]

122 + documents = [

123 + """We present ReaderLM-v2, a compact 1.5 billion parameter language model designed for efficient

124 + web content extraction. Our model processes documents up to 512K tokens, transforming messy HTML

125 + into clean Markdown or JSON formats with high accuracy -- making it an ideal tool for grounding

126 + large language models. The models effectiveness results from two key innovations: (1) a three-stage

127 + data synthesis pipeline that generates high quality, diverse training data by iteratively drafting,

128 + refining, and critiquing web content extraction; and (2) a unified training framework combining

129 + continuous pre-training with multi-objective optimization. Intensive evaluation demonstrates that

130 + ReaderLM-v2 outperforms GPT-4o-2024-08-06 and other larger models by 15-20% on carefully curated

131 + benchmarks, particularly excelling at documents exceeding 100K tokens, while maintaining significantly

132 + lower computational requirements.""", # noqa: E501

133 + "数据提取么？为什么不用正则啊，你用正则不就全解决了么？",

134 + ]

135 +

136 + hf_outputs = hf_reranker(model_name, query, documents, "image", "text")

137 + vllm_outputs = vllm_reranker(model_name, query, documents, "image", "text")

138 +

139 + assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)

140 + assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)

141 +

142 +

143 +# Image Querying for Image Documents

144 +@pytest.mark.parametrize("model_name", [model_name])

145 +def test_model_image_image(model_name):

146 +

147 + query = [

148 + "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"

149 + ]

150 + documents = [

151 + "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png",

152 + "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png",

153 + ]

154 +

155 + hf_outputs = hf_reranker(model_name, query, documents, "image", "image")

156 + vllm_outputs = vllm_reranker(model_name, query, documents, "image",

157 + "image")

158 +

159 + assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)

160 + assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)

tests/models/registry.py

@@ -432,6 +432,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {

432 trust_remote_code=True), # noqa: E501

433 "MllamaForConditionalGeneration": _HfExamplesInfo("meta-llama/Llama-3.2-11B-Vision-Instruct"), # noqa: E501

434 "WhisperForConditionalGeneration": _HfExamplesInfo("openai/whisper-large-v3"), # noqa: E501

435 +

436 + # [Cross-encoder]

437 + "JinaVLForRanking": _HfExamplesInfo("jinaai/jina-reranker-m0"), # noqa: E501

438 }

439

440 _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

vllm/entrypoints/llm.py

+122 -61

@@ -28,8 +28,11 @@ from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,

28 apply_mistral_chat_template,

29 parse_chat_messages,

30 resolve_chat_template_content_format)

31 -from vllm.entrypoints.score_utils import (_cosine_similarity,

32 - _validate_score_input_lens)

31 +from vllm.entrypoints.score_utils import (ScoreContentPartParam,

32 + ScoreMultiModalParam,

33 + _cosine_similarity,

34 + _validate_score_input_lens,

35 + get_score_prompt)

36 from vllm.entrypoints.utils import _validate_truncation_size

37 from vllm.inputs import PromptType, SingletonPrompt, TextPrompt, TokensPrompt

38 from vllm.inputs.parse import parse_and_batch_prompt

@@ -1187,8 +1190,8 @@ class LLM:

1190 def _cross_encoding_score(

1191 self,

1192 tokenizer: AnyTokenizer,

1190 - text_1: list[str],

1191 - text_2: list[str],

1193 + data_1: Union[list[str], list[ScoreContentPartParam]],

1194 + data_2: Union[list[str], list[ScoreContentPartParam]],

1195 truncate_prompt_tokens: Optional[int] = None,

1196 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1197 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

@@ -1199,10 +1202,8 @@ class LLM:

1202 raise ValueError(

1203 "Score API is only enabled for `--task embed or score`")

1204

1202 - if len(text_1) == 1:

1203 - text_1 = text_1 * len(text_2)

1204 -

1205 - input_pairs = [(t1, t2) for t1, t2 in zip(text_1, text_2)]

1205 + if len(data_1) == 1:

1206 + data_1 = data_1 * len(data_2)

1207

1208 pooling_params = PoolingParams(use_cross_encoder=True)

1209 tokenization_kwargs: dict[str, Any] = {}

@@ -1211,19 +1212,41 @@ class LLM:

1212

1213 parsed_prompts = []

1214

1214 - for q, t in input_pairs:

1215 - if self.llm_engine.model_config.use_pad_token:

1216 - # cross_encoder models defaults to using pad_token.

1217 - prompt_inputs = tokenizer(text=q,

1218 - text_pair=t,

1219 - **tokenization_kwargs)

1220 - else:

1221 - # `llm as reranker` models defaults to not using pad_token.

1222 - prompt_inputs = tokenizer(text=q + t, **tokenization_kwargs)

1223 - engine_prompt = TokensPrompt(

1224 - prompt_token_ids=prompt_inputs["input_ids"],

1225 - token_type_ids=prompt_inputs.get("token_type_ids"))

1226 - parsed_prompts.append(engine_prompt)

1215 + input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]

1216 +

1217 + if self.llm_engine.model_config.is_multimodal_model:

1218 +

1219 + model_config = self.llm_engine.model_config

1220 +

1221 + for q, d in input_pairs:

1222 + _, engine_prompt = get_score_prompt(

1223 + model_config=model_config,

1224 + data_1=q,

1225 + data_2=d,

1226 + tokenizer=tokenizer,

1227 + tokenization_kwargs=tokenization_kwargs,

1228 + )

1229 +

1230 + parsed_prompts.append(engine_prompt)

1231 +

1232 + else:

1233 +

1234 + for q, t in input_pairs:

1235 + if self.llm_engine.model_config.use_pad_token:

1236 + # cross_encoder models defaults to using pad_token.

1237 + prompt_inputs = tokenizer(

1238 + text=q, # type: ignore[arg-type]

1239 + text_pair=t, # type: ignore[arg-type]

1240 + **tokenization_kwargs)

1241 + else:

1242 + # `llm as reranker` models defaults to not using pad_token.

1243 + prompt_inputs = tokenizer(

1244 + text=q + t, # type: ignore[operator]

1245 + **tokenization_kwargs)

1246 + engine_prompt = TokensPrompt(

1247 + prompt_token_ids=prompt_inputs["input_ids"],

1248 + token_type_ids=prompt_inputs.get("token_type_ids"))

1249 + parsed_prompts.append(engine_prompt)

1250

1251 self._validate_and_add_requests(

1252 prompts=parsed_prompts,

@@ -1241,8 +1264,10 @@ class LLM:

1264

1265 def score(

1266 self,

1244 - text_1: Union[SingletonPrompt, Sequence[SingletonPrompt]],

1245 - text_2: Union[SingletonPrompt, Sequence[SingletonPrompt]],

1267 + data_1: Union[SingletonPrompt, Sequence[SingletonPrompt],

1268 + ScoreMultiModalParam],

1269 + data_2: Union[SingletonPrompt, Sequence[SingletonPrompt],

1270 + ScoreMultiModalParam],

1271 /,

1272 *,

1273 truncate_prompt_tokens: Optional[int] = None,

@@ -1250,22 +1275,30 @@ class LLM:

1275 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1276 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1277 ) -> list[ScoringRequestOutput]:

1253 - """Generate similarity scores for all pairs `<text,text_pair>`.

1278 + """Generate similarity scores for all pairs `<text,text_pair>` or

1279 + `<multi-modal data, multi-modal data pair>`.

1280

1281 The inputs can be `1 -> 1`, `1 -> N` or `N -> N`.

1256 - In the `1 - N` case the `text_1` sentence will be replicated `N`

1257 - times to pair with the `text_2` sentences.

1282 + In the `1 - N` case the `data_1` input will be replicated `N`

1283 + times to pair with the `data_2` inputs.

1284 The input pairs are used to build a list of prompts for the

1285 cross encoder model. This class automatically batches the prompts,

1286 considering the memory constraint. For the best performance, put all

1261 - of your texts into a single list and pass it to this method.

1287 + of your inputs into a single list and pass it to this method.

1288 +

1289 + Supports both text and multi-modal data (images, etc.) when used with

1290 + appropriate multi-modal models. For multi-modal inputs, ensure the

1291 + prompt structure matches the model's expected input format.

1292

1293 Args:

1264 - text_1: can be a single prompt or a list of prompts, in which

1265 - case it has to have the same length as the `text_2` list

1266 - text_2: The texts to pair with the query to form the input

1267 - to the LLM. See [PromptType][vllm.inputs.PromptType] for

1268 - more details about the format of each prompts.

1294 + data_1: Can be a single prompt, a list of prompts or

1295 + `ScoreMultiModalParam`, which can contain either text or

1296 + multi-modal data. When a list, it must have the same length as

1297 + the `data_2` list.

1298 + data_2: The data to pair with the query to form the input to

1299 + the LLM. Can be text or multi-modal data. See [PromptType]

1300 + [vllm.inputs.PromptType] for more details about the format of

1301 + each prompt.

1302 use_tqdm: If `True`, shows a tqdm progress bar.

1303 If a callable (e.g., `functools.partial(tqdm, leave=False)`),

1304 it is used to create the progress bar.

@@ -1306,42 +1339,70 @@ class LLM:

1339 # lists of tokens to the `text` and `text_pair` kwargs

1340 tokenizer = self.get_tokenizer()

1341

1309 - def ensure_str(prompt: SingletonPrompt):

1310 - if isinstance(prompt, dict):

1311 - if "multi_modal_data" in prompt:

1312 - raise ValueError("Multi-modal prompt is not "

1313 - "supported for scoring")

1314 - elif "prompt_token_ids" in prompt:

1315 - prompt = tokenizer.decode(

1316 - cast(TokensPrompt, prompt)["prompt_token_ids"])

1317 - elif "prompt" in prompt:

1318 - prompt = cast(TextPrompt, prompt)["prompt"]

1319 - assert type(prompt) is str

1320 - return prompt

1321 -

1322 - if isinstance(text_1, (str, dict)):

1323 - # Convert a single prompt to a list.

1324 - text_1 = [text_1]

1325 - input_text_1: list[str] = [ensure_str(t) for t in text_1]

1342 + if not self.llm_engine.model_config.is_multimodal_model:

1343

1327 - if isinstance(text_2, (str, dict)):

1328 - # Convert a single prompt to a list.

1329 - text_2 = [text_2]

1330 - input_text_2: list[str] = [ensure_str(t) for t in text_2]

1344 + def check_data_type(data: Union[SingletonPrompt,

1345 + Sequence[SingletonPrompt],

1346 + ScoreMultiModalParam]):

1347 + if isinstance(data, dict) and "content" in data:

1348 + raise ValueError(

1349 + f"ScoreMultiModalParam is not supported for {self.llm_engine.model_config.architecture}", # noqa: E501

1350 + )

1351 +

1352 + check_data_type(data_1)

1353 + check_data_type(data_2)

1354 +

1355 + def ensure_str(prompt: SingletonPrompt):

1356 + if isinstance(prompt, dict):

1357 + if "multi_modal_data" in prompt:

1358 + raise ValueError("Multi-modal prompt is not "

1359 + "supported for scoring")

1360 + elif "prompt_token_ids" in prompt:

1361 + prompt = tokenizer.decode(

1362 + cast(TokensPrompt, prompt)["prompt_token_ids"])

1363 + elif "prompt" in prompt:

1364 + prompt = cast(TextPrompt, prompt)["prompt"]

1365 + assert type(prompt) is str

1366 + return prompt

1367 +

1368 + if isinstance(data_1, (str, dict)):

1369 + # Convert a single prompt to a list.

1370 + data_1 = [data_1] # type: ignore[list-item]

1371 +

1372 + data_1 = [ensure_str(t) for t in data_1]

1373 +

1374 + if isinstance(data_2, (str, dict)):

1375 + # Convert a single prompt to a list.

1376 + data_2 = [data_2] # type: ignore[list-item]

1377 +

1378 + data_2 = [ensure_str(t) for t in data_2]

1379

1332 - _validate_score_input_lens(input_text_1, input_text_2)

1380 + if isinstance(data_1, dict) and "content" in data_1:

1381 + data_1 = data_1.get("content") # type: ignore[assignment]

1382 + elif isinstance(data_1, str):

1383 + data_1 = [data_1]

1384 +

1385 + if isinstance(data_2, dict) and "content" in data_2:

1386 + data_2 = data_2.get("content") # type: ignore[assignment]

1387 + elif isinstance(data_2, str):

1388 + data_2 = [data_2]

1389 +

1390 + _validate_score_input_lens(data_1, data_2) # type: ignore[arg-type]

1391

1392 if self.llm_engine.model_config.is_cross_encoder:

1335 - return self._cross_encoding_score(tokenizer, input_text_1,

1336 - input_text_2,

1337 - truncate_prompt_tokens, use_tqdm,

1338 - lora_request,

1339 - prompt_adapter_request)

1393 + return self._cross_encoding_score(

1394 + tokenizer,

1395 + data_1, # type: ignore[arg-type]

1396 + data_2, # type: ignore[arg-type]

1397 + truncate_prompt_tokens,

1398 + use_tqdm,

1399 + lora_request,

1400 + prompt_adapter_request)

1401 else:

1402 return self._embedding_score(

1403 tokenizer,

1343 - input_text_1, # type: ignore[arg-type]

1344 - input_text_2, # type: ignore[arg-type]

1404 + data_1, # type: ignore[arg-type]

1405 + data_2, # type: ignore[arg-type]

1406 truncate_prompt_tokens,

1407 use_tqdm,

1408 lora_request,

vllm/entrypoints/openai/protocol.py

+19 -5

@@ -24,6 +24,7 @@ from typing_extensions import TypeAlias

24 from vllm import envs

25 from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,

26 random_tool_call_id)

27 +from vllm.entrypoints.score_utils import ScoreMultiModalParam

28 from vllm.logger import init_logger

29 from vllm.pooling_params import PoolingParams

30 from vllm.sampling_params import (BeamSearchParams, GuidedDecodingParams,

@@ -1282,8 +1283,8 @@ PoolingRequest = Union[PoolingCompletionRequest, PoolingChatRequest]

1283

1284 class ScoreRequest(OpenAIBaseModel):

1285 model: Optional[str] = None

1285 - text_1: Union[list[str], str]

1286 - text_2: Union[list[str], str]

1286 + text_1: Union[list[str], str, ScoreMultiModalParam]

1287 + text_2: Union[list[str], str, ScoreMultiModalParam]

1288 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

1289

1290 # --8<-- [start:score-pooling-params]

@@ -1291,6 +1292,12 @@ class ScoreRequest(OpenAIBaseModel):

1292 # --8<-- [end:score-pooling-params]

1293

1294 # --8<-- [start:score-extra-params]

1295 +

1296 + mm_processor_kwargs: Optional[dict[str, Any]] = Field(

1297 + default=None,

1298 + description=("Additional kwargs to pass to the HF processor."),

1299 + )

1300 +

1301 priority: int = Field(

1302 default=0,

1303 description=(

@@ -1308,8 +1315,8 @@ class ScoreRequest(OpenAIBaseModel):

1315

1316 class RerankRequest(OpenAIBaseModel):

1317 model: Optional[str] = None

1311 - query: str

1312 - documents: list[str]

1318 + query: Union[str, ScoreMultiModalParam]

1319 + documents: Union[list[str], ScoreMultiModalParam]

1320 top_n: int = Field(default_factory=lambda: 0)

1321 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None

1322

@@ -1318,6 +1325,12 @@ class RerankRequest(OpenAIBaseModel):

1325 # --8<-- [end:rerank-pooling-params]

1326

1327 # --8<-- [start:rerank-extra-params]

1328 +

1329 + mm_processor_kwargs: Optional[dict[str, Any]] = Field(

1330 + default=None,

1331 + description=("Additional kwargs to pass to the HF processor."),

1332 + )

1333 +

1334 priority: int = Field(

1335 default=0,

1336 description=(

@@ -1334,7 +1347,8 @@ class RerankRequest(OpenAIBaseModel):

1347

1348

1349 class RerankDocument(BaseModel):

1337 - text: str

1350 + text: Optional[str] = None

1351 + multi_modal: Optional[ScoreMultiModalParam] = None

1352

1353

1354 class RerankResult(BaseModel):

vllm/entrypoints/openai/serving_score.py

+110 -48

@@ -17,8 +17,11 @@ from vllm.entrypoints.openai.protocol import (ErrorResponse, RerankDocument,

17 ScoreResponseData, UsageInfo)

18 from vllm.entrypoints.openai.serving_engine import OpenAIServing

19 from vllm.entrypoints.openai.serving_models import OpenAIServingModels

20 -from vllm.entrypoints.score_utils import (_cosine_similarity,

21 - _validate_score_input_lens)

20 +from vllm.entrypoints.score_utils import (ScoreContentPartParam,

21 + ScoreMultiModalParam,

22 + _cosine_similarity,

23 + _validate_score_input_lens,

24 + get_score_prompt)

25 from vllm.entrypoints.utils import _validate_truncation_size

26 from vllm.inputs.data import TokensPrompt

27 from vllm.logger import init_logger

@@ -137,11 +140,34 @@ class ServingScores(OpenAIServing):

140

141 return final_res_batch

142

143 + def _preprocess_score(

144 + self,

145 + request: Union[RerankRequest, ScoreRequest],

146 + tokenizer: AnyTokenizer,

147 + tokenization_kwargs: dict[str, Any],

148 + data_1: Union[str, ScoreContentPartParam],

149 + data_2: Union[str, ScoreContentPartParam],

150 + ) -> tuple[str, TokensPrompt]:

151 +

152 + model_config = self.model_config

153 +

154 + full_prompt, engine_prompt = get_score_prompt(

155 + model_config=model_config,

156 + data_1=data_1,

157 + data_2=data_2,

158 + tokenizer=tokenizer,

159 + tokenization_kwargs=tokenization_kwargs,

160 + )

161 + if request.mm_processor_kwargs is not None:

162 + engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs

163 +

164 + return full_prompt, engine_prompt

165 +

166 async def _cross_encoding_score(

167 self,

168 tokenizer: AnyTokenizer,

143 - texts_1: list[str],

144 - texts_2: list[str],

169 + data_1: Union[list[str], list[ScoreContentPartParam]],

170 + data_2: Union[list[str], list[ScoreContentPartParam]],

171 request: Union[RerankRequest, ScoreRequest],

172 request_id=str,

173 tokenization_kwargs: Optional[dict[str, Any]] = None,

@@ -154,46 +180,67 @@ class ServingScores(OpenAIServing):

180 request_prompts: list[str] = []

181 engine_prompts: list[TokensPrompt] = []

182

157 - if len(texts_1) == 1:

158 - texts_1 = texts_1 * len(texts_2)

159 -

160 - input_pairs = [(t1, t2) for t1, t2 in zip(texts_1, texts_2)]

183 + if len(data_1) == 1:

184 + data_1 = data_1 * len(data_2)

185

186 if isinstance(tokenizer, MistralTokenizer):

187 raise ValueError(

188 "MistralTokenizer not supported for cross-encoding")

189

166 - tokenize_async = make_async(tokenizer.__call__,

167 - executor=self._tokenizer_executor)

168 -

190 tokenization_kwargs = tokenization_kwargs or {}

170 - use_pad_token = self.model_config.use_pad_token

191

172 - if use_pad_token:

173 - # cross_encoder models defaults to using pad_token.

174 - tokenized_prompts = await asyncio.gather(

175 - *(tokenize_async(text=t1, text_pair=t2, **tokenization_kwargs)

176 - for t1, t2 in input_pairs))

177 - else:

178 - # `llm as reranker` models defaults to not using pad_token.

179 - tokenized_prompts = await asyncio.gather(

180 - *(tokenize_async(text=t1 + t2, **tokenization_kwargs)

181 - for t1, t2 in input_pairs))

192 + input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]

193

183 - for prompt_inputs, (t1, t2) in zip(tokenized_prompts, input_pairs):

184 - sep_token = tokenizer.sep_token if (tokenizer.sep_token

185 - and use_pad_token) else ''

186 - request_prompt = f"{t1}{sep_token}{t2}"

194 + if self.model_config.is_multimodal_model:

195

188 - input_ids = prompt_inputs["input_ids"]

189 - text_token_prompt = \

190 - self._validate_input(request, input_ids, request_prompt)

191 - engine_prompt = TokensPrompt(

192 - prompt_token_ids=text_token_prompt["prompt_token_ids"],

193 - token_type_ids=prompt_inputs.get("token_type_ids"))

196 + preprocess_async = make_async(self._preprocess_score,

197 + executor=self._tokenizer_executor)

198 +

199 + preprocessed_prompts = await asyncio.gather(

200 + *(preprocess_async(request=request,

201 + tokenizer=tokenizer,

202 + tokenization_kwargs=tokenization_kwargs,

203 + data_1=t1,

204 + data_2=t2) for t1, t2 in input_pairs))

205 +

206 + for full_prompt, engine_prompt in preprocessed_prompts:

207 + request_prompts.append(full_prompt)

208 + engine_prompts.append(engine_prompt)

209

195 - request_prompts.append(request_prompt)

196 - engine_prompts.append(engine_prompt)

210 + else:

211 + tokenize_async = make_async(tokenizer.__call__,

212 + executor=self._tokenizer_executor)

213 + use_pad_token = self.model_config.use_pad_token

214 +

215 + if use_pad_token:

216 + # cross_encoder models defaults to using pad_token.

217 + tokenized_prompts = await asyncio.gather(*(

218 + tokenize_async(

219 + text=t1, # type: ignore[arg-type]

220 + text_pair=t2, # type: ignore[arg-type]

221 + **tokenization_kwargs) for t1, t2 in input_pairs))

222 + else:

223 + # `llm as reranker` models defaults to not using pad_token.

224 + tokenized_prompts = await asyncio.gather(*(

225 + tokenize_async(

226 + text=t1 + # type: ignore[operator]

227 + t2,

228 + **tokenization_kwargs) for t1, t2 in input_pairs))

229 +

230 + for prompt_inputs, (t1, t2) in zip(tokenized_prompts, input_pairs):

231 + sep_token = tokenizer.sep_token if (tokenizer.sep_token

232 + and use_pad_token) else ''

233 + request_prompt = f"{t1}{sep_token}{t2}"

234 +

235 + input_ids = prompt_inputs["input_ids"]

236 + text_token_prompt = \

237 + self._validate_input(request, input_ids, request_prompt)

238 + engine_prompt = TokensPrompt(

239 + prompt_token_ids=text_token_prompt["prompt_token_ids"],

240 + token_type_ids=prompt_inputs.get("token_type_ids"))

241 +

242 + request_prompts.append(request_prompt)

243 + engine_prompts.append(engine_prompt)

244

245 # Schedule the request and get the result generator.

246 generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []

@@ -233,8 +280,8 @@ class ServingScores(OpenAIServing):

280

281 async def _run_scoring(

282 self,

236 - texts_1: Union[str, list[str]],

237 - texts_2: Union[str, list[str]],

283 + data_1: Union[list[str], str, ScoreMultiModalParam],

284 + data_2: Union[list[str], str, ScoreMultiModalParam],

285 request: Union[ScoreRequest, RerankRequest],

286 request_id: str,

287 raw_request: Optional[Request] = None,

@@ -259,18 +306,29 @@ class ServingScores(OpenAIServing):

306 trace_headers = (None if raw_request is None else await

307 self._get_trace_headers(raw_request.headers))

308

262 - if isinstance(texts_1, str):

263 - texts_1 = [texts_1]

264 - if isinstance(texts_2, str):

265 - texts_2 = [texts_2]

309 + if not self.model_config.is_multimodal_model and (isinstance(

310 + data_1, dict) or isinstance(data_2, dict)):

311 + raise ValueError(

312 + f"MultiModalParam is not supported for {self.model_config.architecture}" # noqa: E501

313 + )

314 +

315 + if isinstance(data_1, str):

316 + data_1 = [data_1]

317 + elif isinstance(data_1, dict):

318 + data_1 = data_1.get("content") # type: ignore[assignment]

319 +

320 + if isinstance(data_2, str):

321 + data_2 = [data_2]

322 + elif isinstance(data_2, dict):

323 + data_2 = data_2.get("content") # type: ignore[assignment]

324

267 - _validate_score_input_lens(texts_1, texts_2)

325 + _validate_score_input_lens(data_1, data_2) # type: ignore[arg-type]

326

327 if self.model_config.is_cross_encoder:

328 return await self._cross_encoding_score(

329 tokenizer=tokenizer,

272 - texts_1=texts_1,

273 - texts_2=texts_2,

330 + data_1=data_1, # type: ignore[arg-type]

331 + data_2=data_2, # type: ignore[arg-type]

332 request=request,

333 request_id=request_id,

334 tokenization_kwargs=tokenization_kwargs,

@@ -281,8 +339,8 @@ class ServingScores(OpenAIServing):

339 else:

340 return await self._embedding_score(

341 tokenizer=tokenizer,

284 - texts_1=texts_1,

285 - texts_2=texts_2,

342 + texts_1=data_1, # type: ignore[arg-type]

343 + texts_2=data_2, # type: ignore[arg-type]

344 request=request,

345 request_id=request_id,

346 tokenization_kwargs=tokenization_kwargs,

@@ -349,7 +407,9 @@ class ServingScores(OpenAIServing):

407

408 request_id = f"rerank-{self._base_request_id(raw_request)}"

409 documents = request.documents

352 - top_n = request.top_n if request.top_n > 0 else len(documents)

410 + top_n = request.top_n if request.top_n > 0 else (

411 + len(documents)

412 + if isinstance(documents, list) else len(documents["content"]))

413

414 try:

415 final_res_batch = await self._run_scoring(

@@ -410,7 +470,7 @@ class ServingScores(OpenAIServing):

470

471 def request_output_to_rerank_response(

472 self, final_res_batch: list[PoolingRequestOutput], request_id: str,

413 - model_name: str, documents: list[str],

473 + model_name: str, documents: Union[list[str], ScoreMultiModalParam],

474 top_n: int) -> RerankResponse:

475 """

476 Convert the output of do_rank to a RerankResponse

@@ -422,7 +482,9 @@ class ServingScores(OpenAIServing):

482

483 result = RerankResult(

484 index=idx,

425 - document=RerankDocument(text=documents[idx]),

485 + document=RerankDocument(text=documents[idx]) if isinstance(

486 + documents, list) else RerankDocument(

487 + multi_modal=documents["content"][idx]),

488 relevance_score=classify_res.outputs.score,

489 )

490 results.append(result)

vllm/entrypoints/score_utils.py

+152 -8

@@ -1,13 +1,41 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Union

3 +from typing import Any, Optional, Union, cast

5 from torch.nn import CosineSimilarity

6 +from typing_extensions import Required, TypeAlias, TypedDict

8 +from vllm.config import ModelConfig

9 +from vllm.entrypoints.chat_utils import (

10 + BaseMultiModalItemTracker, ChatCompletionContentPartImageEmbedsParam,

11 + ChatCompletionContentPartImageParam, ChatCompletionContentPartTextParam,

12 + MultiModalItemTracker, _ContentPart, _parse_chat_message_content_part)

13 +from vllm.inputs import TokensPrompt

14 +from vllm.model_executor.model_loader import get_model_cls

15 +from vllm.model_executor.models.interfaces import supports_score_template

16 +from vllm.multimodal.inputs import MultiModalDataDict

17 from vllm.outputs import PoolingRequestOutput

8 -from vllm.transformers_utils.tokenizer import (PreTrainedTokenizer,

18 +from vllm.transformers_utils.tokenizer import (AnyTokenizer,

19 + PreTrainedTokenizer,

20 PreTrainedTokenizerFast)

22 +ScoreContentPartParam: TypeAlias = Union[

23 + ChatCompletionContentPartImageParam,

24 + ChatCompletionContentPartImageEmbedsParam]

25 +

26 +

27 +class ScoreMultiModalParam(TypedDict, total=False):

28 + """

29 + A specialized parameter type for scoring multimodal content

30 +

31 + The reasons why don't reuse `CustomChatCompletionMessageParam` directly:

32 + 1. Score tasks don't need the 'role' field (user/assistant/system) that's required in chat completions

33 + 2. Including chat-specific fields would confuse users about their purpose in scoring

34 + 3. This is a more focused interface that only exposes what's needed for scoring

35 + """ # noqa: E501

36 + content: Required[list[ScoreContentPartParam]]

37 + """The multimodal contents"""

38 +

40 def _cosine_similarity(

41 tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

@@ -39,12 +67,128 @@ def _cosine_similarity(

69 def _validate_score_input_lens(

42 - texts_1: Union[list[str], list[dict]],

43 - texts_2: Union[list[str], list[dict]],

70 + data_1: Union[list[str], list[ScoreContentPartParam]],

71 + data_2: Union[list[str], list[ScoreContentPartParam]],

72 ):

45 - if len(texts_1) > 1 and len(texts_1) != len(texts_2):

73 + len_1 = len(data_1)

74 + len_2 = len(data_2)

75 +

76 + if len_1 > 1 and len_1 != len_2:

77 raise ValueError("Input lengths must be either 1:1, 1:N or N:N")

47 - if len(texts_1) == 0:

78 + if len_1 == 0:

79 raise ValueError("At least one text element must be given")

49 - if len(texts_2) == 0:

50 - raise ValueError("At least one text_pair element must be given")

\ No newline at end of file

80 + if len_2 == 0:

81 + raise ValueError("At least one text_pair element must be given")

82 +

83 +

84 +def parse_score_data(

85 + data_1: Union[str, ScoreContentPartParam],

86 + data_2: Union[str, ScoreContentPartParam],

87 + model_config: ModelConfig,

88 + tokenizer: AnyTokenizer,

89 +) -> tuple[str, str, Optional[MultiModalDataDict]]:

90 + mm_tracker = MultiModalItemTracker(model_config, tokenizer)

91 +

92 + content_1 = _parse_score_content(data_1, mm_tracker)

93 +

94 + content_2 = _parse_score_content(data_2, mm_tracker)

95 +

96 + def ensure_str(content: Optional[_ContentPart]) -> str:

97 + if content is not None and isinstance(content, str):

98 + return cast(str, content)

99 + else:

100 + raise ValueError(

101 + f"Only string content is supported, but got {content}.")

102 +

103 + prompt_1 = ensure_str(content_1)

104 + prompt_2 = ensure_str(content_2)

105 +

106 + return prompt_1, prompt_2, mm_tracker.all_mm_data()

107 +

108 +

109 +def _parse_score_content(

110 + data: Union[str, ScoreContentPartParam],

111 + mm_tracker: BaseMultiModalItemTracker,

112 +) -> Optional[_ContentPart]:

113 +

114 + if isinstance(data, str):

115 + data = ChatCompletionContentPartTextParam(type="text", text=data)

116 +

117 + mm_parser = mm_tracker.create_parser()

118 +

119 + parse_res = _parse_chat_message_content_part(

120 + data,

121 + mm_parser,

122 + wrap_dicts=False,

123 + interleave_strings=False,

124 + )

125 +

126 + if parse_res:

127 + return parse_res

128 +

129 + mm_placeholder_storage = mm_parser.mm_placeholder_storage()

130 +

131 + if len(mm_placeholder_storage) != 1 or len(

132 + next(iter(mm_placeholder_storage.values()))) != 1:

133 + raise ValueError("Only one multi-modal item is supported")

134 +

135 + return next(iter(mm_placeholder_storage.values()))[0]

136 +

137 +

138 +def apply_score_template(

139 + model_config: ModelConfig,

140 + prompt_1: str,

141 + prompt_2: str,

142 +) -> str:

143 +

144 + model = get_model_cls(model_config)

145 + if supports_score_template(model):

146 + full_prompt = model.get_score_template(prompt_1, prompt_2)

147 + if full_prompt is None:

148 + raise ValueError("Get empty score template from model")

149 + return full_prompt

150 +

151 + raise ValueError(

152 + f"Unsupported model architecture: {model_config.architecture}")

153 +

154 +

155 +def post_process_tokens(

156 + model_config: ModelConfig,

157 + prompt: TokensPrompt,

158 +) -> None:

159 + """

160 + Perform architecture-specific manipulations on the input tokens.

161 +

162 + Note:

163 + This is an in-place operation.

164 + """

165 + model = get_model_cls(model_config)

166 + if supports_score_template(model):

167 + model.post_process_tokens(prompt)

168 +

169 +

170 +def get_score_prompt(

171 + model_config: ModelConfig,

172 + tokenizer: AnyTokenizer,

173 + tokenization_kwargs: dict[str, Any],

174 + data_1: Union[str, ScoreContentPartParam],

175 + data_2: Union[str, ScoreContentPartParam],

176 +) -> tuple[str, TokensPrompt]:

177 + prompt_1, prompt_2, mm_data = parse_score_data(

178 + data_1,

179 + data_2,

180 + model_config,

181 + tokenizer,

182 + )

183 +

184 + full_prompt = apply_score_template(model_config, prompt_1, prompt_2)

185 +

186 + prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)

187 +

188 + engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"])

189 +

190 + post_process_tokens(model_config, engine_prompt)

191 +

192 + if mm_data is not None:

193 + engine_prompt["multi_modal_data"] = mm_data

194 + return full_prompt, engine_prompt

vllm/model_executor/models/config.py

+10

@@ -170,6 +170,15 @@ class Qwen3ForSequenceClassificationConfig(VerifyAndUpdateConfig):

170 vllm_config.model_config.hf_config.method = "from_2_way_softmax"

171

172

173 +class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):

174 +

175 + @staticmethod

176 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:

177 + config = vllm_config.model_config.hf_config

178 +

179 + config.num_labels = 1

180 +

181 +

182 class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):

183

184 @staticmethod

@@ -197,4 +206,5 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {

206 "NomicBertModel": NomicBertModelConfig,

207 "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig,

208 "XLMRobertaModel": JinaRobertaModelConfig,

209 + "JinaVLForRanking": JinaVLForSequenceClassificationConfig,

210 }

vllm/model_executor/models/interfaces.py

+57

@@ -9,6 +9,7 @@ import torch

9 from torch import Tensor

10 from typing_extensions import Self, TypeIs

12 +from vllm.inputs import TokensPrompt

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.quantization.base_config import (

15 QuantizationConfig)

@@ -141,6 +142,62 @@ def supports_multimodal(

142 return isinstance(model, SupportsMultiModal)

143

144

145 +@runtime_checkable

146 +class SupportsScoreTemplate(Protocol):

147 + """The interface required for all models that support score template."""

148 +

149 + supports_score_template: ClassVar[Literal[True]] = True

150 + """

151 + A flag that indicates this model supports score template.

152 +

153 + Note:

154 + There is no need to redefine this flag if this class is in the

155 + MRO of your model class.

156 + """

157 +

158 + @classmethod

159 + def get_score_template(cls, query: str, document: str) -> Optional[str]:

160 + """

161 + Generate a full prompt by populating the score template with query and document content.

162 + """ # noqa: E501

163 + ...

164 +

165 + @classmethod

166 + def post_process_tokens(cls, prompt: TokensPrompt) -> None:

167 + """

168 + Perform architecture-specific manipulations on the input tokens.

169 + """

170 + ...

171 +

172 +

173 +# We can't use runtime_checkable with ClassVar for issubclass checks

174 +# so we need to treat the class as an instance and use isinstance instead

175 +@runtime_checkable

176 +class _SupportsScoreTemplateType(Protocol):

177 + supports_score_template: Literal[True]

178 +

179 +

180 +@overload

181 +def supports_score_template(

182 + model: type[object]) -> TypeIs[type[SupportsScoreTemplate]]:

183 + ...

184 +

185 +

186 +@overload

187 +def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]:

188 + ...

189 +

190 +

191 +def supports_score_template(

192 + model: Union[type[object], object],

193 +) -> Union[TypeIs[type[SupportsScoreTemplate]], TypeIs[SupportsScoreTemplate]]:

194 +

195 + if isinstance(model, type):

196 + return isinstance(model, _SupportsScoreTemplateType)

197 +

198 + return isinstance(model, SupportsScoreTemplate)

199 +

200 +

201 @runtime_checkable

202 class SupportsLoRA(Protocol):

203 """The interface required for all models that support LoRA."""

vllm/model_executor/models/jina_vl.py new

+150

@@ -0,0 +1,150 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Iterable, Mapping

4 +from typing import Optional

5 +

6 +import torch

7 +import torch.nn as nn

8 +from transformers import BatchFeature, PretrainedConfig

9 +

10 +from vllm.config import VllmConfig

11 +from vllm.inputs import TokensPrompt

12 +from vllm.logger import init_logger

13 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

14 + RowParallelLinear)

15 +from vllm.model_executor.layers.pooler import Pooler, PoolingType

16 +from vllm.model_executor.pooling_metadata import PoolingMetadata

17 +from vllm.multimodal import MULTIMODAL_REGISTRY

18 +from vllm.sequence import IntermediateTensors, PoolerOutput

19 +

20 +from .interfaces import (SupportsCrossEncoding, SupportsMultiModal,

21 + SupportsScoreTemplate)

22 +from .qwen2_vl import (Qwen2VLDummyInputsBuilder,

23 + Qwen2VLForConditionalGeneration,

24 + Qwen2VLMultiModalProcessor, Qwen2VLProcessingInfo)

25 +from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix

26 +

27 +logger = init_logger(__name__)

28 +

29 +

30 +class JinaVLScorer(nn.Module):

31 +

32 + def __init__(self, config: PretrainedConfig):

33 + super().__init__()

34 + self.dense = ColumnParallelLinear(config.hidden_size,

35 + config.hidden_size,

36 + bias=True)

37 + self.out_proj = RowParallelLinear(config.hidden_size,

38 + config.num_labels,

39 + bias=True)

40 +

41 + def forward(self, x, **kwargs):

42 + x, _ = self.dense(x)

43 + x = torch.relu(x)

44 + x, _ = self.out_proj(x)

45 + return x

46 +

47 +

48 +class JinaVLMultiModalProcessor(Qwen2VLMultiModalProcessor):

49 +

50 + def _call_hf_processor(

51 + self,

52 + prompt: str,

53 + mm_data: Mapping[str, object],

54 + mm_kwargs: Mapping[str, object],

55 + tok_kwargs: Mapping[str, object],

56 + ) -> BatchFeature:

57 +

58 + # NOTE: We should reverse the order of the mm_data because the

59 + # query prompt is placed after the document prompt in the score

60 + # template for JinaVLForRanking model, but in mm_data they are

61 + # stored in the opposite order (query first, then document).

62 + for _, value in mm_data.items():

63 + value.reverse()

64 + return super()._call_hf_processor(prompt, mm_data, mm_kwargs,

65 + tok_kwargs)

66 +

67 +

68 +@MULTIMODAL_REGISTRY.register_processor(JinaVLMultiModalProcessor,

69 + info=Qwen2VLProcessingInfo,

70 + dummy_inputs=Qwen2VLDummyInputsBuilder)

71 +class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,

72 + SupportsCrossEncoding,

73 + SupportsMultiModal,

74 + SupportsScoreTemplate):

75 + weight_mapper = WeightsMapper(

76 + orig_to_new_prefix={

77 + "score.0.": "score.dense.",

78 + "score.2.": "score.out_proj.",

79 + # mapping for new names in checkpoint saved after transformers v4.52

80 + "model.language_model.": "language_model.model.",

81 + "visual.": "visual.",

82 + # mapping for original checkpoint

83 + "lm_head.": "language_model.lm_head.",

84 + "model.": "language_model.model.",

85 + })

86 +

87 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

88 + super().__init__(vllm_config=vllm_config,

89 + prefix=maybe_prefix(prefix, "qwen2_vl"))

90 + config = vllm_config.model_config.hf_config

91 + pooler_config = vllm_config.model_config.pooler_config

92 +

93 + # logit bias for sigmoid normalization

94 + self.LOGIT_BIAS = 2.65

95 +

96 + self.score = JinaVLScorer(config)

97 +

98 + self._pooler = Pooler.from_config_with_defaults(

99 + pooler_config,

100 + pooling_type=PoolingType.LAST,

101 + normalize=False,

102 + softmax=True)

103 +

104 + @classmethod

105 + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

106 + if modality.startswith("image"):

107 + return "<|vision_start|><|image_pad|><|vision_end|>"

108 +

109 + raise ValueError("Only image modality is supported")

110 +

111 + @classmethod

112 + def get_score_template(cls, query: str, document: str) -> Optional[str]:

113 + return f"**Document**:\n{document}\n**Query**:\n{query}"

114 +

115 + @classmethod

116 + def post_process_tokens(cls, prompt: TokensPrompt) -> None:

117 +

118 + # add score target token at the end of prompt tokens

119 + prompt['prompt_token_ids'].append(100)

120 +

121 + def forward(

122 + self,

123 + input_ids: torch.Tensor,

124 + positions: torch.Tensor,

125 + intermediate_tensors: Optional[IntermediateTensors] = None,

126 + inputs_embeds: Optional[torch.Tensor] = None,

127 + **kwargs: object,

128 + ) -> torch.Tensor:

129 + hidden_states = super().forward(

130 + input_ids=input_ids,

131 + positions=positions,

132 + intermediate_tensors=intermediate_tensors,

133 + inputs_embeds=inputs_embeds,

134 + **kwargs,

135 + )

136 +

137 + logits = self.score(hidden_states) - self.LOGIT_BIAS

138 + return logits

139 +

140 + def pooler(

141 + self,

142 + hidden_states: torch.Tensor,

143 + pooling_metadata: PoolingMetadata,

144 + ) -> Optional[PoolerOutput]:

145 + return self._pooler(hidden_states, pooling_metadata)

146 +

147 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

148 +

149 + loader = AutoWeightsLoader(self)

150 + return loader.load_weights(weights, mapper=self.weight_mapper)

vllm/model_executor/models/registry.py

@@ -181,6 +181,7 @@ _CROSS_ENCODER_MODELS = {

181 "GemmaForSequenceClassification": ("gemma", "GemmaForSequenceClassification"), # noqa: E501

182 "Qwen2ForSequenceClassification": ("qwen2", "Qwen2ForSequenceClassification"), # noqa: E501

183 "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501

184 + "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"), # noqa: E501

185 }

186

187 _MULTIMODAL_MODELS = {