gemma3n_mm.py 29.4 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,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
185
186
187
188
    ) -> MultiModalDataDict:
        num_images = mm_counts.get("image", 0)
        num_audios = mm_counts.get("audio", 0)
        processor = self.info.get_hf_processor()
189
190
        audio_feature_extractor: Gemma3nAudioFeatureExtractor = (
            processor.feature_extractor
191
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
192
193
194
195
196
        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)

197
198
199
        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
200
        return {
201
202
203
204
205
206
207
208
209
            "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
210
211
212
        }


213
class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]):
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
214
215
216
217
218
219
220
221
    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
222
        # We pop `audios` and replace it with `audio` key to suppress
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
223
        # the warning.
224
225
        if "audios" in mm_data:
            mm_data["audio"] = mm_data.pop("audios")
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
226
227
228
229
230
231
        processed_outputs = super()._call_hf_processor(
            prompt,
            mm_data,
            mm_kwargs,
            tok_kwargs,
        )
232

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

            # Unpad features here since we need the output of each item to be
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
240
241
            # independent of other items for the cache to work correctly
            unpadded_features = [
242
243
                f[mask]
                for f, mask in zip(
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
244
245
246
247
248
249
250
251
252
253
254
255
                    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]:
256
257
258
        return dict(
            pixel_values=MultiModalFieldConfig.batched("image"),
            input_features_padded=MultiModalFieldConfig.batched("audio"),
259
260
            input_features_mask=MultiModalFieldConfig.batched("audio"),
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
261
262
263
264
265

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, Any],
266
        out_mm_kwargs: MultiModalKwargsItems,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
    ) -> 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,
290
291
                )
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
292
293
294
295
296
297

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

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

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

        return prompt_updates

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

        # "\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],
        )

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

    def _find_mm_placeholders(
        self,
        new_token_ids: list[int],
351
        mm_prompt_updates: MultiModalPromptUpdates,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
    ) -> 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)))

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

        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,
386
387
                )
                for p in placeholders
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
388
389
390
391
392
393
            ]
            for modality, placeholders in repls.items()
        }


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

    def __init__(
        self,
399
        multimodal_config: Gemma3nAudioConfig | Gemma3nVisionConfig,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
400
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
        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,
439
440
        input_ids: torch.LongTensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
441
442
443
444
445
446
447
448
449
450
451
452
453
    ) -> 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(
454
455
                "You must specify exactly one of input_ids or inputs_embeds"
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
456
457
458
459
460
461
462
463
464
465
466

        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)


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

Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
    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",
500
501
        }
    )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
502
503
504
505
506
507
508
509
510
511
512

    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

513
514
515
516
517
        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
518

519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
        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
542
543

    def _parse_and_validate_image_input(
544
        self, **kwargs: object
545
    ) -> Gemma3nImageInputs | None:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
546
547
548
549
550
551
552
        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

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

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

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

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

    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:
577
578
579
580
581
582
583
584
585
586
587
588
589
590
            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
591
592
593
594
595
596
597
        return mm_input_by_modality

    def _process_image_input(
        self,
        image_input: Gemma3nImageInputs,
    ) -> list[torch.Tensor]:
        pixel_values = image_input["pixel_values"]
598
599
600
        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
601
602
        # TODO try to avoid copy here
        # (batch, channels, height, width) to (batch, height * width, channels)
603
604
605
606
607
608
609
610
611
        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
612
613
614
615
616
617
618
619
620
        # 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]:
621
622
        # Run on padded features to enable batching
        input_features = audio_input["input_features_padded"].squeeze(1)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
623
        input_features_mask = audio_input["input_features_mask"].squeeze(1)
624
625
626
        audio_outputs, audio_mask = self.audio_tower(
            input_features, ~input_features_mask
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
627
628
        audio_features = self.embed_audio(inputs_embeds=audio_outputs)

629
630
631
632
633
634
635
636
637
        # 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
638
        # TODO precompute and cache padding
639
640
641
        audio_padding_toks = torch.tensor(
            [[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
642
        audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)
643
644
645
        audio_features = torch.where(
            audio_mask.unsqueeze(-1), audio_padding_embs, audio_features
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
646

647
648
649
        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
650
        )
651
652
653
654
655
656
657
        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
658
659
660
661

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

662
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
663
        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
        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

681
    def embed_input_ids(
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
682
683
        self,
        input_ids: torch.Tensor,
684
        multimodal_embeddings: MultiModalEmbeddings | None = None,
685
        *,
686
        is_multimodal: torch.Tensor | None = None,
687
        handle_oov_mm_token: bool = False,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
688
689
690
691
    ) -> 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:
692
            per_layer_inputs = self.language_model.model.get_per_layer_input_embeddings(
693
694
                input_ids
            )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
695
            per_layer_inputs = per_layer_inputs.reshape(
696
697
698
699
700
701
702
                -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
703

704
705
        # This is to satisfy the type checker for each overload
        if multimodal_embeddings is None or is_multimodal is None:
706
            return super().embed_input_ids(input_ids)
707

708
        return super().embed_input_ids(
709
710
711
712
713
            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
714

715
716
    def forward(
        self,
717
        input_ids: torch.Tensor | None,
718
        positions: torch.Tensor,
719
720
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
721
722
        **kwargs: object,
    ) -> IntermediateTensors:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
723
724
725
        if intermediate_tensors is not None:
            inputs_embeds = None

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

        hidden_states = self.language_model.model(
            input_ids,
            positions,
            per_layer_inputs=per_layer_inputs,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
739
740
            **kwargs,
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
741
742
743
744
745
746

        return hidden_states

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

750
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
751
752
753
754
755
756
757
758
759
760
        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",
761
762
            tower_model="vision_tower",
        )
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
763
764

    @classmethod
765
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
766
767
768
769
770
771
        if modality == "image":
            return "<image_soft_token>"
        elif modality == "audio":
            return "<audio_soft_token>"
        else:
            raise ValueError(f"Unsupported modality: {modality}")
772
773

    @classmethod
774
775
776
777
    def get_generation_prompt(
        cls,
        audio: np.ndarray,
        stt_config: SpeechToTextConfig,
778
        model_config: ModelConfig,
779
        language: str | None,
780
781
        task_type: Literal["transcribe", "translate"],
        request_prompt: str,
782
        to_language: str | None,
783
    ) -> PromptType:
784
785
        """
        Gemma3n supports "free-form" transcription.
786
        We fix its prompt here to standardize transcriptions/translations
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
        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"

810
811
812
813
        return TextPrompt(
            prompt=prompt,
            multi_modal_data={"audio": (audio, stt_config.sample_rate)},
        )
814
815

    @classmethod
816
    def get_speech_to_text_config(
817
        cls, model_config: ModelConfig, task_type: str
818
    ) -> SpeechToTextConfig:
819
820
821
822
823
824
825
826
        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,
        )