Aluode/PerceptionLabPortable
0
1# Copyright 2025 Microsoft and the HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import math16from typing import Callable, Optional, Union17 18import numpy as np19import torch20import torch.nn.functional as F21from torch import nn22 23from ...activations import ACT2FN24from ...cache_utils import Cache, DynamicCache25from ...configuration_utils import PretrainedConfig26from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask27from ...modeling_attn_mask_utils import _prepare_4d_attention_mask28from ...modeling_outputs import (29 BaseModelOutput,30 BaseModelOutputWithPast,31 BaseModelOutputWithPooling,32 CausalLMOutputWithPast,33)34from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel35from ...processing_utils import Unpack36from ...utils import auto_docstring, logging37from ...utils.generic import TransformersKwargs, check_model_inputs38from ..phi3.configuration_phi3 import Phi3Config39from ..phi3.modeling_phi3 import (40 Phi3DecoderLayer,41 Phi3ForCausalLM,42 Phi3Model,43 Phi3PreTrainedModel,44 Phi3RMSNorm,45 Phi3RotaryEmbedding,46)47from ..siglip.configuration_siglip import SiglipVisionConfig48from ..siglip.modeling_siglip import (49 SiglipEncoder,50 SiglipEncoderLayer,51 SiglipMLP,52 SiglipMultiheadAttentionPoolingHead,53 SiglipPreTrainedModel,54 SiglipVisionEmbeddings,55 default_flax_embed_init,56 lecun_normal_,57)58 59 60logger = logging.get_logger(__name__)61 62 63class Phi4MultimodalVisionConfig(SiglipVisionConfig):64 r"""65 This is the configuration class to store the configuration of a [`Phi4MultimodalVisionModel`]. It is used to instantiate a66 Phi4Multimodal vision encoder according to the specified arguments, defining the model architecture. Instantiating a67 configuration with the defaults will yield a similar configuration to that of the vision encoder of68 [microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct) architecture.69 70 Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the71 documentation from [`PretrainedConfig`] for more information.72 73 Args:74 hidden_size (`int`, *optional*, defaults to 1152):75 Dimensionality of the encoder layers and the pooler layer.76 intermediate_size (`int`, *optional*, defaults to 4304):77 Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.78 num_hidden_layers (`int`, *optional*, defaults to 27):79 Number of hidden layers in the Transformer encoder.80 num_attention_heads (`int`, *optional*, defaults to 16):81 Number of attention heads for each attention layer in the Transformer encoder.82 num_channels (`int`, *optional*, defaults to 3):83 Number of channels in the input images.84 image_size (`int`, *optional*, defaults to 448):85 The size (resolution) of each image.86 patch_size (`int`, *optional*, defaults to 14):87 The size (resolution) of each patch.88 hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):89 The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,90 `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.91 layer_norm_eps (`float`, *optional*, defaults to 1e-06):92 The epsilon used by the layer normalization layers.93 attention_dropout (`float`, *optional*, defaults to 0.0):94 The dropout ratio for the attention probabilities.95 crop_size (`int`, *optional*, defaults to 448):96 Crop size for the input images.97 image_token_id (`int`, *optional*, defaults to 200010):98 The image token id.99 feature_layer (`int`, *optional*, defaults to -2):100 The index of the layer of the encoder from which to extract image features.101 102 Example:103 104 ```python105 >>> from transformers import Phi4MultimodalVisionConfig106 107 >>> # Initializing a Phi4MultimodalVisionConfig with microsoft/Phi-4-multimodal-instruct style configuration108 >>> configuration = Phi4MultimodalVisionConfig()109 ```"""110 111 model_type = "phi4_multimodal_vision"112 113 def __init__(114 self,115 hidden_size=1152,116 intermediate_size=4304,117 num_hidden_layers=27,118 num_attention_heads=16,119 num_channels=3,120 image_size=448,121 patch_size=14,122 hidden_act="gelu_pytorch_tanh",123 layer_norm_eps=1e-6,124 attention_dropout=0.0,125 crop_size: int = 448,126 image_token_id: int = 200010,127 feature_layer: int = -2,128 **kwargs,129 ):130 super().__init__(131 hidden_size=hidden_size,132 intermediate_size=intermediate_size,133 num_hidden_layers=num_hidden_layers,134 num_attention_heads=num_attention_heads,135 num_channels=num_channels,136 image_size=image_size,137 patch_size=patch_size,138 hidden_act=hidden_act,139 layer_norm_eps=layer_norm_eps,140 attention_dropout=attention_dropout,141 **kwargs,142 )143 self.crop_size = crop_size144 self.image_token_id = image_token_id145 self.feature_layer = feature_layer146 147 148class Phi4MultimodalAudioConfig(PretrainedConfig):149 r"""150 This is the configuration class to store the configuration of a [`Phi4MultimodalAudioModel`]. It is used to instantiate a151 Phi4Multimodal audio encoder according to the specified arguments, defining the model architecture. Instantiating a152 configuration with the defaults will yield a similar configuration to that of the audio encoder of153 [microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct) architecture.154 155 Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the156 documentation from [`PretrainedConfig`] for more information.157 158 Args:159 hidden_size (`int`, *optional*, defaults to 1024):160 Dimensionality of the encoder layers.161 intermediate_size (`int`, *optional*, defaults to 1536):162 Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.163 num_blocks (`int`, *optional*, defaults to 24):164 Number of hidden layers in the Transformer encoder.165 num_attention_heads (`int`, *optional*, defaults to 16):166 Number of attention heads for each attention layer in the Transformer encoder.167 activation (`str`, *optional*, defaults to `"swish"`):168 The non-linear activation function in the MLPs.169 chunk_size (`int`, *optional*, defaults to -1):170 The chunk size to create the masks.171 left_chunk (`int`, *optional*, defaults to 18):172 The left chunk to create the masks.173 dropout_rate (`float`, *optional*, defaults to 0.0):174 The dropout ratio.175 ext_pw_out_channel (`int`, *optional*, defaults to 1024):176 Number of out channels in the point-wise conv modules.177 depthwise_separable_out_channel (`int`, *optional*, defaults to 1024):178 Number of out channels in the depth-wise separable conv modules.179 depthwise_multiplier (`int`, *optional*, defaults to 1):180 Input size multiplier for the depth-wise separable conv modules.181 kernel_size (`int`, *optional*, defaults to 3):182 Kernel size for the depth-wise separable conv modules.183 conv_activation (`str`, *optional*, defaults to `"swish"`):184 The non-linear activation function in the conv modules.185 input_size (`int`, *optional*, defaults to 80):186 Input size for the audio model.187 conv_glu_type (`str`, *optional*, defaults to `"swish"`):188 The non-linear activation function in the point-wise conv modules.189 time_reduction (`int`, *optional*, defaults to 8):190 Time reduction (subsampling factor).191 bias_max_distance (`int`, *optional*, defaults to 1000):192 Max distance for the relative attention bias module.193 bias_symmetric (`bool`, *optional*, defaults to `False`):194 Whether the relative attention bias should be symmetric or not.195 nemo_activation (`str`, *optional*, defaults to `"relu"`):196 The non-linear activation function in the nemo conv modules.197 nemo_conv_channels (`int`, *optional*, defaults to 1024):198 Number of channels in the nemo conv modules.199 downsample_rate (`int`, *optional*, defaults to 1):200 Downsample rate for the audio feature extractor.201 initializer_range (`float`, *optional*, defaults to 0.02):202 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.203 audio_token_id (`int`, *optional*, defaults to 200011):204 The audio token id.205 feature_layer (`int`, *optional*, defaults to -2):206 The index of the layer of the encoder from which to extract audio features.207 208 Example:209 210 ```python211 >>> from transformers import Phi4MultimodalAudioConfig212 213 >>> # Initializing a Phi4MultimodalAudioConfig with microsoft/Phi-4-multimodal-instruct style configuration214 >>> configuration = Phi4MultimodalAudioConfig()215 ```"""216 217 model_type = "phi4_multimodal_audio"218 219 def __init__(220 self,221 hidden_size: int = 1024,222 intermediate_size: int = 1536,223 num_blocks: int = 24,224 num_attention_heads: int = 16,225 activation: str = "swish",226 chunk_size: int = -1,227 left_chunk: int = 18,228 dropout_rate: float = 0.0,229 ext_pw_out_channel: int = 1024,230 depthwise_separable_out_channel: int = 1024,231 depthwise_multiplier: int = 1,232 kernel_size: int = 3,233 conv_activation: str = "swish",234 input_size: int = 80,235 conv_glu_type: str = "swish",236 time_reduction: int = 8,237 bias_max_distance: int = 1000,238 bias_symmetric: bool = False,239 nemo_activation: str = "relu",240 nemo_conv_channels: int = 1024,241 downsample_rate: int = 1,242 initializer_range: float = 0.02,243 audio_token_id: int = 200011,244 feature_layer: int = -2,245 **kwargs,246 ):247 super().__init__(**kwargs)248 self.hidden_size = hidden_size249 self.num_attention_heads = num_attention_heads250 self.intermediate_size = intermediate_size251 self.activation = activation252 self.chunk_size = chunk_size253 self.left_chunk = left_chunk254 self.num_blocks = num_blocks255 self.dropout_rate = dropout_rate256 self.ext_pw_out_channel = ext_pw_out_channel257 self.depthwise_separable_out_channel = depthwise_separable_out_channel258 self.depthwise_multiplier = depthwise_multiplier259 self.kernel_size = kernel_size260 self.conv_activation = conv_activation261 self.input_size = input_size262 self.conv_glu_type = conv_glu_type263 self.time_reduction = time_reduction264 self.bias_max_distance = bias_max_distance265 self.bias_symmetric = bias_symmetric266 self.nemo_activation = nemo_activation267 self.nemo_conv_channels = nemo_conv_channels268 self.downsample_rate = downsample_rate269 self.audio_token_id = audio_token_id270 self.initializer_range = initializer_range271 self.feature_layer = feature_layer272 273 if time_reduction % 2 != 0:274 raise ValueError("`time_reduction` should be a multiple of 2!")275 length = input_size276 for _ in range(int(math.log2(time_reduction))):277 length = math.floor((length - 1) / 2 + 1)278 self.nemo_final_size = length279 280 281class Phi4MultimodalConfig(Phi3Config):282 r"""283 This is the configuration class to store the configuration of a [`Phi4MultimodalModel`]. It is used to instantiate a284 Phi4Multimodal model according to the specified arguments, defining the model architecture. Instantiating a configuration285 with the defaults will yield a similar configuration to that of the286 [microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct) architecture.287 288 Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the289 documentation from [`PretrainedConfig`] for more information.290 291 Args:292 vocab_size (`int`, *optional*, defaults to 200064):293 Vocabulary size of the Phi-3 model. Defines the number of different tokens that can be represented by the294 `inputs_ids` passed when calling [`Phi3Model`].295 hidden_size (`int`, *optional*, defaults to 3072):296 Dimension of the hidden representations.297 intermediate_size (`int`, *optional*, defaults to 8192):298 Dimension of the MLP representations.299 num_hidden_layers (`int`, *optional*, defaults to 32):300 Number of hidden layers in the Transformer decoder.301 num_attention_heads (`int`, *optional*, defaults to 32):302 Number of attention heads for each attention layer in the Transformer decoder.303 num_key_value_heads (`int`, *optional*, defaults to 8):304 This is the number of key_value heads that should be used to implement Grouped Query Attention. If305 `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if306 `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When307 converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed308 by meanpooling all the original heads within that group. For more details, check out [this309 paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to310 `num_attention_heads`.311 resid_pdrop (`float`, *optional*, defaults to 0.0):312 Dropout probability for mlp outputs.313 embd_pdrop (`int`, *optional*, defaults to 0.0):314 The dropout ratio for the embeddings.315 attention_dropout (`float`, *optional*, defaults to 0.0):316 The dropout ratio after computing the attention scores.317 hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):318 The non-linear activation function (function or string) in the decoder.319 max_position_embeddings (`int`, *optional*, defaults to 131072):320 The maximum sequence length that this model might ever be used with.321 initializer_range (`float`, *optional*, defaults to 0.02):322 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.323 rms_norm_eps (`float`, *optional*, defaults to 1e-05):324 The epsilon value used for the RMSNorm.325 use_cache (`bool`, *optional*, defaults to `True`):326 Whether or not the model should return the last key/values attentions (not used by all models). Only327 relevant if `config.is_decoder=True`. Whether to tie weight embeddings or not.328 tie_word_embeddings (`bool`, *optional*, defaults to `False`):329 Whether to tie weight embeddings330 rope_theta (`float`, *optional*, defaults to 10000.0):331 The base period of the RoPE embeddings.332 rope_scaling (`dict`, *optional*):333 The scaling strategy for the RoPE embeddings. If `None`, no scaling is applied. If a dictionary, it must334 contain the following keys: `type`, `short_factor` and `long_factor`. The `type` must be `longrope` and335 the `short_factor` and `long_factor` must be lists of numbers with the same length as the hidden size336 divided by the number of attention heads divided by 2.337 partial_rotary_factor (`float`, *optional*, defaults to `1.0`):338 Percentage of the query and keys which will have rotary embedding. Must be between 0.0 and 1.0.339 bos_token_id (`int`, *optional*, defaults to 199999):340 The id of the "beginning-of-sequence" token.341 eos_token_id (`int` or `list[int]`, *optional*, defaults to `[199999, 200020]`):342 The id of the "end-of-sequence" token.343 pad_token_id (`int`, *optional*, defaults to 199999):344 The id of the padding token.345 original_max_position_embeddings (`int`, *optional*, defaults to 4096):346 The maximum sequence length that this model was trained with. This is used to determine the size of the347 original RoPE embeddings when using long scaling.348 sliding_window (`int`, *optional*):349 Sliding window attention window size. If `None`, no sliding window is applied.350 vision_config (`Phi4MultimodalVisionConfig` or `dict`, *optional*):351 The vision config for the underlying image embedding model. If not provided, will default to the configuration352 used to instantiate a model similar in architecture as353 [microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct).354 audio_config (`Phi4MultimodalAudioConfig` or `dict`, *optional*):355 The audio config for the underlying audio embedding model. If not provided, will default to the configuration356 used to instantiate a model similar in architecture as357 [microsoft/Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct).358 359 Example:360 361 ```python362 >>> from transformers import Phi4MultimodalModel, Phi4MultimodalConfig363 364 >>> # Initializing a Phi4Multimodal style configuration365 >>> configuration = Phi4MultimodalConfig.from_pretrained("microsoft/Phi-4-multimodal-instruct")366 367 >>> # Initializing a model from the configuration368 >>> model = Phi4MultimodalModel(configuration)369 370 >>> # Accessing the model configuration371 >>> configuration = model.config372 ```"""373 374 sub_configs = {"audio_config": Phi4MultimodalAudioConfig, "vision_config": Phi4MultimodalVisionConfig}375 376 def __init__(377 self,378 vocab_size=200064,379 hidden_size=3072,380 intermediate_size=8192,381 num_hidden_layers=32,382 num_attention_heads=32,383 num_key_value_heads=8,384 resid_pdrop=0.0,385 embd_pdrop=0.0,386 attention_dropout=0.0,387 hidden_act="silu",388 max_position_embeddings=131072,389 initializer_range=0.02,390 rms_norm_eps=1e-5,391 use_cache=True,392 tie_word_embeddings=False,393 rope_theta=10000.0,394 rope_scaling=None,395 partial_rotary_factor=1,396 bos_token_id=199999,397 eos_token_id=[199999, 200020],398 pad_token_id=199999,399 original_max_position_embeddings=4096,400 sliding_window=None,401 vision_config=None,402 audio_config=None,403 **kwargs,404 ):405 super().__init__(406 vocab_size=vocab_size,407 hidden_size=hidden_size,408 intermediate_size=intermediate_size,409 num_hidden_layers=num_hidden_layers,410 num_attention_heads=num_attention_heads,411 num_key_value_heads=num_key_value_heads,412 resid_pdrop=resid_pdrop,413 embd_pdrop=embd_pdrop,414 attention_dropout=attention_dropout,415 hidden_act=hidden_act,416 max_position_embeddings=max_position_embeddings,417 initializer_range=initializer_range,418 rms_norm_eps=rms_norm_eps,419 use_cache=use_cache,420 tie_word_embeddings=tie_word_embeddings,421 rope_theta=rope_theta,422 rope_scaling=rope_scaling,423 partial_rotary_factor=partial_rotary_factor,424 bos_token_id=bos_token_id,425 eos_token_id=eos_token_id,426 pad_token_id=pad_token_id,427 original_max_position_embeddings=original_max_position_embeddings,428 sliding_window=sliding_window,429 **kwargs,430 )431 432 if isinstance(vision_config, dict):433 vision_config = Phi4MultimodalVisionConfig(**vision_config)434 elif vision_config is None:435 Phi4MultimodalVisionConfig()436 self.vision_config = vision_config437 438 if isinstance(audio_config, dict):439 audio_config = Phi4MultimodalAudioConfig(**audio_config)440 elif vision_config is None:441 audio_config = Phi4MultimodalAudioConfig()442 self.audio_config = audio_config443 444 445class Phi4MultimodalVisionMLP(SiglipMLP):446 pass447 448 449def simple_eager_attention_forward(450 module: nn.Module,451 query_states: torch.Tensor,452 key_states: torch.Tensor,453 value_states: torch.Tensor,454 attention_mask: Optional[torch.Tensor],455 scaling: float,456 dropout: float = 0.0,457 **kwargs: Unpack[TransformersKwargs],458):459 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * scaling460 if attention_mask is not None:461 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]462 attn_weights = attn_weights + causal_mask463 464 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)465 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)466 attn_output = torch.matmul(attn_weights, value_states)467 attn_output = attn_output.transpose(1, 2).contiguous()468 469 return attn_output, attn_weights470 471 472class Phi4MultimodalVisionAttention(nn.Module):473 def __init__(self, config: Phi4MultimodalVisionConfig):474 super().__init__()475 self.config = config476 self.embed_dim = config.hidden_size477 self.num_heads = config.num_attention_heads478 self.head_dim = self.embed_dim // self.num_heads479 self.scaling = self.head_dim**-0.5480 self.is_causal = True481 self.attention_dropout = config.attention_dropout482 483 self.k_proj = nn.Linear(config.hidden_size, config.hidden_size)484 self.v_proj = nn.Linear(config.hidden_size, config.hidden_size)485 self.q_proj = nn.Linear(config.hidden_size, config.hidden_size)486 self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)487 488 def forward(489 self,490 hidden_states: torch.Tensor,491 attention_mask: Optional[torch.Tensor] = None,492 **kwargs: Unpack[TransformersKwargs],493 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:494 """Input shape: Batch x Time x Channel"""495 input_shape = hidden_states.shape[:-1]496 hidden_shape = (*input_shape, -1, self.head_dim)497 498 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)499 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)500 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)501 502 attention_interface: Callable = simple_eager_attention_forward503 if self.config._attn_implementation != "eager":504 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]505 506 attn_output, attn_weights = attention_interface(507 self,508 query_states,509 key_states,510 value_states,511 attention_mask,512 dropout=0.0 if not self.training else self.attention_dropout,513 scaling=self.scaling,514 **kwargs,515 )516 517 attn_output = attn_output.reshape(*input_shape, -1)518 attn_output = self.out_proj(attn_output)519 return attn_output, attn_weights520 521 522class Phi4MultimodalVisionEncoderLayer(SiglipEncoderLayer):523 def __init__(self, config: Phi4MultimodalVisionConfig):524 super().__init__(config)525 self.self_attn = Phi4MultimodalVisionAttention(config)526 self.mlp = Phi4MultimodalVisionMLP(config)527 528 529class Phi4MultimodalVisionEncoder(SiglipEncoder):530 def __init__(self, config: Phi4MultimodalVisionConfig):531 super().__init__(config)532 self.layers = nn.ModuleList(533 [Phi4MultimodalVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]534 )535 536 537class Phi4MultimodalVisionPreTrainedModel(SiglipPreTrainedModel):538 config: Phi4MultimodalVisionConfig539 base_model_prefix = "phi4_vision"540 supports_gradient_checkpointing = True541 542 _no_split_modules = ["Phi4MultimodalVisionEncoderLayer"]543 _supports_flash_attn = True544 _supports_sdpa = True545 _supports_flex_attn = True546 547 _can_record_outputs = {548 "hidden_states": Phi4MultimodalVisionEncoderLayer,549 "attentions": Phi4MultimodalVisionAttention,550 }551 552 def _init_weights(self, module):553 """Initialize the weights"""554 if isinstance(module, Phi4MultimodalVisionEmbeddings):555 width = (556 self.config.hidden_size557 if isinstance(self.config, Phi4MultimodalVisionConfig)558 else self.config.hidden_size559 )560 nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width))561 elif isinstance(module, nn.Embedding):562 default_flax_embed_init(module.weight)563 elif isinstance(module, Phi4MultimodalVisionAttention):564 nn.init.normal_(module.q_proj.weight)565 nn.init.normal_(module.k_proj.weight)566 nn.init.normal_(module.v_proj.weight)567 nn.init.normal_(module.out_proj.weight)568 nn.init.zeros_(module.q_proj.bias)569 nn.init.zeros_(module.k_proj.bias)570 nn.init.zeros_(module.v_proj.bias)571 nn.init.zeros_(module.out_proj.bias)572 elif isinstance(module, Phi4MultimodalVisionMLP):573 nn.init.normal_(module.fc1.weight)574 nn.init.normal_(module.fc2.weight)575 nn.init.normal_(module.fc1.bias, std=1e-6)576 nn.init.normal_(module.fc2.bias, std=1e-6)577 elif isinstance(module, Phi4MultimodalVisionMultiheadAttentionPoolingHead):578 nn.init.normal_(module.probe.data)579 nn.init.normal_(module.attention.in_proj_weight.data)580 nn.init.zeros_(module.attention.in_proj_bias.data)581 elif isinstance(module, (nn.Linear, nn.Conv2d)):582 lecun_normal_(module.weight)583 if module.bias is not None:584 nn.init.zeros_(module.bias)585 elif isinstance(module, nn.LayerNorm):586 module.bias.data.zero_()587 module.weight.data.fill_(1.0)588 589 590class Phi4MultimodalVisionEmbeddings(SiglipVisionEmbeddings):591 def __init__(self, config: Phi4MultimodalVisionConfig):592 nn.Module.__init__(self)593 self.config = config594 self.patch_size = config.patch_size595 self.num_patches_per_side = config.image_size // self.patch_size596 597 self.patch_embedding = nn.Conv2d(598 in_channels=config.num_channels,599 out_channels=config.hidden_size,600 kernel_size=self.patch_size,601 stride=self.patch_size,602 padding="valid",603 )604 self.position_embedding = nn.Embedding(self.num_patches_per_side**2, config.hidden_size)605 606 def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.BoolTensor) -> torch.Tensor:607 batch_size = pixel_values.size(0)608 609 patch_embeds = self.patch_embedding(pixel_values)610 embeddings = patch_embeds.flatten(2).transpose(1, 2)611 612 max_im_h, max_im_w = pixel_values.size(2), pixel_values.size(3)613 max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size614 boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side)615 position_ids = torch.full((batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0)616 617 for batch_idx, p_attn_mask in enumerate(patch_attention_mask):618 nb_patches_h = p_attn_mask[:, 0].sum()619 nb_patches_w = p_attn_mask[0].sum()620 621 fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h)622 fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w)623 624 bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)625 bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)626 627 pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()628 position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids629 630 position_ids = position_ids.to(self.position_embedding.weight.device)631 632 embeddings = embeddings + self.position_embedding(position_ids)633 return embeddings634 635 636class Phi4MultimodalVisionMultiheadAttentionPoolingHead(SiglipMultiheadAttentionPoolingHead):637 def __init__(self, config: Phi4MultimodalVisionConfig):638 super().__init__(config)639 self.mlp = Phi4MultimodalVisionMLP(config)640 641 def forward(self, hidden_state, attention_mask):642 batch_size = hidden_state.shape[0]643 probe = self.probe.repeat(batch_size, 1, 1)644 645 hidden_state = self.attention(646 query=probe, key=hidden_state, value=hidden_state, key_padding_mask=~attention_mask647 )[0]648 649 residual = hidden_state650 hidden_state = self.layernorm(hidden_state)651 hidden_state = residual + self.mlp(hidden_state)652 653 return hidden_state[:, 0]654 655 656class Phi4MultimodalVisionModel(Phi4MultimodalVisionPreTrainedModel):657 config: Phi4MultimodalVisionConfig658 main_input_name = "pixel_values"659 660 def __init__(self, config: Phi4MultimodalVisionConfig):661 super().__init__(config)662 self.config = config663 664 self.embeddings = Phi4MultimodalVisionEmbeddings(config)665 self.encoder = Phi4MultimodalVisionEncoder(config)666 self.post_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)667 self.head = Phi4MultimodalVisionMultiheadAttentionPoolingHead(config)668 669 # Initialize weights and apply final processing670 self.post_init()671 672 def get_input_embeddings(self) -> nn.Module:673 return self.embeddings.patch_embedding674 675 @check_model_inputs(tie_last_hidden_states=False)676 def forward(677 self,678 pixel_values,679 patch_attention_mask: Optional[torch.BoolTensor] = None,680 **kwargs: Unpack[TransformersKwargs],681 ) -> BaseModelOutputWithPooling:682 batch_size = pixel_values.size(0)683 if patch_attention_mask is None:684 patch_attention_mask = torch.ones(685 size=(686 batch_size,687 pixel_values.size(2) // self.config.patch_size,688 pixel_values.size(3) // self.config.patch_size,689 ),690 dtype=torch.bool,691 device=pixel_values.device,692 )693 694 hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)695 696 patch_attention_mask = patch_attention_mask.view(batch_size, -1)697 # The call to `_upad_input` in `_flash_attention_forward` is expensive698 # So when the `patch_attention_mask` is full of 1s (i.e. attending to the whole sequence),699 # avoiding passing the attention_mask, which is equivalent to attending to the full sequence700 if not torch.any(~patch_attention_mask):701 attention_mask = None702 else:703 attention_mask = (704 _prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype)705 if self.config._attn_implementation != "flash_attention_2"706 else patch_attention_mask707 )708 709 encoder_outputs: BaseModelOutput = self.encoder(710 inputs_embeds=hidden_states,711 attention_mask=attention_mask,712 **kwargs,713 )714 715 last_hidden_state = encoder_outputs.last_hidden_state716 last_hidden_state = self.post_layernorm(last_hidden_state)717 718 pooled_output = self.head(719 hidden_state=last_hidden_state,720 attention_mask=patch_attention_mask,721 )722 723 return BaseModelOutputWithPooling(724 last_hidden_state=last_hidden_state,725 pooler_output=pooled_output,726 )727 728 729class Phi4MultimodalImageEmbedding(nn.Module):730 """Image embedding."""731 732 def __init__(self, config: Phi4MultimodalConfig):733 super().__init__()734 self.config = config735 self.layer_idx = config.vision_config.feature_layer736 self.crop_size = config.vision_config.crop_size737 self.image_dim_out = config.vision_config.hidden_size738 739 n_patches = config.vision_config.image_size // config.vision_config.patch_size740 if n_patches % 2 != 0:741 self.img_processor_padding = nn.ReflectionPad2d((0, 1, 0, 1))742 n_patches += 1743 self.num_img_tokens = (n_patches // 2) ** 2744 745 self.drop = nn.Dropout(config.embd_pdrop)746 self.img_processor = Phi4MultimodalVisionModel._from_config(config.vision_config)747 self.image_token_compression = nn.AvgPool2d(kernel_size=2, stride=2)748 self.img_projection_up = nn.Linear(self.image_dim_out, config.hidden_size)749 self.img_projection_down = nn.Linear(config.hidden_size, config.hidden_size)750 self.global_img_feature_extensor = nn.Parameter(torch.zeros([1, 1, self.image_dim_out]))751 self.sub_img_feature_extensor = nn.Parameter(torch.zeros([1, 1, 1, self.image_dim_out]))752 753 def get_img_features(self, img_embeds: torch.FloatTensor, attention_mask=None) -> torch.FloatTensor:754 img_processor_output = self.img_processor(755 img_embeds, patch_attention_mask=attention_mask, output_hidden_states=True756 )757 img_feature = img_processor_output.hidden_states[self.layer_idx]758 759 patch_feature = img_feature760 # reshape to 2D tensor761 width = int(math.sqrt(patch_feature.size(1)))762 patch_feature = patch_feature.view(-1, width, width, patch_feature.size(-1))763 # convert to NCHW764 patch_feature = patch_feature.permute(0, 3, 1, 2)765 if getattr(self, "img_processor_padding", None) is not None:766 patch_feature = self.img_processor_padding(patch_feature)767 patch_feature = self.image_token_compression(patch_feature)768 # convert to NHWC769 patch_feature = patch_feature.permute(0, 2, 3, 1)770 patch_feature = patch_feature.view(-1, patch_feature.size(1) * patch_feature.size(2), patch_feature.size(-1))771 return patch_feature772 773 def forward(774 self,775 input_ids: torch.LongTensor,776 inputs_embeds: torch.Tensor,777 image_pixel_values: torch.FloatTensor,778 image_sizes: Optional[torch.Tensor] = None,779 image_attention_mask: Optional[torch.Tensor] = None,780 ) -> torch.FloatTensor:781 image_pixel_values = image_pixel_values.to(self.img_processor.embeddings.patch_embedding.weight.dtype)782 783 target_device = self.img_projection_up.bias.device784 target_dtype = self.img_projection_up.bias.dtype785 786 batch_size = image_pixel_values.shape[0]787 788 img_features = self.get_img_features(789 image_pixel_values.flatten(0, 1),790 attention_mask=image_attention_mask.flatten(0, 1).to(dtype=bool, device=target_device),791 )792 base_feat_size = int(np.sqrt(img_features.shape[1]))793 img_features = img_features.view(batch_size, -1, base_feat_size**2, self.image_dim_out)794 image_sizes = image_sizes.view(-1, 2)795 796 output_imgs = []797 for idx in range(batch_size):798 height, width = image_sizes[idx]799 height_ratio = height // self.crop_size800 width_ratio = width // self.crop_size801 area_ratio = height_ratio * width_ratio802 803 global_img = img_features[idx, :1]804 global_img = global_img.reshape(1, base_feat_size, base_feat_size, self.image_dim_out).contiguous()805 temporary_extensor = self.sub_img_feature_extensor.repeat(1, base_feat_size, 1, 1)806 global_img = torch.cat([global_img, temporary_extensor], dim=2).reshape(1, -1, self.image_dim_out)807 808 sub_img = img_features[idx, 1:]809 sub_img = sub_img[:area_ratio]810 sub_img = (811 sub_img.reshape(height_ratio, width_ratio, base_feat_size, base_feat_size, self.image_dim_out)812 .transpose(1, 2)813 .reshape(1, height_ratio * base_feat_size, width_ratio * base_feat_size, self.image_dim_out)814 .contiguous()815 )816 817 if image_attention_mask is not None:818 reshaped_image_attention_mask = (819 image_attention_mask[idx, 1 : area_ratio + 1, 0::2, 0::2]820 .reshape(height_ratio, width_ratio, base_feat_size, base_feat_size)821 .transpose(1, 2)822 .reshape(1, height_ratio * base_feat_size, width_ratio * base_feat_size)823 )824 useful_height = int(reshaped_image_attention_mask[0, :, 0].sum().item())825 useful_width = int(reshaped_image_attention_mask[0, 0, :].sum().item())826 sub_img = sub_img[:, :useful_height, :useful_width]827 temporary_extensor = self.sub_img_feature_extensor.repeat(1, useful_height, 1, 1)828 else:829 temporary_extensor = self.sub_img_feature_extensor.repeat(1, height_ratio * base_feat_size, 1, 1)830 831 sub_img = torch.cat([sub_img, temporary_extensor], dim=2).reshape(1, -1, self.image_dim_out)832 833 # Merge global and sub834 output_imgs.append(torch.cat([sub_img, self.global_img_feature_extensor, global_img], dim=1))835 836 img_set_tensor = []837 for output_img in output_imgs:838 output_img = output_img.to(device=target_device, dtype=target_dtype)839 img_feature_proj = self.img_projection_up(output_img)840 img_feature_proj = nn.functional.gelu(img_feature_proj)841 img_feature_proj = self.img_projection_down(img_feature_proj)842 img_set_tensor.append(img_feature_proj)843 844 merged_img_set_tensor = torch.cat(img_set_tensor, dim=1).squeeze(0)845 merged_img_set_tensor = merged_img_set_tensor.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)846 847 with torch.no_grad():848 positions_tuple = torch.nonzero(input_ids == self.config.vision_config.image_token_id, as_tuple=True)849 850 # Temporarily disable autocast to avoid issue on bf16 tensors851 # Ref: https://github.com/pytorch/pytorch/issues/132715852 with torch.autocast(device_type=inputs_embeds.device.type, enabled=False):853 image_embeds = inputs_embeds.index_put(854 indices=positions_tuple, values=merged_img_set_tensor, accumulate=False855 )856 857 image_embeds = self.drop(image_embeds)858 859 return image_embeds860 861 862########################################################## AUDIO #############################################863 864 865class Phi4MultimodalAudioMLP(nn.Module):866 def __init__(self, config: Phi4MultimodalAudioConfig):867 super().__init__()868 self.layer_norm = nn.LayerNorm(config.hidden_size)869 self.act_fn = ACT2FN[config.activation]870 self.gate_up_proj = nn.Linear(config.hidden_size, config.intermediate_size * 2)871 self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size)872 self.dropout = nn.Dropout(config.dropout_rate)873 874 def forward(self, hidden_states):875 hidden_states = self.layer_norm(hidden_states)876 up_states = self.gate_up_proj(hidden_states)877 up_states, gate = up_states.chunk(2, dim=-1)878 up_states = up_states * self.act_fn(gate)879 up_states = self.dropout(up_states)880 hidden_states = self.down_proj(up_states)881 out = self.dropout(hidden_states)882 883 return out884 885 886class Phi4MultimodalAudioAttention(nn.Module):887 def __init__(self, config: Phi4MultimodalAudioConfig):888 super().__init__()889 self.config = config890 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)891 self.scaling = self.head_dim**-0.5892 self.attention_dropout = config.dropout_rate893 self.is_causal = True894 895 self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True)896 self.k_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True)897 self.v_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True)898 self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=True)899 900 def forward(901 self,902 hidden_states: torch.Tensor,903 attention_mask: torch.Tensor,904 **kwargs,905 ):906 input_shape = hidden_states.shape[:-1]907 hidden_shape = (*input_shape, -1, self.head_dim)908 909 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)910 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)911 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)912 913 attention_interface: Callable = simple_eager_attention_forward914 if self.config._attn_implementation != "eager":915 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]916 917 attn_output, _ = attention_interface(918 self,919 query_states,920 key_states,921 value_states,922 attention_mask,923 dropout=0.0 if not self.training else self.attention_dropout,924 scaling=self.scaling,925 **kwargs,926 )927 928 attn_output = attn_output.reshape(*input_shape, -1).contiguous()929 attn_output = self.o_proj(attn_output)930 return attn_output931 932 933class Phi4MultimodalAudioDepthWiseSeparableConv1d(nn.Module):934 def __init__(self, config: Phi4MultimodalAudioConfig, padding: int = 0):935 super().__init__()936 self.dw_conv = nn.Conv1d(937 config.hidden_size,938 config.hidden_size * config.depthwise_multiplier,939 config.kernel_size,940 1,941 padding=padding,942 groups=config.hidden_size,943 )944 self.pw_conv = nn.Conv1d(945 config.hidden_size * config.depthwise_multiplier, config.depthwise_separable_out_channel, 1, 1, 0946 )947 948 def forward(self, hidden_states):949 return self.pw_conv(self.dw_conv(hidden_states))950 951 952class Phi4MultimodalAudioGluPointWiseConv(nn.Module):953 def __init__(self, config: Phi4MultimodalAudioConfig):954 super().__init__()955 self.config = config956 self.output_dim = config.ext_pw_out_channel957 958 self.ext_pw_conv_1d = nn.Conv1d(config.hidden_size, config.ext_pw_out_channel * 2, kernel_size=1, stride=1)959 self.glu_act = ACT2FN[config.conv_glu_type]960 self.b1 = nn.Parameter(torch.zeros(1, config.ext_pw_out_channel, 1))961 self.b2 = nn.Parameter(torch.zeros(1, config.ext_pw_out_channel, 1))962 963 def forward(self, hidden_states):964 # we assume the input always has the #channel (#dim) in the last dimension of the965 # tensor, so need to switch the dimension first for 1D-Conv case966 hidden_states = hidden_states.permute([0, 2, 1])967 hidden_states = self.ext_pw_conv_1d(hidden_states)968 out = hidden_states[:, 0 : self.output_dim, :] + self.b1969 out = out * self.glu_act(hidden_states[:, self.output_dim : self.output_dim * 2, :] + self.b2)970 return out.permute([0, 2, 1])971 972 973class Phi4MultimodalAudioConvModule(nn.Module):974 def __init__(self, config: Phi4MultimodalAudioConfig):975 super().__init__()976 self.config = config977 self.kernel_size = config.kernel_size978 979 self.layer_norm = nn.LayerNorm(config.hidden_size)980 self.glu = Phi4MultimodalAudioGluPointWiseConv(config)981 self.dw_sep_conv_1d = Phi4MultimodalAudioDepthWiseSeparableConv1d(config, padding=config.kernel_size - 1)982 self.act = ACT2FN[config.conv_activation]983 self.ext_pw_conv_1d = nn.Conv1d(config.hidden_size, config.ext_pw_out_channel, kernel_size=1, stride=1)984 self.dropout = nn.Dropout(config.dropout_rate)985 986 def forward(self, hidden_states: torch.Tensor):987 hidden_states = self.glu(self.layer_norm(hidden_states))988 hidden_states = self.dw_sep_conv_1d(hidden_states.permute([0, 2, 1]))989 990 if self.kernel_size > 1:991 hidden_states = hidden_states[:, :, : -(self.kernel_size - 1)]992 993 hidden_states = self.act(hidden_states)994 hidden_states = self.ext_pw_conv_1d(hidden_states)995 out = self.dropout(hidden_states.permute([0, 2, 1]))996 return out997 998 999class Phi4MultimodalAudioConformerEncoderLayer(nn.Module):1000 def __init__(self, config: Phi4MultimodalAudioConfig):1001 super().__init__()1002 1003 self.feed_forward_in = Phi4MultimodalAudioMLP(config)1004 self.self_attn = Phi4MultimodalAudioAttention(config)1005 self.conv = Phi4MultimodalAudioConvModule(config)1006 self.feed_forward_out = Phi4MultimodalAudioMLP(config)1007 self.layer_norm_att = nn.LayerNorm(config.hidden_size)1008 self.layer_norm = nn.LayerNorm(config.hidden_size)1009 1010 def forward(1011 self,1012 hidden_states: torch.Tensor,1013 attention_mask: torch.Tensor,1014 ):1015 residual = hidden_states + 0.5 * self.feed_forward_in(hidden_states)1016 hidden_states = self.layer_norm_att(residual)1017 1018 hidden_states = residual + self.self_attn(hidden_states, attention_mask)1019 hidden_states = hidden_states + self.conv(hidden_states)1020 hidden_states = hidden_states + 0.5 * self.feed_forward_out(hidden_states)1021 1022 out = self.layer_norm(hidden_states)1023 1024 return out1025 1026 1027class Phi4MultimodalAudioNemoConvSubsampling(torch.nn.Module):1028 def __init__(self, config: Phi4MultimodalAudioConfig):1029 super().__init__()1030 self.subsampling_factor = config.time_reduction1031 self.sampling_num = int(math.log2(self.subsampling_factor))1032 self.act_fn = ACT2FN[config.nemo_activation]1033 conv_channels = config.nemo_conv_channels1034 1035 layers = [1036 nn.Conv2d(1, conv_channels, kernel_size=3, stride=2, padding=1),1037 self.act_fn,1038 ]1039 for _ in range(self.sampling_num - 1):1040 layers.extend(1041 [1042 nn.Conv2d(conv_channels, conv_channels, kernel_size=3, stride=2, padding=1, groups=conv_channels),1043 nn.Conv2d(conv_channels, conv_channels, kernel_size=1, stride=1, padding=0, groups=1),1044 self.act_fn,1045 ]1046 )1047 1048 # Aggregate the layers1049 self.conv = torch.nn.Sequential(*layers)1050 self.out = torch.nn.Linear(conv_channels * config.nemo_final_size, config.hidden_size)1051 1052 def forward(self, hidden_states: torch.Tensor, mask: Optional[torch.Tensor]):1053 # Unsqueeze Channel Axis1054 hidden_states = hidden_states.unsqueeze(1)1055 hidden_states = self.conv(hidden_states)1056 1057 # Flatten Channel and Frequency Axes1058 b, _, t, _ = hidden_states.size()1059 hidden_states = self.out(hidden_states.transpose(1, 2).reshape(b, t, -1))1060 1061 if mask is None:1062 return hidden_states, None1063 1064 max_audio_length = hidden_states.shape[1]1065 feature_lens = mask.sum(1)1066 padding_length = torch.ceil(feature_lens / self.subsampling_factor)1067 arange_ = torch.arange(0, max_audio_length, device=hidden_states.device)1068 pad_mask = arange_.expand(padding_length.size(0), -1) < padding_length.unsqueeze(1)1069 return hidden_states, pad_mask.unsqueeze(1)1070 1071 1072class Phi4MultimodalAudioRelativeAttentionBias(nn.Module):1073 def __init__(self, config: Phi4MultimodalAudioConfig):1074 super().__init__()1075 1076 self.max_distance = config.bias_max_distance1077 self.symmetric = config.bias_symmetric1078 self.num_buckets = self.max_distance1079 if not config.bias_symmetric:1080 self.num_buckets *= 21081 self.bias_values = nn.Embedding(self.num_buckets, config.num_attention_heads)1082 1083 def forward(self, x):1084 # instantiate bias compatible with shape of x1085 max_pos = x.size(1)1086 context_position = torch.arange(max_pos, device=x.device, dtype=torch.long)[:, None]1087 memory_position = torch.arange(max_pos, device=x.device, dtype=torch.long)[None, :]1088 relative_position = memory_position - context_position1089 # clipping to a maximum distance using ops that play well with ONNX export1090 relative_position = relative_position.masked_fill(relative_position < -self.max_distance, -self.max_distance)1091 relative_position = relative_position.masked_fill(1092 relative_position > self.max_distance - 1, self.max_distance - 11093 )1094 1095 # mapping from relative position to index in the bias parameter1096 bias_idx = relative_position1097 bias_idx = bias_idx.abs() if self.symmetric else bias_idx + self.num_buckets // 21098 1099 att_bias = self.bias_values(bias_idx)1100 att_bias = att_bias.permute(2, 0, 1).unsqueeze(0)1101 1102 return att_bias1103 1104 1105class Phi4MultimodalAudioMeanVarianceNormLayer(nn.Module):1106 def __init__(self, config: Phi4MultimodalAudioConfig):1107 super().__init__()1108 self.register_buffer("global_mean", torch.zeros(config.input_size))1109 self.register_buffer("global_invstd", torch.ones(config.input_size))1110 1111 def forward(self, x):1112 return (x - self.global_mean) * self.global_invstd1113 1114 1115@auto_docstring1116class Phi4MultimodalAudioPreTrainedModel(PreTrainedModel):1117 config: Phi4MultimodalAudioConfig1118 supports_gradient_checkpointing = True1119 _no_split_modules = ["Phi4MultimodalAudioConformerEncoderLayer"]1120 _supports_flash_attn = True1121 _supports_sdpa = True1122 _supports_flex_attn = True1123 1124 def _init_weights(self, module):1125 super()._init_weights(module)1126 if isinstance(module, Phi4MultimodalAudioGluPointWiseConv):1127 module.b1.data.zero_()1128 module.b2.data.zero_()1129 1130 1131class Phi4MultimodalAudioModel(Phi4MultimodalAudioPreTrainedModel):1132 def __init__(self, config: Phi4MultimodalAudioConfig):1133 super().__init__(config)1134 self.config = config1135 1136 self.encoder_embedding = Phi4MultimodalAudioMeanVarianceNormLayer(config)1137 self.embed = Phi4MultimodalAudioNemoConvSubsampling(config)1138 self.relative_attention_bias_layer = Phi4MultimodalAudioRelativeAttentionBias(config)1139 self.encoders = nn.ModuleList(1140 [Phi4MultimodalAudioConformerEncoderLayer(config) for _ in range(config.num_blocks)]1141 )1142 self.gradient_checkpointing = False1143 1144 # Initialize weights and apply final processing1145 self.post_init()1146 1147 def _streaming_mask(self, seq_len, batch_size, chunk_size, left_chunk):1148 # Create mask matrix for streaming1149 # S stores start index. if chunksize is 18, s is [0,18,36,....]1150 chunk_start_idx = np.arange(0, seq_len, chunk_size)1151 # avoid randomness when run evaluation or decoding1152 if self.training and np.random.rand() > 0.5:1153 # Either first or last chunk is not complete.1154 # If only the last one is not complete, EOS is not effective1155 chunk_start_idx = seq_len - chunk_start_idx1156 chunk_start_idx = chunk_start_idx[::-1]1157 chunk_start_idx = chunk_start_idx[:-1]1158 chunk_start_idx = np.insert(chunk_start_idx, 0, 0)1159 1160 enc_streaming_mask = (1161 adaptive_enc_mask(seq_len, chunk_start_idx, left_window=left_chunk)1162 .unsqueeze(0)1163 .expand([batch_size, -1, -1])1164 )1165 return enc_streaming_mask1166 1167 def forward_embeddings(self, hidden_states, masks):1168 """Forwarding the inputs through the top embedding layers"""1169 seq_len = math.ceil(hidden_states.shape[1] / self.config.time_reduction)1170 if seq_len <= 0:1171 raise ValueError(1172 f"The sequence length after time reduction is invalid: {seq_len}. Your input feature is too short."1173 )1174 1175 batch_size = hidden_states.shape[0]1176 1177 enc_streaming_mask = self._streaming_mask(seq_len, batch_size, self.config.chunk_size, self.config.left_chunk)1178 enc_streaming_mask = enc_streaming_mask.to(hidden_states.device)1179 1180 hidden_states, masks = self.embed(hidden_states, masks)1181 1182 streaming_mask = enc_streaming_mask1183 if streaming_mask is not None and masks is not None:1184 hs_mask = masks & streaming_mask1185 elif masks is not None:1186 hs_mask = masks1187 else:1188 hs_mask = streaming_mask1189 1190 return hidden_states, hs_mask, masks1191 1192 def calculate_hs_mask(self, hidden_states, device, mask):1193 max_audio_length = hidden_states.shape[1]1194 batch_size = hidden_states.shape[0]1195 enc_streaming_mask = self._streaming_mask(1196 max_audio_length, batch_size, self.config.chunk_size, self.config.left_chunk1197 )1198 enc_streaming_mask = enc_streaming_mask.to(device)1199 if mask is None:1200 return enc_streaming_mask