"vscode:/vscode.git/clone" did not exist on "5de8d9f111242671beb6fdc28367df8796ce638f"
vision_language_pooling.py 11.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
Cyrus Leung's avatar
Cyrus Leung committed
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 pooling.
Cyrus Leung's avatar
Cyrus Leung committed
6
7
8
9

For most models, the prompt format should follow corresponding examples
on HuggingFace model repository.
"""
10

Cyrus Leung's avatar
Cyrus Leung committed
11
from argparse import Namespace
12
from dataclasses import asdict
13
from pathlib import Path
Cyrus Leung's avatar
Cyrus Leung committed
14
15
16
17
from typing import Literal, NamedTuple, Optional, TypedDict, Union, get_args

from PIL.Image import Image

18
from vllm import LLM, EngineArgs
19
from vllm.entrypoints.score_utils import ScoreMultiModalParam
Cyrus Leung's avatar
Cyrus Leung committed
20
21
22
from vllm.multimodal.utils import fetch_image
from vllm.utils import FlexibleArgumentParser

23
24
25
ROOT_DIR = Path(__file__).parent.parent.parent
EXAMPLES_DIR = ROOT_DIR / "examples"

Cyrus Leung's avatar
Cyrus Leung committed
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

class TextQuery(TypedDict):
    modality: Literal["text"]
    text: str


class ImageQuery(TypedDict):
    modality: Literal["image"]
    image: Image


class TextImageQuery(TypedDict):
    modality: Literal["text+image"]
    text: str
    image: Image


43
44
45
46
47
48
49
50
class TextImagesQuery(TypedDict):
    modality: Literal["text+images"]
    text: str
    image: ScoreMultiModalParam


QueryModality = Literal["text", "image", "text+image", "text+images"]
Query = Union[TextQuery, ImageQuery, TextImageQuery, TextImagesQuery]
Cyrus Leung's avatar
Cyrus Leung committed
51
52
53


class ModelRequestData(NamedTuple):
54
    engine_args: EngineArgs
55
56
57
58
    prompt: Optional[str] = None
    image: Optional[Image] = None
    query: Optional[str] = None
    documents: Optional[ScoreMultiModalParam] = None
Cyrus Leung's avatar
Cyrus Leung committed
59
60


61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def run_clip(query: Query) -> ModelRequestData:
    if query["modality"] == "text":
        prompt = query["text"]
        image = None
    elif query["modality"] == "image":
        prompt = ""  # For image input, make sure that the prompt text is empty
        image = query["image"]
    else:
        modality = query["modality"]
        raise ValueError(f"Unsupported query modality: '{modality}'")

    engine_args = EngineArgs(
        model="openai/clip-vit-base-patch32",
        runner="pooling",
        limit_mm_per_prompt={"image": 1},
    )

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
        image=image,
    )


85
def run_e5_v(query: Query) -> ModelRequestData:
86
    llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n"  # noqa: E501
Cyrus Leung's avatar
Cyrus Leung committed
87
88
89

    if query["modality"] == "text":
        text = query["text"]
90
        prompt = llama3_template.format(f"{text}\nSummary above sentence in one word: ")
Cyrus Leung's avatar
Cyrus Leung committed
91
92
        image = None
    elif query["modality"] == "image":
93
        prompt = llama3_template.format("<image>\nSummary above image in one word: ")
Cyrus Leung's avatar
Cyrus Leung committed
94
95
        image = query["image"]
    else:
96
        modality = query["modality"]
Cyrus Leung's avatar
Cyrus Leung committed
97
98
        raise ValueError(f"Unsupported query modality: '{modality}'")

99
    engine_args = EngineArgs(
Cyrus Leung's avatar
Cyrus Leung committed
100
        model="royokong/e5-v",
101
        runner="pooling",
Cyrus Leung's avatar
Cyrus Leung committed
102
        max_model_len=4096,
103
        limit_mm_per_prompt={"image": 1},
Cyrus Leung's avatar
Cyrus Leung committed
104
105
106
    )

    return ModelRequestData(
107
        engine_args=engine_args,
Cyrus Leung's avatar
Cyrus Leung committed
108
109
110
111
112
        prompt=prompt,
        image=image,
    )


113
def _get_vlm2vec_prompt_image(query: Query, image_token: str):
Cyrus Leung's avatar
Cyrus Leung committed
114
115
116
117
118
    if query["modality"] == "text":
        text = query["text"]
        prompt = f"Find me an everyday image that matches the given caption: {text}"  # noqa: E501
        image = None
    elif query["modality"] == "image":
119
        prompt = f"{image_token} Find a day-to-day image that looks similar to the provided image."  # noqa: E501
Cyrus Leung's avatar
Cyrus Leung committed
120
121
122
        image = query["image"]
    elif query["modality"] == "text+image":
        text = query["text"]
123
        prompt = f"{image_token} Represent the given image with the following question: {text}"  # noqa: E501
Cyrus Leung's avatar
Cyrus Leung committed
124
125
        image = query["image"]
    else:
126
        modality = query["modality"]
127
128
129
130
131
132
133
        raise ValueError(f"Unsupported query modality: {modality!r}")

    return prompt, image


def run_vlm2vec_phi3v(query: Query) -> ModelRequestData:
    prompt, image = _get_vlm2vec_prompt_image(query, "<|image_1|>")
Cyrus Leung's avatar
Cyrus Leung committed
134

135
    engine_args = EngineArgs(
Cyrus Leung's avatar
Cyrus Leung committed
136
        model="TIGER-Lab/VLM2Vec-Full",
137
        runner="pooling",
138
        max_model_len=4096,
Cyrus Leung's avatar
Cyrus Leung committed
139
140
        trust_remote_code=True,
        mm_processor_kwargs={"num_crops": 4},
141
        limit_mm_per_prompt={"image": 1},
Cyrus Leung's avatar
Cyrus Leung committed
142
143
144
    )

    return ModelRequestData(
145
        engine_args=engine_args,
Cyrus Leung's avatar
Cyrus Leung committed
146
147
148
149
150
        prompt=prompt,
        image=image,
    )


151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def run_vlm2vec_qwen2vl(query: Query) -> ModelRequestData:
    # vLLM does not support LoRA adapters on multi-modal encoder,
    # so we merge the weights first
    from huggingface_hub.constants import HF_HUB_CACHE
    from peft import PeftConfig, PeftModel
    from transformers import AutoModelForImageTextToText, AutoProcessor

    from vllm.entrypoints.chat_utils import load_chat_template

    model_id = "TIGER-Lab/VLM2Vec-Qwen2VL-2B"

    base_model = AutoModelForImageTextToText.from_pretrained(model_id)
    lora_model = PeftModel.from_pretrained(
        base_model,
        model_id,
        config=PeftConfig.from_pretrained(model_id),
    )
    model = lora_model.merge_and_unload().to(dtype=base_model.dtype)
    model._hf_peft_config_loaded = False  # Needed to save the merged model

    processor = AutoProcessor.from_pretrained(
        model_id,
173
174
        # `min_pixels` and `max_pixels` are deprecated for
        # transformers `preprocessor_config.json`
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
        size={"shortest_edge": 3136, "longest_edge": 12845056},
    )
    processor.chat_template = load_chat_template(
        # The original chat template is not correct
        EXAMPLES_DIR / "template_vlm2vec_qwen2vl.jinja",
    )

    merged_path = str(
        Path(HF_HUB_CACHE) / ("models--" + model_id.replace("/", "--") + "-vllm")
    )
    print(f"Saving merged model to {merged_path}...")
    print(
        "NOTE: This directory is not tracked by `huggingface_hub` "
        "so you have to delete this manually if you don't want it anymore."
    )
    model.save_pretrained(merged_path)
    processor.save_pretrained(merged_path)
    print("Done!")

    prompt, image = _get_vlm2vec_prompt_image(query, "<|image_pad|>")

    engine_args = EngineArgs(
        model=merged_path,
        runner="pooling",
        max_model_len=4096,
200
201
202
203
        mm_processor_kwargs={
            "min_pixels": 3136,
            "max_pixels": 12845056,
        },
204
205
206
207
208
209
210
211
212
213
        limit_mm_per_prompt={"image": 1},
    )

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
        image=image,
    )


214
215
216
217
218
219
def run_jinavl_reranker(query: Query) -> ModelRequestData:
    if query["modality"] != "text+images":
        raise ValueError(f"Unsupported query modality: '{query['modality']}'")

    engine_args = EngineArgs(
        model="jinaai/jina-reranker-m0",
220
        runner="pooling",
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
        max_model_len=32768,
        trust_remote_code=True,
        mm_processor_kwargs={
            "min_pixels": 3136,
            "max_pixels": 602112,
        },
        limit_mm_per_prompt={"image": 1},
    )

    return ModelRequestData(
        engine_args=engine_args,
        query=query["text"],
        documents=query["image"],
    )


Cyrus Leung's avatar
Cyrus Leung committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
def get_query(modality: QueryModality):
    if modality == "text":
        return TextQuery(modality="text", text="A dog sitting in the grass")

    if modality == "image":
        return ImageQuery(
            modality="image",
            image=fetch_image(
                "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/American_Eskimo_Dog.jpg/360px-American_Eskimo_Dog.jpg"  # noqa: E501
            ),
        )

    if modality == "text+image":
        return TextImageQuery(
            modality="text+image",
            text="A cat standing in the snow.",
            image=fetch_image(
                "https://upload.wikimedia.org/wikipedia/commons/thumb/b/b6/Felis_catus-cat_on_snow.jpg/179px-Felis_catus-cat_on_snow.jpg"  # noqa: E501
            ),
        )

258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
    if modality == "text+images":
        return TextImagesQuery(
            modality="text+images",
            text="slm markdown",
            image={
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
                        },
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
                        },
                    },
                ]
            },
        )

Cyrus Leung's avatar
Cyrus Leung committed
280
281
282
283
    msg = f"Modality {modality} is not supported."
    raise ValueError(msg)


284
def run_encode(model: str, modality: QueryModality, seed: Optional[int]):
Cyrus Leung's avatar
Cyrus Leung committed
285
286
287
    query = get_query(modality)
    req_data = model_example_map[model](query)

288
289
290
    # Disable other modalities to save memory
    default_limits = {"image": 0, "video": 0, "audio": 0}
    req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
291
292
        req_data.engine_args.limit_mm_per_prompt or {}
    )
293

294
295
296
    engine_args = asdict(req_data.engine_args) | {"seed": seed}
    llm = LLM(**engine_args)

Cyrus Leung's avatar
Cyrus Leung committed
297
298
299
300
    mm_data = {}
    if req_data.image is not None:
        mm_data["image"] = req_data.image

301
302
303
304
305
306
    outputs = llm.embed(
        {
            "prompt": req_data.prompt,
            "multi_modal_data": mm_data,
        }
    )
Cyrus Leung's avatar
Cyrus Leung committed
307

308
    print("-" * 50)
Cyrus Leung's avatar
Cyrus Leung committed
309
310
    for output in outputs:
        print(output.outputs.embedding)
311
        print("-" * 50)
Cyrus Leung's avatar
Cyrus Leung committed
312
313


314
315
316
317
318
319
320
321
322
323
324
325
326
327
def run_score(model: str, modality: QueryModality, seed: Optional[int]):
    query = get_query(modality)
    req_data = model_example_map[model](query)

    engine_args = asdict(req_data.engine_args) | {"seed": seed}
    llm = LLM(**engine_args)

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

    print("-" * 30)
    print([output.outputs.score for output in outputs])
    print("-" * 30)


Cyrus Leung's avatar
Cyrus Leung committed
328
model_example_map = {
329
    "clip": run_clip,
Cyrus Leung's avatar
Cyrus Leung committed
330
    "e5_v": run_e5_v,
331
332
    "vlm2vec_phi3v": run_vlm2vec_phi3v,
    "vlm2vec_qwen2vl": run_vlm2vec_qwen2vl,
333
    "jinavl_reranker": run_jinavl_reranker,
Cyrus Leung's avatar
Cyrus Leung committed
334
335
}

336
337

def parse_args():
Cyrus Leung's avatar
Cyrus Leung committed
338
    parser = FlexibleArgumentParser(
339
        description="Demo on using vLLM for offline inference with "
340
        "vision language models for multimodal pooling tasks."
341
342
343
344
345
    )
    parser.add_argument(
        "--model-name",
        "-m",
        type=str,
346
        default="vlm2vec_phi3v",
347
348
349
        choices=model_example_map.keys(),
        help="The name of the embedding model.",
    )
350
351
352
353
354
355
356
357
    parser.add_argument(
        "--task",
        "-t",
        type=str,
        default="embedding",
        choices=["embedding", "scoring"],
        help="The task type.",
    )
358
359
360
361
362
363
364
365
366
367
368
369
370
    parser.add_argument(
        "--modality",
        type=str,
        default="image",
        choices=get_args(QueryModality),
        help="Modality of the input.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Set the seed when initializing `vllm.LLM`.",
    )
371
    return parser.parse_args()
372

373
374

def main(args: Namespace):
375
376
377
378
379
380
    if args.task == "embedding":
        run_encode(args.model_name, args.modality, args.seed)
    elif args.task == "scoring":
        run_score(args.model_name, args.modality, args.seed)
    else:
        raise ValueError(f"Unsupported task: {args.task}")
381
382
383
384


if __name__ == "__main__":
    args = parse_args()
Cyrus Leung's avatar
Cyrus Leung committed
385
    main(args)