gemma3n_mm.py 29.7 KB
Newer Older
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
1
2
3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Iterable, Mapping, Sequence
4
from typing import Annotated, Any, Literal
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
5

6
import numpy as np
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
7
8
9
import torch
from torch import nn
from transformers import AutoModel, BatchFeature
10
11
12
13
14
15
16
17
from transformers.models.gemma3n import (
    Gemma3nAudioConfig,
    Gemma3nAudioFeatureExtractor,
    Gemma3nConfig,
    Gemma3nProcessor,
    Gemma3nTextConfig,
    Gemma3nVisionConfig,
)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
18
19
from transformers.models.siglip import SiglipImageProcessorFast

20
from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
21
from vllm.config.multimodal import BaseDummyOptions
22
from vllm.inputs.data import PromptType, TextPrompt
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
23
24
25
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import RowParallelLinear
26
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
27
from vllm.model_executor.models.gemma3n import Gemma3nForCausalLM
28
29
30
from vllm.model_executor.models.gemma3n_audio_utils import (
    adjust_audio_features_to_expected_length,
)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
31
from vllm.model_executor.models.module_mapping import MultiModelKeys
32
from vllm.model_executor.models.whisper import ISO639_1_SUPPORTED_LANGS
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
33
from vllm.multimodal import MULTIMODAL_REGISTRY
34
35
36
37
38
39
40
41
42
43
from vllm.multimodal.inputs import (
    MultiModalDataDict,
    MultiModalFieldConfig,
    MultiModalKwargsItems,
)
from vllm.multimodal.parse import (
    ImageProcessorItems,
    MultiModalDataItems,
    MultiModalDataParser,
)
44
45
from vllm.multimodal.processing import BaseDummyInputsBuilder
from vllm.multimodal.processing.processor import (
46
47
48
49
50
51
52
53
54
55
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    MultiModalPromptUpdates,
    MultiModalPromptUpdatesApplyResult,
    PlaceholderFeaturesInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
    replace_token_matches,
)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
56
from vllm.sequence import IntermediateTensors
57
from vllm.utils.tensor_schema import TensorSchema, TensorShape
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
58

59
60
61
62
63
64
65
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsTranscription
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    init_vllm_registered_model,
    maybe_prefix,
)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
66
67
68
69
70
71
72
73

logger = init_logger(__name__)

# This should be based on model config but we hardcode them for now.
TOKENS_PER_IMAGE = 256
TOKENS_PER_AUDIO = 188


74
75
76
77
78
79
80
81
class Gemma3nImagePixelInputs(TensorSchema):
    """
    Dimensions:
        - bn: Batch size * number of images
        - c: Number of channels (3)
        - h: Height of each patch
        - w: Width of each patch
    """
82

83
84
    type: Literal["pixel_values"] = "pixel_values"
    pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
85
86


87
88
89
90
91
92
93
class Gemma3nAudioInputs(TensorSchema):
    """
    Dimensions:
        - bn: Batch size * number of audios
        - s: seq_length
        - f: num_features
    """
94

95
96
97
    type: Literal["audio"] = "audio"
    input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")]
    input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")]
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
98
99
100
101
102
103
104
105
106
107
108
109


Gemma3nImageInputs = Gemma3nImagePixelInputs


class Gemma3nProcessingInfo(BaseProcessingInfo):
    def get_hf_config(self):
        return self.ctx.get_hf_config(Gemma3nConfig)

    def get_hf_processor(self, **kwargs: object):
        return self.ctx.get_hf_processor(Gemma3nProcessor, **kwargs)

110
111
112
113
114
115
116
117
118
119
120
    def get_feature_extractor(self, **kwargs: object) -> Gemma3nAudioFeatureExtractor:
        return self.get_hf_processor(**kwargs).feature_extractor

    def get_data_parser(self):
        feature_extractor = self.get_feature_extractor()

        return MultiModalDataParser(
            target_sr=feature_extractor.sampling_rate,
            expected_hidden_size=self._get_expected_hidden_size(),
        )

121
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
122
123
124
        return {"image": None, "audio": None}

    def get_max_tokens_per_item(
125
        self, seq_len: int, mm_counts: Mapping[str, int]
126
    ) -> Mapping[str, int] | None:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
127
128
129
130
131
132
133
        return {"image": TOKENS_PER_IMAGE, "audio": TOKENS_PER_AUDIO}

    def get_image_repl(
        self,
        *,
        image_width: int,
        image_height: int,
134
        processor: Gemma3nProcessor | None,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
135
136
137
    ) -> str:
        """
        Get the replacement text for image tokens.
138

Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
139
140
141
142
143
144
145
        For Gemma3n, this should return the full_image_sequence which includes
        BOI token, repeated image tokens, and EOI token.
        """
        if processor is None:
            processor = self.get_hf_processor()

        return PromptUpdateDetails.select_token_id(
146
147
            processor.full_image_sequence, processor.image_token_id
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
148
149
150
151

    def get_audio_repl(
        self,
        *,
152
        processor: Gemma3nProcessor | None,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
153
154
155
    ) -> str:
        """
        Get the replacement text for audio tokens.
156

Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
157
158
159
160
161
162
163
164
        For Gemma3n, this should return the full_audio_sequence which includes
        BOA token, repeated audio tokens, and EOA token.
        """
        if processor is None:
            processor = self.get_hf_processor()

        # Return the full audio sequence as defined by the processor
        return PromptUpdateDetails.select_token_id(
165
166
            processor.full_audio_sequence, processor.audio_token_id
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183


class Gemma3nDummyInputsBuilder(BaseDummyInputsBuilder[Gemma3nProcessingInfo]):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_images = mm_counts.get("image", 0)
        num_audios = mm_counts.get("audio", 0)

        processor = self.info.get_hf_processor()
        image_token = processor.image_token
        audio_token = processor.audio_token

        return image_token * num_images + audio_token * num_audios

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
184
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
185
        mm_processor_kwargs: Mapping[str, object] | None = None,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
186
187
188
189
    ) -> MultiModalDataDict:
        num_images = mm_counts.get("image", 0)
        num_audios = mm_counts.get("audio", 0)
        processor = self.info.get_hf_processor()
190
191
        audio_feature_extractor: Gemma3nAudioFeatureExtractor = (
            processor.feature_extractor
192
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
193
194
195
196
197
        audio_len = audio_feature_extractor.fft_length
        image_processor: SiglipImageProcessorFast = processor.image_processor
        img_width = image_processor.size.get("width", 224)
        img_height = image_processor.size.get("height", 224)

198
199
200
        image_overrides = mm_options.get("image") if mm_options else None
        audio_overrides = mm_options.get("audio") if mm_options else None

Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
201
        return {
202
203
204
205
206
207
208
209
210
            "image": self._get_dummy_images(
                width=img_width,
                height=img_height,
                num_images=num_images,
                overrides=image_overrides,
            ),
            "audio": self._get_dummy_audios(
                length=audio_len, num_audios=num_audios, overrides=audio_overrides
            ),
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
211
212
213
        }


214
class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]):
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
215
216
217
218
219
220
221
222
    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        # HF Transformers audio processor no longer accepts `audios` key.
co63oc's avatar
co63oc committed
223
        # We pop `audios` and replace it with `audio` key to suppress
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
224
        # the warning.
225
226
        if "audios" in mm_data:
            mm_data["audio"] = mm_data.pop("audios")
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
227
228
229
230
231
232
        processed_outputs = super()._call_hf_processor(
            prompt,
            mm_data,
            mm_kwargs,
            tok_kwargs,
        )
233

234
        if "input_features" in processed_outputs:
235
            # Padding enables audio_tower to run in batched mode
236
237
238
            processed_outputs["input_features_padded"] = processed_outputs[
                "input_features"
            ]
239
240

            # Unpad features here since we need the output of each item to be
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
241
242
            # independent of other items for the cache to work correctly
            unpadded_features = [
243
244
                f[mask]
                for f, mask in zip(
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
245
246
247
248
249
250
251
252
253
254
255
256
                    processed_outputs["input_features"],
                    processed_outputs["input_features_mask"],
                )
            ]
            processed_outputs["input_features"] = unpadded_features
        return processed_outputs

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
257
258
259
        return dict(
            pixel_values=MultiModalFieldConfig.batched("image"),
            input_features_padded=MultiModalFieldConfig.batched("audio"),
260
261
            input_features_mask=MultiModalFieldConfig.batched("audio"),
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
262
263
264
265
266

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, Any],
267
        out_mm_kwargs: MultiModalKwargsItems,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
    ) -> Sequence[PromptUpdate]:
        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

        prompt_updates = []

        # Handle image tokens
        if "image" in mm_items:
            image_token = hf_processor.image_token

            def get_replacement_image(item_idx: int):
                images = mm_items.get_items("image", ImageProcessorItems)
                image_size = images.get_image_size(item_idx)
                return self.info.get_image_repl(
                    image_width=image_size.width,
                    image_height=image_size.height,
                    processor=hf_processor,
                )

            prompt_updates.append(
                PromptReplacement(
                    modality="image",
                    target=image_token,
                    replacement=get_replacement_image,
291
292
                )
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
293
294
295
296
297
298

        # Handle audio tokens
        if "audio" in mm_items:
            audio_token = hf_processor.audio_token

            def get_replacement_audio(item_idx: int):
299
300
301
                return self.info.get_audio_repl(
                    processor=hf_processor,
                )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
302
303
304
305
306
307

            prompt_updates.append(
                PromptReplacement(
                    modality="audio",
                    target=audio_token,
                    replacement=get_replacement_audio,
308
309
                )
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
310
311
312
313
314
315

        return prompt_updates

    def _apply_token_matches(
        self,
        prompt: list[int],
316
317
        mm_prompt_updates: MultiModalPromptUpdates,
    ) -> tuple[list[int], MultiModalPromptUpdatesApplyResult]:
318
        token_ids, res = super()._apply_token_matches(prompt, mm_prompt_updates)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
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

        # "\n\n\n" and "\n\n\n\n" are single tokens
        # Since our replacement can insert "\n\n" next to "\n"
        # tokens, we have to combine them to be consistent with
        # the output of the tokenizer
        tokenizer = self.info.get_tokenizer()
        vocab = tokenizer.get_vocab()
        newline_1 = vocab["\n"]
        newline_2 = vocab["\n\n"]
        newline_3 = vocab["\n\n\n"]
        newline_4 = vocab["\n\n\n\n"]

        token_ids = replace_token_matches(
            token_ids,
            [newline_1, newline_2],
            [newline_3],
        )
        token_ids = replace_token_matches(
            token_ids,
            [newline_2, newline_1],
            [newline_3],
        )
        token_ids = replace_token_matches(
            token_ids,
            [newline_2, newline_2],
            [newline_4],
        )

347
        return token_ids, res
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
348
349
350
351

    def _find_mm_placeholders(
        self,
        new_token_ids: list[int],
352
        mm_prompt_updates: MultiModalPromptUpdates,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
    ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
        # We need to detect "\n\n" inside "\n\n\n" and "\n\n\n\n"
        tokenizer = self.info.get_tokenizer()
        vocab = tokenizer.get_vocab()
        newline_1 = vocab["\n"]
        newline_2 = vocab["\n\n"]
        newline_3 = vocab["\n\n\n"]
        newline_4 = vocab["\n\n\n\n"]

        def get_repl_toks(tok: int) -> list[int]:
            if tok == newline_3:
                return [newline_1, newline_2]
            if tok == newline_4:
                return [newline_2, newline_2]

            return [tok]

        repl_token_ids = list[int]()
        repl_orig_idxs = list[int]()
        for orig_idx, orig_tok in enumerate(new_token_ids):
            repl_toks = get_repl_toks(orig_tok)
            repl_token_ids.extend(repl_toks)
            repl_orig_idxs.extend(orig_idx for _ in range(len(repl_toks)))

377
        repls = super()._find_mm_placeholders(repl_token_ids, mm_prompt_updates)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
378
379
380
381
382
383
384
385
386

        return {
            modality: [
                PlaceholderFeaturesInfo(
                    modality=p.modality,
                    item_idx=p.item_idx,
                    start_idx=repl_orig_idxs[p.start_idx],
                    tokens=p.tokens,
                    is_embed=p.is_embed,
387
388
                )
                for p in placeholders
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
389
390
391
392
393
394
            ]
            for modality, placeholders in repls.items()
        }


class Gemma3nMultimodalEmbedder(nn.Module):
395
    """Embeds token ids or soft tokens for multimodal content into language
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
396
397
398
399
    model space."""

    def __init__(
        self,
400
        multimodal_config: Gemma3nAudioConfig | Gemma3nVisionConfig,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
        text_config: Gemma3nTextConfig,
    ):
        super().__init__()

        self.multimodal_hidden_size = multimodal_config.hidden_size
        self.eps = multimodal_config.rms_norm_eps
        self.vocab_offset = multimodal_config.vocab_offset
        self.vocab_size = multimodal_config.vocab_size
        self.text_hidden_size = text_config.hidden_size

        self.embedding = VocabParallelEmbedding(
            self.vocab_size,
            self.multimodal_hidden_size,
        )

        self.hard_embedding_norm = RMSNorm(
            self.multimodal_hidden_size,
            eps=self.eps,
        )

        self.soft_embedding_norm = RMSNorm(
            self.multimodal_hidden_size,
            eps=self.eps,
        )

        self.embedding_projection = RowParallelLinear(
            self.multimodal_hidden_size,
            self.text_hidden_size,
            bias=False,
        )

        self.embedding_post_projection_norm = RMSNorm(
            self.text_hidden_size,
            eps=self.eps,
            has_weight=False,
        )

    def forward(
        self,
440
441
        input_ids: torch.LongTensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
442
443
444
445
446
447
448
449
450
451
452
453
454
    ) -> torch.Tensor:
        """Embeds token ids or soft tokens for multimodal content into language model space.

        Args:
            input_ids: A torch.LongTensor containing the token ids to embed. Values should be in the range
                `[vocab_offset, vocab_offset + vocab_size)`.
            inputs_embeds: A torch.Tensor containing the soft tokens to embed.

        Returns:
            A torch.Tensor of embeddings with  shape `[batch_size, seq_len, self.config.text_config.hidden_size]`.
        """  # noqa: E501
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
455
456
                "You must specify exactly one of input_ids or inputs_embeds"
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
457
458
459
460
461
462
463
464
465
466
467

        if inputs_embeds is not None:
            emb_norm = self.soft_embedding_norm(inputs_embeds)
        else:
            hard_emb = self.embedding(input_ids - self.vocab_offset)
            emb_norm = self.hard_embedding_norm(hard_emb)

        emb_norm_proj, _ = self.embedding_projection(emb_norm)
        return self.embedding_post_projection_norm(emb_norm_proj)


468
469
470
471
472
473
474
475
@MULTIMODAL_REGISTRY.register_processor(
    Gemma3nMultiModalProcessor,
    info=Gemma3nProcessingInfo,
    dummy_inputs=Gemma3nDummyInputsBuilder,
)
class Gemma3nForConditionalGeneration(
    nn.Module, SupportsMultiModal, SupportsTranscription
):
476
477
    supported_languages = ISO639_1_SUPPORTED_LANGS

Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            # mapping for new names in checkpoint saved after transformers v4.52
            "model.embed_audio.": "embed_audio.",
            "model.embed_vision.": "embed_vision.",
            "model.language_model.": "language_model.model.",
            "model.vision_tower.": "vision_tower.",
            "model.audio_tower.": "audio_tower.",
            "model.multi_modal_projector.": "multi_modal_projector.",
            "lm_head.": "language_model.lm_head.",
            "model": "language_model.model",
501
502
        }
    )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
503
504
505
506
507
508
509
510
511
512
513

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
        self.config = config
        self.quant_config = quant_config
        self.multimodal_config = multimodal_config
        self.vocab_size = config.text_config.vocab_size

514
515
516
517
518
        with self._mark_tower_model(vllm_config, "image"):
            self.vision_tower = AutoModel.from_config(config=config.vision_config)
            self.embed_vision = Gemma3nMultimodalEmbedder(
                config.vision_config, config.text_config
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
519

520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
        with self._mark_tower_model(vllm_config, "audio"):
            self.audio_tower = AutoModel.from_config(config=config.audio_config)
            self.embed_audio = Gemma3nMultimodalEmbedder(
                config.audio_config, config.text_config
            )

        with self._mark_language_model(vllm_config):
            self.language_model: Gemma3nForCausalLM = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.text_config,
                prefix=maybe_prefix(prefix, "language_model"),
                architectures=["Gemma3nForCausalLM"],
            )

            # NOTE (NickLucche) In order to be compatible with cudagraph, the
            # buffer needs to be consistent, so we pre-allocate here.
            self.per_layer_embeddings = torch.zeros(
                vllm_config.scheduler_config.max_num_batched_tokens,
                self.config.text_config.num_hidden_layers,
                self.config.text_config.hidden_size_per_layer_input,
                device=self.language_model.model.embed_tokens.weight.device,
                dtype=self.language_model.model.embed_tokens.weight.dtype,
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
543
544

    def _parse_and_validate_image_input(
545
        self, **kwargs: object
546
    ) -> Gemma3nImageInputs | None:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
547
548
549
550
551
552
553
        pixel_values = kwargs.pop("pixel_values", None)
        image_embeds = kwargs.pop("image_embeds", None)
        # TODO is this the case?
        assert image_embeds is None, "Gemma3n does not support image_embeds."
        if pixel_values is None:
            return None

554
        return Gemma3nImagePixelInputs(pixel_values=pixel_values)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
555
556

    def _parse_and_validate_audio_input(
557
        self, **kwargs: object
558
    ) -> Gemma3nAudioInputs | None:
559
560
        input_features_padded = kwargs.pop("input_features_padded", None)
        if input_features_padded is None:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
561
562
563
564
565
566
567
            return None

        input_features_mask = kwargs.pop("input_features_mask", None)
        if input_features_mask is None:
            return None

        return Gemma3nAudioInputs(
568
            input_features_padded=input_features_padded,
569
            input_features_mask=input_features_mask,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
570
571
572
573
574
575
576
577
        )

    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
        mm_input_by_modality = {}

        # Preserve the order of modalities if there are multiple of them
        # from the order of kwargs.
        for input_key in kwargs:
578
579
580
581
582
583
584
585
586
587
588
589
590
591
            if (
                input_key in ("pixel_values", "image_embeds")
                and "image" not in mm_input_by_modality
            ):
                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
                    **kwargs
                )
            if (
                input_key == "input_features_padded"
                and "audio" not in mm_input_by_modality
            ):
                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
                    **kwargs
                )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
592
593
594
595
596
597
598
        return mm_input_by_modality

    def _process_image_input(
        self,
        image_input: Gemma3nImageInputs,
    ) -> list[torch.Tensor]:
        pixel_values = image_input["pixel_values"]
599
600
601
        vision_outputs = self.vision_tower(
            pixel_values=pixel_values, do_pooling=False, return_dict=True
        ).last_hidden_state
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
602
603
        # TODO try to avoid copy here
        # (batch, channels, height, width) to (batch, height * width, channels)
604
605
606
607
608
609
610
611
612
        vision_outputs = (
            vision_outputs.reshape(
                vision_outputs.shape[0],
                self.config.vision_config.hidden_size,
                self.config.vision_soft_tokens_per_image,
            )
            .permute(0, 2, 1)
            .contiguous()
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
613
614
615
616
617
618
619
620
621
        # Normalize and embed the soft tokens into language model space.
        vision_outputs *= self.config.vision_config.hidden_size**0.5
        # Return a list of embeddings instead of a batched tensor
        return self.embed_vision(inputs_embeds=vision_outputs).unbind(0)

    def _process_audio_input(
        self,
        audio_input: Gemma3nAudioInputs,
    ) -> list[torch.Tensor]:
622
623
        # Run on padded features to enable batching
        input_features = audio_input["input_features_padded"].squeeze(1)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
624
        input_features_mask = audio_input["input_features_mask"].squeeze(1)
625
626
627
628
629
630
631
632
633
        audio_outputs = self.audio_tower(input_features, ~input_features_mask)
        if isinstance(audio_outputs, tuple):
            # Transformers v4
            audio_encodings, audio_mask = audio_outputs
        else:
            # Transformers v5
            audio_encodings = audio_outputs.last_hidden_state
            audio_mask = audio_outputs.audio_mel_mask
        audio_features = self.embed_audio(inputs_embeds=audio_encodings)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
634

635
636
637
638
639
640
641
642
643
        # The Gemma3nProcessor expects all audio will be 30s in length and
        # inserts 188 audio soft tokens into the text to account for this.
        # However, the audio preprocessing and encoder do not guarantee they
        # will produce exactly 188 soft tokens; they may produce fewer tokens
        # (for shorter audio) or more tokens (for longer audio or due to
        # BOA/EOA special tokens in the placeholder sequence).
        # We handle both cases:
        # - If fewer tokens: pad with the embedding of the last vocab token
        # - If more tokens: truncate to the expected count
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
644
        # TODO precompute and cache padding
645
646
647
        audio_padding_toks = torch.tensor(
            [[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
648
        audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)
649
650
651
        audio_features = torch.where(
            audio_mask.unsqueeze(-1), audio_padding_embs, audio_features
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
652

653
654
655
        expected_tokens = self.config.audio_soft_tokens_per_image
        audio_features, tokens_truncated = adjust_audio_features_to_expected_length(
            audio_features, expected_tokens, audio_padding_embs
656
        )
657
658
659
660
661
662
663
        if tokens_truncated > 0:
            logger.warning(
                "Gemma3n audio encoder produced %d extra tokens. "
                "Truncating to match placeholder count of %d.",
                tokens_truncated,
                expected_tokens,
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
664
665
666
667

        # Return a list of embeddings instead of a batched tensor
        return audio_features.unbind(0)

668
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
669
        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
        if mm_input_by_modality is None:
            return []

        multimodal_embeddings: list[torch.Tensor] = []

        # NOTE: It is important to iterate over the keys in this dictionary
        # to preserve the order of the modalities.
        for modality in mm_input_by_modality:
            multimodal_input = mm_input_by_modality[modality]
            if modality == "image":
                vision_embeddings = self._process_image_input(multimodal_input)
                multimodal_embeddings.extend(vision_embeddings)
            if modality == "audio":
                audio_embeddings = self._process_audio_input(multimodal_input)
                multimodal_embeddings.extend(audio_embeddings)
        return multimodal_embeddings

687
    def embed_input_ids(
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
688
689
        self,
        input_ids: torch.Tensor,
690
        multimodal_embeddings: MultiModalEmbeddings | None = None,
691
        *,
692
        is_multimodal: torch.Tensor | None = None,
693
        handle_oov_mm_token: bool = False,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
694
695
696
697
    ) -> torch.Tensor:
        # NOTE (NickLucche) Each pass needs tokens to compute PLE so we cache
        # them here, as the model  forward has only access to the input_embeds.
        if input_ids is not None:
698
            per_layer_inputs = self.language_model.model.get_per_layer_input_embeddings(
699
700
                input_ids
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
701
            per_layer_inputs = per_layer_inputs.reshape(
702
703
704
705
706
707
708
                -1,
                self.config.text_config.num_hidden_layers,
                self.config.text_config.hidden_size_per_layer_input,
            )
            self.per_layer_embeddings[: per_layer_inputs.shape[0]].copy_(
                per_layer_inputs
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
709

710
711
        # This is to satisfy the type checker for each overload
        if multimodal_embeddings is None or is_multimodal is None:
712
            return super().embed_input_ids(input_ids)
713

714
        return super().embed_input_ids(
715
716
717
718
719
            input_ids,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
            handle_oov_mm_token=handle_oov_mm_token,
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
720

721
722
    def forward(
        self,
723
        input_ids: torch.Tensor | None,
724
        positions: torch.Tensor,
725
726
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
727
728
        **kwargs: object,
    ) -> IntermediateTensors:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
729
730
731
        if intermediate_tensors is not None:
            inputs_embeds = None

732
        # NOTE (NickLucche) During profiling, `embed_input_ids` is not
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
733
734
        # called, hence we don't have input_ids to compute PLEs. We simply
        # select a chunk of pre-allocated PLEs. During normal execution,
735
        # `embed_input_ids` is called before forward, hence this slice
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
736
        # will contain PLEs computed from the actual input_ids.
737
        per_layer_inputs = self.per_layer_embeddings[: inputs_embeds.shape[0]]
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
738
739
740
741
742
743
744

        hidden_states = self.language_model.model(
            input_ids,
            positions,
            per_layer_inputs=per_layer_inputs,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
745
746
            **kwargs,
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
747
748
749
750
751
752

        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
753
    ) -> torch.Tensor | None:
754
        return self.language_model.compute_logits(hidden_states)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
755

756
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
757
758
759
760
761
762
763
764
765
766
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    def get_mm_mapping(self) -> MultiModelKeys:
        """
        Get the module prefix in multimodal models
        """
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector="multi_modal_projector",
767
768
            tower_model="vision_tower",
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
769
770

    @classmethod
771
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
772
773
774
775
776
777
        if modality == "image":
            return "<image_soft_token>"
        elif modality == "audio":
            return "<audio_soft_token>"
        else:
            raise ValueError(f"Unsupported modality: {modality}")
778
779

    @classmethod
780
781
782
783
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
784
        model_config: ModelConfig,
785
        language: str | None,
786
787
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
788
        to_language: str | None,
789
    ) -> PromptType:
790
791
        """
        Gemma3n supports "free-form" transcription.
792
        We fix its prompt here to standardize transcriptions/translations
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
        requests.
        """
        # Transcribe this audio [into <>] | for transcription
        # Translate this audio [from <> into <>] | for translation
        prompt = "<start_of_turn>user\n"
        prompt += "Transcribe" if task_type == "transcribe" else "Translate"
        prompt += " this audio"

        # We assume the language is a valid ISO 639-1 code.
        full_lang_name = cls.supported_languages.get(language, "")
        # Translation only for now
        full_lang_name_to = cls.supported_languages.get(to_language, "")

        if task_type == "transcribe" and full_lang_name:
            prompt += f" into {full_lang_name}"
        elif task_type == "translate":
            if full_lang_name:
                prompt += f" from {full_lang_name}"
            if full_lang_name_to:
                prompt += f" into {full_lang_name_to}"

        prompt += ": <audio_soft_token><end_of_turn>\n<start_of_turn>model\n"

816
817
818
819
        return TextPrompt(
            prompt=prompt,
            multi_modal_data={"audio": (audio, stt_config.sample_rate)},
        )
820
821

    @classmethod
822
    def get_speech_to_text_config(
823
        cls, model_config: ModelConfig, task_type: str
824
    ) -> SpeechToTextConfig:
825
826
827
828
829
830
831
832
        return SpeechToTextConfig(
            # Let's set this to 30 as suggested in the docs for now, although
            # the model is only limited by its context length.
            max_audio_clip_s=30,
            sample_rate=16000,
            # TODO enable chunking after more thorough testing.
            min_energy_split_window_size=None,
        )