main
py 154 lines 4.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Copyright 2025 The vLLM team.
5 # Copyright 2025 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
6 # reserved.
7 #
8 # Licensed under the Apache License, Version 2.0 (the "License");
9 # you may not use this file except in compliance with the License.
10 # You may obtain a copy of the License at
11 #
12 # http://www.apache.org/licenses/LICENSE-2.0
13 #
14 # Unless required by applicable law or agreed to in writing, software
15 # distributed under the License is distributed on an "AS IS" BASIS,
16 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 # See the License for the specific language governing permissions and
18 # limitations under the License.
19
20 from unittest.mock import MagicMock
21
22 import numpy as np
23 import pytest
24 import torch
25 from transformers import PreTrainedConfig
26
27 from tests.models.registry import HF_EXAMPLE_MODELS
28 from vllm.config.multimodal import MultiModalDummyOptions
29
30
31 class MockAudioFlamingo3Config(PreTrainedConfig):
32 model_type = "audioflamingo3"
33
34 def __init__(self, **kwargs):
35 super().__init__(**kwargs)
36 self.audio_config = PreTrainedConfig()
37 self.text_config = PreTrainedConfig()
38
39
40 class MockAudioFlamingo3Processor:
41 def __init__(self):
42 self.audio_token = "<sound>"
43 self.audio_token_id = 12345
44 self.max_audio_len = 60
45 self.feature_extractor = MockFeatureExtractor()
46 self.tokenizer = self._tokenize
47
48 def __call__(self, text=None, audio=None, **kwargs):
49 return {
50 "input_ids": torch.tensor([[1, 2, 3]], dtype=torch.long),
51 "input_features": torch.zeros((3, 80, 3000)),
52 "input_features_mask": torch.ones((3, 3000), dtype=torch.long),
53 }
54
55 def _tokenize(self, text, **kwargs):
56 return {"input_ids": torch.tensor([[1, 2, 3]], dtype=torch.long)}
57
58
59 class MockFeatureExtractor:
60 def __init__(self):
61 self.sampling_rate = 16000
62 self.chunk_length = 30
63 self.hop_length = 160
64
65 def __call__(self, audios, **kwargs):
66 return {
67 "input_features": torch.zeros((len(audios), 80, 3000)),
68 "attention_mask": torch.ones((len(audios), 3000), dtype=torch.long),
69 }
70
71
72 @pytest.fixture
73 def mock_ctx():
74 config = MockAudioFlamingo3Config()
75
76 ctx = MagicMock()
77 ctx.get_hf_config.return_value = config
78 ctx.get_hf_processor.return_value = MockAudioFlamingo3Processor()
79 ctx.call_hf_processor.side_effect = lambda processor, data, kwargs: processor(
80 **data, **kwargs
81 )
82 ctx.model_config.hf_config = config
83 return ctx
84
85
86 @pytest.fixture(autouse=True)
87 def check_transformers_version():
88 model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
89 model_info.check_transformers_version(on_fail="skip")
90
91
92 def test_audio_chunk_counting(mock_ctx):
93 from vllm.model_executor.models.audioflamingo3 import (
94 AudioFlamingo3DummyInputsBuilder,
95 AudioFlamingo3MultiModalProcessor,
96 AudioFlamingo3ProcessingInfo,
97 )
98
99 info = AudioFlamingo3ProcessingInfo(mock_ctx)
100 processor = AudioFlamingo3MultiModalProcessor(
101 info, AudioFlamingo3DummyInputsBuilder(info)
102 )
103
104 sr = 16000
105 audio_1 = np.zeros(30 * sr)
106 audio_2 = np.zeros(75 * sr)
107
108 mm_data = {"audio": [audio_1, audio_2]}
109
110 mm_items = processor.info.parse_mm_data(mm_data, validate=False)
111 processed = processor._apply_hf_processor_main(mm_items, {})
112
113 chunk_counts = processed["chunk_counts"]
114
115 assert chunk_counts[0].item() == 1
116 assert chunk_counts[1].item() == 2
117 assert len(chunk_counts) == 2
118 assert processed["feature_attention_mask"].shape == (3, 3000)
119
120
121 def test_dummy_data_generation(mock_ctx):
122 from vllm.model_executor.models.audioflamingo3 import (
123 AudioFlamingo3DummyInputsBuilder,
124 AudioFlamingo3ProcessingInfo,
125 )
126
127 info = AudioFlamingo3ProcessingInfo(mock_ctx)
128 builder = AudioFlamingo3DummyInputsBuilder(info)
129
130 mm_counts = {"audio": 2}
131 dummy_data = builder.get_dummy_mm_data(100, mm_counts, MultiModalDummyOptions())
132
133 assert "audio" in dummy_data
134 assert len(dummy_data["audio"]) == 2
135
136 expected_len = 60 * 16000
137 assert len(dummy_data["audio"][0]) == expected_len
138
139
140 def test_audio_token_count_matches_hf_processor_math():
141 from vllm.model_executor.models.audioflamingo3 import (
142 _count_audio_tokens_from_mask,
143 )
144
145 feature_attention_mask = torch.zeros((3, 3000), dtype=torch.long)
146 feature_attention_mask[0, :2999] = 1
147 feature_attention_mask[1, :2999] = 1
148 feature_attention_mask[2, :1500] = 1
149 chunk_counts = torch.tensor([2, 1], dtype=torch.long)
150
151 assert (
152 _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 0) == 1499
153 )
154 assert _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 1) == 375