@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 155 lines 4.34 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import Any
5
6 import pytest
7 import torch
8 from transformers import SiglipModel
9
10 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
11 from ...utils import check_embeddings_close
12
13 HF_TEXT_PROMPTS = [
14 "a photo of a stop sign",
15 "a photo of a cherry blossom",
16 ]
17
18 HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
19 {
20 "stop_sign": "",
21 "cherry_blossom": "",
22 }
23 )
24
25 MODELS = [
26 "google/siglip-base-patch16-224",
27 "google/siglip2-base-patch16-224",
28 # Different image embedding dim than text_config.hidden_size
29 "google/siglip2-giant-opt-patch16-384",
30 ]
31
32
33 def _run_test(
34 hf_runner: type[HfRunner],
35 vllm_runner: type[VllmRunner],
36 input_cases: list[tuple[list[str], PromptImageInput, dict[str, Any]]],
37 model: str,
38 *,
39 dtype: str,
40 ) -> None:
41 with vllm_runner(
42 model,
43 runner="pooling",
44 dtype=dtype,
45 enforce_eager=True,
46 max_model_len=64,
47 gpu_memory_utilization=0.7,
48 ) as vllm_model:
49 vllm_outputs_per_case = [
50 vllm_model.embed(
51 input_texts,
52 images=input_images,
53 tokenization_kwargs=tokenization_kwargs,
54 )
55 for input_texts, input_images, tokenization_kwargs in input_cases
56 ]
57
58 texts = [HF_TEXT_PROMPTS[0]]
59 images = [input_cases[1][1][0]]
60 with pytest.raises(ValueError, match="not both"):
61 vllm_model.embed(texts, images=images)
62
63 vllm_model.embed(texts)
64 vllm_model.embed([""], images=images)
65
66 mixed_outputs = vllm_outputs_per_case[2]
67 check_embeddings_close(
68 embeddings_0_lst=[vllm_outputs_per_case[0][0]],
69 embeddings_1_lst=[mixed_outputs[0]],
70 name_0="text_only",
71 name_1="mixed_text",
72 )
73 check_embeddings_close(
74 embeddings_0_lst=[vllm_outputs_per_case[1][0]],
75 embeddings_1_lst=[mixed_outputs[1]],
76 name_0="image_only",
77 name_1="mixed_image",
78 )
79
80 with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model:
81 hf_outputs_per_case = []
82 for input_texts, input_images, tokenization_kwargs in input_cases:
83 all_inputs = hf_model.get_inputs(
84 input_texts,
85 images=input_images,
86 tokenization_kwargs=tokenization_kwargs,
87 )
88
89 hf_outputs = []
90 for inputs in all_inputs:
91 inputs = hf_model.wrap_device(inputs)
92
93 if "pixel_values" in inputs:
94 pooled_output = hf_model.model.get_image_features(
95 pixel_values=inputs.pixel_values,
96 )
97 else:
98 pooled_output = hf_model.model.get_text_features(
99 input_ids=inputs.input_ids,
100 )
101
102 if not isinstance(pooled_output, torch.Tensor):
103 pooled_output = pooled_output.pooler_output
104 pooled_output = pooled_output.squeeze(0)
105 hf_outputs.append(pooled_output.tolist())
106
107 hf_outputs_per_case.append(hf_outputs)
108
109 for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
110 check_embeddings_close(
111 embeddings_0_lst=hf_outputs,
112 embeddings_1_lst=vllm_outputs,
113 name_0="hf",
114 name_1="vllm",
115 )
116
117
118 @pytest.mark.parametrize("model", MODELS)
119 @pytest.mark.parametrize("dtype", ["float"])
120 def test_models(
121 hf_runner,
122 vllm_runner,
123 image_assets,
124 model: str,
125 dtype: str,
126 ) -> None:
127 text_images = [None] * len(HF_TEXT_PROMPTS)
128 images = [asset.pil_image for asset in image_assets]
129 input_cases = [
130 (
131 HF_TEXT_PROMPTS,
132 text_images,
133 {
134 "padding": "max_length",
135 "max_length": 64,
136 },
137 ),
138 (HF_IMAGE_PROMPTS, images, {}),
139 (
140 [HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]],
141 [None, images[0]],
142 {
143 "padding": "max_length",
144 "max_length": 64,
145 },
146 ),
147 ]
148
149 _run_test(
150 hf_runner,
151 vllm_runner,
152 input_cases,
153 model,
154 dtype=dtype,
155 )