@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 189 lines 5.74 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # ruff: noqa
5 import argparse
6
7 from vllm import LLM
8 from vllm.sampling_params import SamplingParams
9 from vllm.assets.image import ImageAsset
10 from vllm.multimodal.utils import encode_image_url
11
12 # This script is an offline demo for running Mistral-Small-3.1
13 #
14 # If you want to run a server/client setup, please follow this code:
15 #
16 # - Server:
17 #
18 # ```bash
19 # # Mistral format
20 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
21 # --tokenizer-mode mistral --config-format mistral --load-format mistral \
22 # --limit-mm-per-prompt.image 4 --max-model-len 16384
23 #
24 # # HF format
25 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
26 # --limit-mm-per-prompt.image 4 --max-model-len 16384
27 # ```
28 #
29 # - Client:
30 #
31 # ```bash
32 # curl --location 'http://<your-node-url>:8000/v1/chat/completions' \
33 # --header 'Content-Type: application/json' \
34 # --header 'Authorization: Bearer token' \
35 # --data '{
36 # "model": "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
37 # "messages": [
38 # {
39 # "role": "user",
40 # "content": [
41 # {"type" : "text", "text": "Describe this image in detail please."},
42 # {"type": "image_url", "image_url": {"url": "https://s3.amazonaws.com/cms.ipressroom.com/338/files/201808/5b894ee1a138352221103195_A680%7Ejogging-edit/A680%7Ejogging-edit_hero.jpg"}},
43 # {"type" : "text", "text": "and this one as well. Answer in French."},
44 # {"type": "image_url", "image_url": {"url": "https://www.wolframcloud.com/obj/resourcesystem/images/a0e/a0ee3983-46c6-4c92-b85d-059044639928/6af8cfb971db031b.png"}}
45 # ]
46 # }
47 # ]
48 # }'
49 # ```
50 #
51 # Usage:
52 # python demo.py simple
53 # python demo.py advanced
54
55 # Lower max_model_len and/or max_num_seqs on low-VRAM GPUs.
56 # These scripts have been tested on 2x L40 GPUs
57
58
59 def run_simple_demo(args: argparse.Namespace):
60 model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
61 sampling_params = SamplingParams(max_tokens=8192)
62
63 llm = LLM(
64 model=model_name,
65 tokenizer_mode="mistral" if args.format == "mistral" else "hf",
66 config_format="mistral" if args.format == "mistral" else "hf",
67 load_format="mistral" if args.format == "mistral" else "hf",
68 limit_mm_per_prompt={"image": 1},
69 max_model_len=4096,
70 max_num_seqs=2,
71 tensor_parallel_size=2,
72 mm_processor_cache_gb=0 if args.disable_mm_processor_cache else 4,
73 )
74
75 prompt = "Describe this image in one sentence."
76
77 messages = [
78 {
79 "role": "user",
80 "content": [
81 {"type": "text", "text": prompt},
82 {
83 "type": "image_url",
84 "image_url": {
85 "url": encode_image_url(ImageAsset("cherry_blossom").pil_image)
86 },
87 },
88 ],
89 },
90 ]
91 outputs = llm.chat(messages, sampling_params=sampling_params)
92 print("-" * 50)
93 print(outputs[0].outputs[0].text)
94 print("-" * 50)
95
96
97 def run_advanced_demo(args: argparse.Namespace):
98 model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
99 max_img_per_msg = 3
100 max_tokens_per_img = 4096
101
102 sampling_params = SamplingParams(max_tokens=8192, temperature=0.7)
103 llm = LLM(
104 model=model_name,
105 tokenizer_mode="mistral" if args.format == "mistral" else "hf",
106 config_format="mistral" if args.format == "mistral" else "hf",
107 load_format="mistral" if args.format == "mistral" else "hf",
108 limit_mm_per_prompt={"image": max_img_per_msg},
109 max_model_len=max_img_per_msg * max_tokens_per_img,
110 tensor_parallel_size=2,
111 mm_processor_cache_gb=0 if args.disable_mm_processor_cache else 4,
112 )
113
114 prompt = "Describe the following image."
115
116 url_1 = "https://huggingface.co/datasets/patrickvonplaten/random_img/resolve/main/yosemite.png"
117 url_2 = "https://picsum.photos/seed/picsum/200/300"
118 url_3 = "https://picsum.photos/id/32/512/512"
119
120 messages = [
121 {
122 "role": "user",
123 "content": [
124 {"type": "text", "text": prompt},
125 {"type": "image_url", "image_url": {"url": url_1}},
126 {"type": "image_url", "image_url": {"url": url_2}},
127 ],
128 },
129 {
130 "role": "assistant",
131 "content": "The images show nature.",
132 },
133 {
134 "role": "user",
135 "content": "More details please and answer only in French!.",
136 },
137 {
138 "role": "user",
139 "content": [
140 {"type": "image_url", "image_url": {"url": url_3}},
141 ],
142 },
143 ]
144
145 outputs = llm.chat(messages=messages, sampling_params=sampling_params)
146 print("-" * 50)
147 print(outputs[0].outputs[0].text)
148 print("-" * 50)
149
150
151 def parse_args():
152 parser = argparse.ArgumentParser(
153 description="Run a demo in simple or advanced mode."
154 )
155
156 parser.add_argument(
157 "mode",
158 choices=["simple", "advanced"],
159 help="Specify the demo mode: 'simple' or 'advanced'",
160 )
161
162 parser.add_argument(
163 "--format",
164 choices=["mistral", "hf"],
165 default="mistral",
166 help="Specify the format of the model to load.",
167 )
168
169 parser.add_argument(
170 "--disable-mm-processor-cache",
171 action="store_true",
172 help="If True, disables caching of multi-modal processor.",
173 )
174 return parser.parse_args()
175
176
177 def main():
178 args = parse_args()
179
180 if args.mode == "simple":
181 print("Running simple demo...")
182 run_simple_demo(args)
183 elif args.mode == "advanced":
184 print("Running advanced demo...")
185 run_advanced_demo(args)
186
187
188 if __name__ == "__main__":
189 main()