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
from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataParser
42
from vllm.multimodal.processing import (BaseMultiModalProcessor,
43
                                        MultiModalDataItems, ProcessingMixin,
44
                                        PromptReplacement)
45
from vllm.multimodal.profiling import BaseProfilingInfo, 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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
class Qwen2AudioProcessingMixin(ProcessingMixin):

    def _get_hf_config(self):
        return self.ctx.get_hf_config(Qwen2AudioConfig)

    def _get_hf_processor(
        self,
        *,
        # Ignored in initialization
        sampling_rate: Optional[int] = None,
    ) -> Qwen2AudioProcessor:
        return self.ctx.get_hf_processor(Qwen2AudioProcessor)

    def _get_feature_extractor(
        self,
        *,
        # Ignored in initialization
        sampling_rate: Optional[int] = None,
    ) -> WhisperFeatureExtractor:
        hf_processor = self._get_hf_processor(sampling_rate=sampling_rate)
        feature_extractor = hf_processor.feature_extractor  # type: ignore
        assert isinstance(feature_extractor, WhisperFeatureExtractor)
        return feature_extractor


class Qwen2AudioProfilingInfo(Qwen2AudioProcessingMixin, BaseProfilingInfo):
109

110
111
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"audio": None}
112

113
    def get_mm_max_tokens_per_item(self, seq_len: int) -> Mapping[str, int]:
114
        hf_config = self._get_hf_config()
115
116
117
118
        max_source_positions = hf_config.audio_config.max_source_positions
        max_output_lengths = (max_source_positions - 2) // 2 + 1

        return {"audio": max_output_lengths}
119

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

        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,
        )

141

142
143
144
145
146
class Qwen2AudioMultiModalProcessor(Qwen2AudioProcessingMixin,
                                    BaseMultiModalProcessor):

    def _get_profiling_info(self) -> BaseProfilingInfo:
        return Qwen2AudioProfilingInfo(self.ctx)
147

148
    def _get_data_parser(self) -> MultiModalDataParser:
149
        feature_extractor = self._get_feature_extractor()
150
        return MultiModalDataParser(target_sr=feature_extractor.sampling_rate)
151

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

161
        if audios:
162
            mm_data["audios"] = audios
163

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

173
        processed_outputs = super()._call_hf_processor(
174
            prompt=prompt,
175
176
177
178
179
180
181
182
183
184
185
186
187
188
            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"),
189
190
191
192
193
        )

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

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

208
209
            audio_output_lengths = audio_output_lens.tolist()

210
        def get_replacement_qwen2_audio(item_idx: int):
211
212
213
214
215
216
217
218
219
            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
220
221
222
223
224
225
226

        return [
            PromptReplacement(
                modality="audio",
                target=[placeholder],
                replacement=get_replacement_qwen2_audio,
            )
227
        ]
228

229
    def _always_apply_prompt_replacements(self) -> bool:
230
231
232
        # Qwen2-Audio processor will start inserting placeholder tokens
        # in an upcoming release:
        # https://github.com/huggingface/transformers/pull/35534
233
234
235
236
        # 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)
237
        return not hasattr(self._get_hf_processor(), "audio_token")
238

239
240

@MULTIMODAL_REGISTRY.register_processor(Qwen2AudioMultiModalProcessor)
241
242
243
class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,
                                         SupportsPP):

244
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
245
        super().__init__()
246
247
248
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
249
250
251
252
253
254
255
256
257
        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

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

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

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

        return get_sampler()

275
    def _validate_and_reshape_mm_tensor(self, mm_input: object,
276
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
                                        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
338
        audio_output_lengths = audio_output_lengths.unsqueeze(1)
339
        audio_features_mask = torch.arange(max_audio_tokens).expand(
340
341
            num_audios, max_audio_tokens).to(
                audio_output_lengths.device) < audio_output_lengths
342
343
344
        masked_audio_features = audio_features[audio_features_mask].view(
            -1, embed_dim)

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

349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
    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

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

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

        # 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

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

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

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

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