main
py 79 lines 2.67 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from vllm.assets.image import ImageAsset
7 from vllm.assets.video import VideoAsset
8 from vllm.multimodal.image import convert_image_mode
9
10 models = ["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"]
11
12
13 def base_prompt(modalities_str: str) -> str:
14 return f"<|im_start|>user {modalities_str}\nDescribe what you see from these items.<|im_end|><|im_start|>assistant\n" # noqa: E501
15
16
17 INTERLEAVED_PROMPT = base_prompt("<image><video><image>\n")
18 NONINTERLEAVED_PROMPT = base_prompt("<image><image><video>\n")
19
20
21 @pytest.mark.core_model
22 @pytest.mark.parametrize("model", models)
23 @pytest.mark.parametrize("dtype", ["float16"])
24 @pytest.mark.parametrize("max_tokens", [128])
25 def test_models(vllm_runner, model, dtype: str, max_tokens: int) -> None:
26 """This is a simple test to check if interleaved and non-interleaved prompts
27 give the same result.
28 """
29 image_cherry = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
30 image_stop = convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB")
31 images = [image_cherry, image_stop]
32 video = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays
33
34 inputs = [
35 (
36 [INTERLEAVED_PROMPT],
37 [images],
38 [video],
39 ),
40 (
41 [NONINTERLEAVED_PROMPT],
42 [images],
43 [video],
44 ),
45 ]
46
47 with vllm_runner(
48 model,
49 runner="generate",
50 dtype=dtype,
51 limit_mm_per_prompt={"image": 2},
52 max_model_len=32768,
53 max_num_seqs=2,
54 tensor_parallel_size=1,
55 enforce_eager=True,
56 ) as vllm_model:
57 vllm_outputs_per_case = [
58 vllm_model.generate_greedy(
59 prompts, max_tokens, images=images, videos=videos
60 )
61 for prompts, images, videos in inputs
62 ]
63
64 all_results = [output[0][1] for output in vllm_outputs_per_case]
65 outputs = [
66 (total_str, total_str.find("assistant\n") + len("assistant\n"))
67 for total_str in all_results
68 ]
69 prompt_lengths = [prompt_len for _, prompt_len in outputs]
70 generated_strs = [total_str[prompt_len:] for total_str, prompt_len in outputs]
71 interleaved_prompt_len, noninterleaved_prompt_len = prompt_lengths
72 interleaved_output_str, noninterleaved_output_str = generated_strs
73
74 # The two prompts are identical except for the order of modality tokens.
75 assert interleaved_prompt_len == noninterleaved_prompt_len
76
77 # The two generated strings should be different because of the
78 # interleaved modality tokens.
79 assert interleaved_output_str != noninterleaved_output_str