@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 164 lines 5.97 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import numpy as np
4 import pytest
5
6 from vllm.config import ModelConfig
7 from vllm.multimodal import MULTIMODAL_REGISTRY
8
9 AUDIO_MODEL_SETTINGS = {
10 "ibm-granite/granite-speech-3.3-2b": {
11 "prompt": (
12 "<|start_of_role|>system<|end_of_role|>"
13 "You are a helpful AI assistant<|end_of_text|>\n"
14 "<|start_of_role|>user<|end_of_role|>"
15 "<|audio|>can you transcribe the speech into a written format?"
16 "<|end_of_text|>\n"
17 "<|start_of_role|>assistant<|end_of_role|>"
18 ),
19 },
20 "nvidia/audio-flamingo-3-hf": {
21 "prompt": (
22 "<|im_start|>system\n"
23 "You are a helpful assistant.<|im_end|>\n"
24 "<|im_start|>user\n"
25 "<sound>Transcribe the input speech.<|im_end|>\n"
26 "<|im_start|>assistant\n"
27 ),
28 },
29 "mistralai/Voxtral-Mini-3B-2507": {
30 "prompt": ("[INST][AUDIO]What can you tell me about this audio?[/INST]"),
31 },
32 "microsoft/VibeVoice-ASR-HF": {
33 "prompt": (
34 "<|im_start|>system\n"
35 "You are a helpful assistant that transcribes audio input "
36 "into text output in JSON format.<|im_end|>\n"
37 "<|im_start|>user\n"
38 "<|object_ref_start|><|box_start|><|object_ref_end|>\n"
39 "This is a 1.0 seconds audio, please transcribe it with "
40 "these keys: Start time, End time, Speaker ID, Content"
41 "<|im_end|>\n"
42 "<|im_start|>assistant\n"
43 ),
44 },
45 "zai-org/GLM-ASR-Nano-2512": {
46 "prompt": (
47 "<|user|>\n"
48 "<|begin_of_audio|><|pad|><|end_of_audio|><|user|>\n"
49 "Please transcribe this audio into text"
50 "<|assistant|>\n"
51 ),
52 },
53 }
54
55
56 @pytest.mark.parametrize(
57 "model_id",
58 [
59 "ibm-granite/granite-speech-3.3-2b",
60 "nvidia/audio-flamingo-3-hf",
61 pytest.param(
62 "mistralai/Voxtral-Mini-3B-2507",
63 marks=pytest.mark.xfail(
64 reason="Voxtral's mistral_common processor does not compose with "
65 "the Transformers modelling backend. Loading it currently fails "
66 "outright, because MistralCommonBackend.from_pretrained rejects "
67 "the kwargs vLLM passes, and it implements no "
68 "`replace_audio_token`, so it would report no replacement "
69 "offsets. Both fixes belong in mistral_common or transformers.",
70 strict=False,
71 ),
72 ),
73 "microsoft/VibeVoice-ASR-HF",
74 "zai-org/GLM-ASR-Nano-2512",
75 ],
76 )
77 def test_audio_multimodal_processor(model_id):
78 settings = AUDIO_MODEL_SETTINGS[model_id]
79
80 mm_processor = MULTIMODAL_REGISTRY.create_processor(
81 ModelConfig(model=model_id, model_impl="transformers")
82 )
83
84 audio = np.zeros(16000, dtype=np.float32)
85 mm_data = {"audio": (audio, 16000)}
86
87 result = mm_processor(
88 prompt=settings["prompt"],
89 mm_items=mm_processor.info.parse_mm_data(mm_data),
90 hf_processor_mm_kwargs={},
91 )
92
93 assert "prompt_token_ids" in result
94 assert len(result["prompt_token_ids"]) > 0
95
96 mm_placeholders = result.get("mm_placeholders", {})
97 assert "audio" in mm_placeholders, f"No audio placeholders found for {model_id}"
98 assert len(mm_placeholders["audio"]) == 1
99
100 placeholder = mm_placeholders["audio"][0]
101 assert placeholder.length > 0
102 assert placeholder.offset >= 0
103
104 audio_items = result.get("mm_kwargs", {}).get("audio", [])
105 assert len(audio_items) == 1, f"Expected 1 audio item, got {len(audio_items)}"
106 item_keys = list(audio_items[0].keys())
107 has_features = "input_features" in item_keys or "input_values" in item_keys
108 assert has_features, (
109 f"No audio features (input_features/input_values) in {item_keys} for {model_id}"
110 )
111
112
113 @pytest.mark.parametrize("separator", [" and ", ""])
114 def test_audio_multiple_inputs(separator):
115 """Multiple audios per prompt are each detected as a separate placeholder
116 and multi-modal item by the Transformers modelling backend."""
117 model_id = "ibm-granite/granite-speech-3.3-2b"
118 mm_processor = MULTIMODAL_REGISTRY.create_processor(
119 ModelConfig(model=model_id, model_impl="transformers")
120 )
121
122 audio_token = mm_processor.info.get_hf_processor().audio_token
123 # One token per audio; the processor expands each to its placeholder run.
124 prompt = (
125 "<|start_of_role|>user<|end_of_role|>"
126 f"{audio_token}{separator}{audio_token} transcribe<|end_of_text|>\n"
127 )
128 audios = [np.zeros(16000, dtype=np.float32), np.zeros(24000, dtype=np.float32)]
129
130 result = mm_processor(
131 prompt=prompt,
132 mm_items=mm_processor.info.parse_mm_data({"audio": audios}),
133 hf_processor_mm_kwargs={},
134 )
135
136 assert len(result["mm_placeholders"]["audio"]) == 2
137 assert len(result["mm_kwargs"]["audio"]) == 2
138
139
140 def test_audio_fields_not_claimed_by_image():
141 """Audio fields survive when the image branch is also active."""
142 model_id = "ibm-granite/granite-speech-3.3-2b"
143 model_config = ModelConfig(model=model_id, model_impl="transformers")
144 mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)
145
146 audio_keys = ["input_features", "input_features_mask"]
147 owned = mm_processor._partition_keys_by_modality(audio_keys, ["audio", "image"])
148
149 assert owned["audio"] == audio_keys
150 assert owned["image"] == []
151
152
153 def test_unclaimed_fields_warn_rather_than_raise():
154 """Keys no sub-processor declares are dropped with a warning, not an error."""
155 model_id = "ibm-granite/granite-speech-3.3-2b"
156 model_config = ModelConfig(model=model_id, model_impl="transformers")
157 mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)
158
159 owned = mm_processor._partition_keys_by_modality(
160 ["input_features", "surprise_field"], ["audio", "image"]
161 )
162
163 assert owned["audio"] == ["input_features"]
164 assert owned["image"] == []