qwen2_audio.py 16.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# Copyright 2024 The Qwen team.
# Copyright 2023 The vLLM team.
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Inference-only Qwen2-Audio model compatible with HuggingFace weights."""
22
from functools import cached_property
23
24
from typing import (Any, Iterable, List, Mapping, Optional, Set, Tuple,
                    TypedDict, Union)
25
26
27

import torch
import torch.nn as nn
28
from transformers import BatchFeature
29
30
31
32
from transformers.models.qwen2_audio import (Qwen2AudioConfig,
                                             Qwen2AudioEncoder,
                                             Qwen2AudioProcessor)
from transformers.models.whisper import WhisperFeatureExtractor
33
34

from vllm.attention import AttentionMetadata
35
from vllm.config import VllmConfig
Joe Runde's avatar
Joe Runde committed
36
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
37
from vllm.model_executor.sampling_metadata import SamplingMetadata
38
from vllm.multimodal import MULTIMODAL_REGISTRY
39
40
from vllm.multimodal.inputs import (MultiModalFieldConfig, MultiModalKwargs,
                                    NestedTensors)
41
42
from vllm.multimodal.parse import (AudioProcessorItems, MultiModalDataItems,
                                   MultiModalDataParser)
43
from vllm.multimodal.processing import (BaseMultiModalProcessor,
44
45
                                        BaseProcessingInfo, PromptReplacement)
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
46
from vllm.sequence import IntermediateTensors
47
48

from .interfaces import SupportsMultiModal, SupportsPP
49
50
from .utils import (AutoWeightsLoader, init_vllm_registered_model,
                    maybe_prefix, merge_multimodal_embeddings)
51
52
53
54
55


# # === Audio Inputs === #
class Qwen2AudioInputs(TypedDict):
    input_features: torch.Tensor
56
    """Shape: `(num_audios, num_mel_bins, 3000)`"""
57
58

    feature_attention_mask: torch.Tensor
59
    """Shape: `(num_audios, 3000)`"""
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75


# === Audio Encoder === #


class Qwen2AudioMultiModalProjector(nn.Module):

    def __init__(self, audio_hidden_size: int, text_hidden_size: int):
        super().__init__()
        self.linear = nn.Linear(audio_hidden_size, text_hidden_size, bias=True)

    def forward(self, audio_features):
        hidden_states = self.linear(audio_features)
        return hidden_states


76
# From Qwen2AudioEncoder._get_feat_extract_output_lengths
77
def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
78
79
80
    feat_lengths = (input_lengths - 1) // 2 + 1
    output_lengths = (feat_lengths - 2) // 2 + 1
    return feat_lengths, output_lengths
81
82


83
class Qwen2AudioProcessingInfo(BaseProcessingInfo):
84

85
    def get_hf_config(self):
86
87
        return self.ctx.get_hf_config(Qwen2AudioConfig)

88
    def get_hf_processor(
89
90
91
92
93
94
95
        self,
        *,
        # Ignored in initialization
        sampling_rate: Optional[int] = None,
    ) -> Qwen2AudioProcessor:
        return self.ctx.get_hf_processor(Qwen2AudioProcessor)

96
    def get_feature_extractor(
97
98
99
100
101
        self,
        *,
        # Ignored in initialization
        sampling_rate: Optional[int] = None,
    ) -> WhisperFeatureExtractor:
102
        hf_processor = self.get_hf_processor(sampling_rate=sampling_rate)
103
104
105
106
        feature_extractor = hf_processor.feature_extractor  # type: ignore
        assert isinstance(feature_extractor, WhisperFeatureExtractor)
        return feature_extractor

107
108
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"audio": None}
109

110
    def get_mm_max_tokens_per_item(self, seq_len: int) -> Mapping[str, int]:
111
        hf_config = self.get_hf_config()
112
113
114
115
        max_source_positions = hf_config.audio_config.max_source_positions
        max_output_lengths = (max_source_positions - 2) // 2 + 1

        return {"audio": max_output_lengths}
116

117
118
119
120

class Qwen2AudioDummyInputsBuilder(
        BaseDummyInputsBuilder[Qwen2AudioProcessingInfo]):

121
    def get_dummy_processor_inputs(
122
        self,
123
124
125
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> ProcessorInputs:
126
        feature_extractor = self.info.get_feature_extractor()
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

        sampling_rate = feature_extractor.sampling_rate
        audio_len = feature_extractor.chunk_length * sampling_rate
        num_audios = mm_counts.get("audio", 0)

        mm_data = {
            "audio":
            self._get_dummy_audios(length=audio_len, num_audios=num_audios)
        }

        return ProcessorInputs(
            prompt_text="<|AUDIO|>" * num_audios,
            mm_data=mm_data,
        )

142

143
144
class Qwen2AudioMultiModalProcessor(
        BaseMultiModalProcessor[Qwen2AudioProcessingInfo]):
145

146
    def _get_data_parser(self) -> MultiModalDataParser:
147
        feature_extractor = self.info.get_feature_extractor()
148
        return MultiModalDataParser(target_sr=feature_extractor.sampling_rate)
149

150
151
152
    def _call_hf_processor(
        self,
        prompt: str,
153
        mm_data: Mapping[str, object],
154
        mm_kwargs: Mapping[str, Any],
155
    ) -> BatchFeature:
156
157
        mm_data = dict(mm_data)
        audios = mm_data.pop("audios", [])
158

159
        if audios:
160
            mm_data["audios"] = audios
161

162
            feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
163
164
            mm_kwargs = dict(
                **mm_kwargs,
165
166
167
168
169
170
                sampling_rate=feature_extractor.sampling_rate,
            )
        else:
            # NOTE: WhisperFeatureExtractor cannot handle empty list of audios
            pass

171
        processed_outputs = super()._call_hf_processor(
172
            prompt=prompt,
173
174
175
176
177
178
179
180
181
182
183
184
185
186
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
        )

        return processed_outputs

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        return dict(
            input_features=MultiModalFieldConfig.batched("audio"),
            feature_attention_mask=MultiModalFieldConfig.batched("audio"),
187
188
189
190
191
        )

    def _get_prompt_replacements(
        self,
        mm_items: MultiModalDataItems,
192
193
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargs,
194
    ) -> list[PromptReplacement]:
195
        hf_config = self.info.get_hf_config()
196
197
        placeholder = hf_config.audio_token_index

198
        feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
199
200
201
        if feature_attention_mask is None:
            audio_output_lengths = []
        else:
202
            assert isinstance(feature_attention_mask, torch.Tensor)
203
            _, audio_output_lens = _get_feat_extract_output_lengths(
204
205
                feature_attention_mask.sum(-1))

206
207
            audio_output_lengths = audio_output_lens.tolist()

208
        def get_replacement_qwen2_audio(item_idx: int):
209
210
211
212
213
214
215
216
217
            num_placeholders = audio_output_lengths[item_idx]
            if num_placeholders == 0:
                audios = mm_items.get_items("audio", AudioProcessorItems)
                audio = audios.get(item_idx)
                raise ValueError(
                    f"The audio {audio} (len={len(audio)}) is too short "
                    "to be represented inside the model")

            return [placeholder] * num_placeholders
218
219
220
221
222
223
224

        return [
            PromptReplacement(
                modality="audio",
                target=[placeholder],
                replacement=get_replacement_qwen2_audio,
            )
225
        ]
226

227
    def _always_apply_prompt_replacements(self) -> bool:
228
229
230
        # Qwen2-Audio processor will start inserting placeholder tokens
        # in an upcoming release:
        # https://github.com/huggingface/transformers/pull/35534
231
232
233
234
        # NOTE: `_find_placeholders_by_modality` may incorrectly think that HF
        # has already performed processing for multi-audio input when the input
        # audios are short (the corresponding placeholders may take up fewer
        # tokens than the number of audio items)
235
        return not hasattr(self.info.get_hf_processor(), "audio_token")
236

237

238
239
240
241
@MULTIMODAL_REGISTRY.register_processor(
    Qwen2AudioMultiModalProcessor,
    info=Qwen2AudioProcessingInfo,
    dummy_inputs=Qwen2AudioDummyInputsBuilder)
242
243
244
class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,
                                         SupportsPP):

245
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
246
        super().__init__()
247
248
249
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
250
251
252
253
254
255
256
257
258
        self.config = config
        self.multimodal_config = multimodal_config

        self.audio_tower = Qwen2AudioEncoder(config.audio_config)
        self.multi_modal_projector = Qwen2AudioMultiModalProjector(
            config.audio_config.d_model, config.text_config.hidden_size)

        self.quant_config = quant_config

259
260
261
262
263
264
        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
            architectures=["Qwen2ForCausalLM"],
        )
265
266
267
268

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

269
270
271
272
273
274
275
    @cached_property
    def sampler(self):
        if hasattr(self.language_model, "sampler"):
            return self.language_model.sampler

        return get_sampler()

276
    def _validate_and_reshape_mm_tensor(self, mm_input: object,
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
                                        name: str) -> torch.Tensor:
        if not isinstance(mm_input, (torch.Tensor, list)):
            raise ValueError(f"Incorrect type of {name}. "
                             f"Got type: {type(mm_input)}")
        if isinstance(mm_input, torch.Tensor):
            return torch.concat(list(mm_input))
        else:
            return torch.concat(mm_input)

    def _parse_and_validate_audio_input(
            self, **kwargs: object) -> Optional[Qwen2AudioInputs]:
        input_features = kwargs.pop('input_features', None)
        feature_attention_mask = kwargs.pop('feature_attention_mask', None)
        if input_features is None:
            return None
        input_features = self._validate_and_reshape_mm_tensor(
            input_features, 'input_features')
        feature_attention_mask = self._validate_and_reshape_mm_tensor(
            feature_attention_mask, 'feature_attention_mask')
        if not isinstance(input_features, (torch.Tensor, list)):
            raise ValueError("Incorrect type of audio input features. "
                             f"Got type: {type(input_features)}")
        return Qwen2AudioInputs(input_features=input_features,
                                feature_attention_mask=feature_attention_mask)

    def _process_audio_input(self,
                             audio_input: Qwen2AudioInputs) -> torch.Tensor:

        input_features = audio_input["input_features"]
        feature_attention_mask = audio_input["feature_attention_mask"]

        audio_feat_lengths, audio_output_lengths = (
            self.audio_tower._get_feat_extract_output_lengths(
                feature_attention_mask.sum(-1)))

        batch_size, _, max_mel_seq_len = input_features.shape
        max_seq_len = (max_mel_seq_len - 2) // 2 + 1
        # Create a sequence tensor of shape (batch_size, max_seq_len)
        seq_range = (torch.arange(
            0,
            max_seq_len,
            dtype=audio_feat_lengths.dtype,
            device=audio_feat_lengths.device).unsqueeze(0).expand(
                batch_size, max_seq_len))
        lengths_expand = audio_feat_lengths.unsqueeze(-1).expand(
            batch_size, max_seq_len)
        # Create mask
        padding_mask = seq_range >= lengths_expand

        audio_attention_mask_ = padding_mask.view(
            batch_size, 1, 1, max_seq_len).expand(batch_size, 1, max_seq_len,
                                                  max_seq_len)
        audio_attention_mask = audio_attention_mask_.to(
            dtype=self.audio_tower.conv1.weight.dtype,
            device=self.audio_tower.conv1.weight.device)
        audio_attention_mask[audio_attention_mask_] = float("-inf")

        audio_outputs = self.audio_tower(input_features,
                                         attention_mask=audio_attention_mask)
        selected_audio_feature = audio_outputs.last_hidden_state
        audio_features = self.multi_modal_projector(selected_audio_feature)
        num_audios, max_audio_tokens, embed_dim = audio_features.shape
339
        audio_output_lengths = audio_output_lengths.unsqueeze(1)
340
        audio_features_mask = torch.arange(max_audio_tokens).expand(
341
342
            num_audios, max_audio_tokens).to(
                audio_output_lengths.device) < audio_output_lengths
343
344
345
        masked_audio_features = audio_features[audio_features_mask].view(
            -1, embed_dim)

346
347
348
        # Split to tuple of embeddings for individual audio input.
        return torch.split(masked_audio_features,
                           audio_output_lengths.flatten().tolist())
349

350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
    def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:
        audio_input = self._parse_and_validate_audio_input(**kwargs)
        if audio_input is None:
            return None
        masked_audio_features = self._process_audio_input(audio_input)
        return masked_audio_features

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: Optional[NestedTensors] = None,
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
        if multimodal_embeddings is not None:
            inputs_embeds = merge_multimodal_embeddings(
                input_ids, inputs_embeds, multimodal_embeddings,
                self.config.audio_token_index)
        return inputs_embeds

369
370
371
372
373
374
375
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
376
        inputs_embeds: Optional[torch.Tensor] = None,
377
378
        **kwargs: object,
    ) -> Union[torch.Tensor, IntermediateTensors]:
379

380
381
        if intermediate_tensors is not None:
            inputs_embeds = None
382
383
384
385
386
387
388
389
390

        # NOTE: In v1, inputs_embeds is always generated at model runner, this
        # condition is for v0 compatibility.
        elif inputs_embeds is None:
            multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
            inputs_embeds = self.get_input_embeddings(input_ids,
                                                      multimodal_embeddings)
            input_ids = None

391
392
393
394
395
396
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
                                                  intermediate_tensors,
                                                  inputs_embeds=inputs_embeds)
397
398
        return hidden_states

399
400
401
402
403
404
405
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
406
407
408
409
410
411

    def sample(
        self,
        logits: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
412
        return self.language_model.sample(logits, sampling_metadata)
413

414
415
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
416
417
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)