qwen2_audio.py 17.1 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,
                                        PromptReplacementDetails)
46
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
47
from vllm.sequence import IntermediateTensors
48
49

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


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

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


# === 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


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


84
class Qwen2AudioProcessingInfo(BaseProcessingInfo):
85

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

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

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

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

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

        return {"audio": max_output_lengths}
117

118
119
120
121

class Qwen2AudioDummyInputsBuilder(
        BaseDummyInputsBuilder[Qwen2AudioProcessingInfo]):

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

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

143

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

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

151
152
153
    def _call_hf_processor(
        self,
        prompt: str,
154
        mm_data: Mapping[str, object],
155
        mm_kwargs: Mapping[str, Any],
156
    ) -> BatchFeature:
157
158
159
160
161
162
163
164
165
166
167
        # Text-only input not supported in composite processor
        if not mm_data or not mm_data.get("audios", []):
            prompt_ids = self.info.get_tokenizer().encode(prompt)
            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

        feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
        mm_kwargs = dict(
            **mm_kwargs,
            sampling_rate=feature_extractor.sampling_rate,
        )
168

169
        return super()._call_hf_processor(
170
            prompt=prompt,
171
172
173
174
175
176
177
178
179
180
181
182
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
        )

    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"),
183
184
185
186
187
        )

    def _get_prompt_replacements(
        self,
        mm_items: MultiModalDataItems,
188
189
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargs,
190
    ) -> list[PromptReplacement]:
191
192
193
        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
        tokenizer = self.info.get_tokenizer()
        vocab = tokenizer.get_vocab()
194
195
196
197
198
199
200

        # Use getattr with default to be compatible with transformers<4.48
        audio_token = getattr(processor, "audio_token", "<|AUDIO|>")
        audio_bos_token = getattr(processor, "audio_bos_token",
                                  "<|audio_bos|>")
        audio_eos_token = getattr(processor, "audio_eos_token",
                                  "<|audio_eos|>")
201

202
203
204
205
        audio_token_id = vocab[audio_token]
        audio_bos_id = vocab[audio_bos_token]
        audio_eos_id = vocab[audio_eos_token]

206
        feature_attention_mask = out_mm_kwargs.get("feature_attention_mask")
207
208
209
        if feature_attention_mask is None:
            audio_output_lengths = []
        else:
210
            assert isinstance(feature_attention_mask, torch.Tensor)
211
            _, audio_output_lens = _get_feat_extract_output_lengths(
212
213
                feature_attention_mask.sum(-1))

214
215
            audio_output_lengths = audio_output_lens.tolist()

216
        def get_replacement_qwen2_audio(item_idx: int):
217
218
            num_features = audio_output_lengths[item_idx]
            if num_features == 0:
219
220
221
222
223
224
                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")

225
            audio_tokens = [audio_token_id] * num_features
226
227

            return PromptReplacementDetails(
228
                full=[audio_bos_id] + audio_tokens + [audio_eos_id],
229
230
                features=audio_tokens,
            )
231
232
233
234

        return [
            PromptReplacement(
                modality="audio",
235
                target=audio_token,
236
237
                replacement=get_replacement_qwen2_audio,
            )
238
        ]
239

240
    def _always_apply_prompt_replacements(self) -> bool:
241
242
243
        # Qwen2-Audio processor will start inserting placeholder tokens
        # in an upcoming release:
        # https://github.com/huggingface/transformers/pull/35534
244
245
246
247
        # 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)
248
        return not hasattr(self.info.get_hf_processor(), "audio_token")
249

250

251
252
253
254
@MULTIMODAL_REGISTRY.register_processor(
    Qwen2AudioMultiModalProcessor,
    info=Qwen2AudioProcessingInfo,
    dummy_inputs=Qwen2AudioDummyInputsBuilder)
255
256
257
class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,
                                         SupportsPP):

258
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
259
        super().__init__()
260
261
262
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
263
264
265
266
267
268
269
270
271
        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

272
273
274
275
276
277
        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
            architectures=["Qwen2ForCausalLM"],
        )
278
279
280
281

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

282
283
284
285
286
287
288
    @cached_property
    def sampler(self):
        if hasattr(self.language_model, "sampler"):
            return self.language_model.sampler

        return get_sampler()

289
    def _validate_and_reshape_mm_tensor(self, mm_input: object,
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
339
340
341
342
343
344
345
346
347
348
349
350
351
                                        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
352
        audio_output_lengths = audio_output_lengths.unsqueeze(1)
353
        audio_features_mask = torch.arange(max_audio_tokens).expand(
354
355
            num_audios, max_audio_tokens).to(
                audio_output_lengths.device) < audio_output_lengths
356
357
358
        masked_audio_features = audio_features[audio_features_mask].view(
            -1, embed_dim)

359
360
361
        # Split to tuple of embeddings for individual audio input.
        return torch.split(masked_audio_features,
                           audio_output_lengths.flatten().tolist())
362

363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
    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

382
383
384
385
386
387
388
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors] = None,
389
        inputs_embeds: Optional[torch.Tensor] = None,
390
391
        **kwargs: object,
    ) -> Union[torch.Tensor, IntermediateTensors]:
392

393
394
        if intermediate_tensors is not None:
            inputs_embeds = None
395
396
397
398
399
400
401
402
403

        # 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

404
405
406
407
408
409
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
                                                  intermediate_tensors,
                                                  inputs_embeds=inputs_embeds)
410
411
        return hidden_states

412
413
414
415
416
417
418
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
419
420
421
422
423
424

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

427
428
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
429
430
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)