# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/embedding_gemma2/modular_embedding_gemma2.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_embedding_gemma2.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 the HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from collections.abc import Callable from dataclasses import dataclass import torch from torch import nn from ... import initialization as init from ...activations import ACT2FN from ...configuration_utils import PreTrainedConfig from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_bidirectional_mask, create_bidirectional_sliding_window_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ( TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check, ) from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ..auto import AutoModel from .configuration_embedding_gemma2 import EmbeddingGemma2Config, EmbeddingGemma2TextConfig class EmbeddingGemma2RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6, with_scale: bool = True): super().__init__() self.eps = eps self.with_scale = with_scale if self.with_scale: self.weight = nn.Parameter(torch.ones(dim), requires_grad=True) def _norm(self, hidden_states: torch.Tensor): mean_squared = hidden_states.pow(2).mean(-1, keepdim=True) + self.eps # Use torch.pow() (over torch.sqrt() or torch.rsqrt()) to address compiler differences between Torch and JAX return hidden_states * torch.pow(mean_squared, -0.5) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: normed_output = self._norm(hidden_states.float()) if self.with_scale: normed_output = normed_output * self.weight.float() return normed_output.type_as(hidden_states) class EmbeddingGemma2RotaryEmbedding(nn.Module): def __init__(self, config: EmbeddingGemma2TextConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.layer_types = sorted(set(config.layer_types)) self.rope_init_fns: dict[str, Callable[..., tuple[torch.Tensor, float]]] = {} self.rope_type: dict[str, str] = {} for layer_type in self.layer_types: rope_params = self.config.rope_parameters[layer_type] if rope_params is None: continue if (rope_type := rope_params["rope_type"]) != "default": rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type] else: rope_init_fn = self.compute_default_rope_parameters self.rope_init_fns[layer_type] = rope_init_fn self.rope_type[layer_type] = rope_type # `inv_freq` depends on the head dim, which varies by layer type, so initialise # from a config resolved for this layer type rather than the global one. rope_config = config.per_layer_config[layer_type] curr_inv_freq, curr_attention_scaling = rope_init_fn(rope_config, device, layer_type=layer_type) setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False)) setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False)) setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling) @staticmethod def compute_default_rope_parameters( config: EmbeddingGemma2TextConfig, device=None, layer_type: str | None = None, **kwargs ) -> tuple[torch.Tensor, float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. layer_type (`str`): The current layer type if the model has different RoPE parameters per type. Should not be used unless `config.layer_types is not None` Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). """ # For backward compatibility standardize the `rope_parameters_dict` if it uses old format base = config.rope_parameters[layer_type]["rope_theta"] dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) return inv_freq.to(device), attention_factor @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids, layer_type): inv_freq = getattr(self, f"{layer_type}_inv_freq") attention_scaling = getattr(self, f"{layer_type}_attention_scaling") freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * attention_scaling sin = emb.sin() * attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) class EmbeddingGemma2TextScaledWordEmbedding(nn.Embedding): """ This module overrides nn.Embeddings' forward by multiplying with embeddings scale. """ def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float = 1.0): super().__init__(num_embeddings, embedding_dim, padding_idx) self.scalar_embed_scale = embed_scale self.embed_scale = nn.Buffer(torch.tensor(embed_scale), persistent=False) def forward(self, input_ids: torch.Tensor): return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) class EmbeddingGemma2MLP(nn.Module): def __init__(self, config: EmbeddingGemma2TextConfig): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_activation] def forward(self, x): down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj class EmbeddingGemma2TextPLE(nn.Module): """Produces the per-layer embeddings (PLE) that every decoder layer is gated with. EmbeddingGemma 2 uses *projection-only* PLE: the signal is derived from `inputs_embeds` alone, with no token-identity lookup and no blending scale. """ def __init__(self, config: EmbeddingGemma2TextConfig): super().__init__() self.num_hidden_layers = config.num_hidden_layers self.hidden_size_per_layer_input = config.hidden_size_per_layer_input self.per_layer_model_projection = nn.Linear( config.hidden_size, config.num_hidden_layers * config.hidden_size_per_layer_input, bias=False, ) self.per_layer_model_projection_scale = config.hidden_size**-0.5 self.per_layer_projection_norm = EmbeddingGemma2RMSNorm( config.hidden_size_per_layer_input, eps=config.rms_norm_eps ) def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: per_layer_projection = self.per_layer_model_projection(inputs_embeds) * self.per_layer_model_projection_scale per_layer_projection = per_layer_projection.reshape( *inputs_embeds.shape[:-1], self.num_hidden_layers, self.hidden_size_per_layer_input, ) return self.per_layer_projection_norm(per_layer_projection) class EmbeddingGemma2TextPLEBlock(nn.Module): """Mixes one layer's slice of the per-layer embeddings into the residual stream. The third residual sub-block of a decoder layer, after attention and the MLP. The defining operation is the elementwise gate against `per_layer_input`; the two linears are a bottleneck down to `hidden_size_per_layer_input` and back. """ def __init__(self, config: EmbeddingGemma2TextConfig): super().__init__() self.hidden_size = config.hidden_size self.hidden_size_per_layer_input = config.hidden_size_per_layer_input self.act_fn = ACT2FN[config.hidden_activation] self.per_layer_input_gate = nn.Linear(self.hidden_size, self.hidden_size_per_layer_input, bias=False) self.per_layer_projection = nn.Linear(self.hidden_size_per_layer_input, self.hidden_size, bias=False) self.post_per_layer_input_norm = EmbeddingGemma2RMSNorm(self.hidden_size, eps=config.rms_norm_eps) def forward(self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor) -> torch.Tensor: residual = hidden_states hidden_states = self.per_layer_input_gate(hidden_states) hidden_states = self.act_fn(hidden_states) hidden_states = hidden_states * per_layer_input hidden_states = self.per_layer_projection(hidden_states) hidden_states = self.post_per_layer_input_norm(hidden_states) return residual + hidden_states def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @use_kernel_forward_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, dropout: float | int = 0.0, scaling: float | None = None, softcap: float | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: if scaling is None: scaling = module.head_dim**-0.5 key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if softcap is not None: attn_weights = attn_weights / softcap attn_weights = torch.tanh(attn_weights) attn_weights = attn_weights * softcap if attention_mask is not None: attn_weights = attn_weights + attention_mask # upcast attention to fp32 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights @use_kernelized_func(apply_rotary_pos_emb) class EmbeddingGemma2Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper. EmbeddingGemma 2 is an encoder: attention is bidirectional on every layer and no key-value state is carried between calls. """ def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx # +1 is needed because flash attention sets inclusive boundaries (see modeling_flash_attention_utils.py) self.sliding_window = ( config.sliding_window + 1 if config.layer_types[layer_idx] == "sliding_attention" else None ) layer_config = config.per_layer_config[layer_idx] self.head_dim = layer_config.head_dim self.num_key_value_groups = config.num_attention_heads // layer_config.num_key_value_heads self.scaling = 1.0 self.attention_dropout = config.attention_dropout self.is_causal = False self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear( config.hidden_size, layer_config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.v_proj = nn.Linear( config.hidden_size, layer_config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.o_proj = nn.Linear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias ) self.q_norm = EmbeddingGemma2RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps) self.k_norm = EmbeddingGemma2RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps) self.v_norm = EmbeddingGemma2RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps, with_scale=False) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) value_states = self.v_norm(value_states) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=self.attention_dropout if self.training else 0.0, scaling=self.scaling, sliding_window=self.sliding_window, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class EmbeddingGemma2EncoderLayer(GradientCheckpointingLayer): """A single transformer block. Named an *encoder* layer because attention is bidirectional.""" def __init__(self, config: EmbeddingGemma2TextConfig, layer_idx: int): super().__init__() self.config = config self.hidden_size = config.hidden_size self.layer_idx = layer_idx self.self_attn = EmbeddingGemma2Attention(config=config, layer_idx=layer_idx) self.mlp = EmbeddingGemma2MLP(config) self.input_layernorm = EmbeddingGemma2RMSNorm(self.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = EmbeddingGemma2RMSNorm(self.hidden_size, eps=config.rms_norm_eps) self.pre_feedforward_layernorm = EmbeddingGemma2RMSNorm(self.hidden_size, eps=config.rms_norm_eps) self.post_feedforward_layernorm = EmbeddingGemma2RMSNorm(self.hidden_size, eps=config.rms_norm_eps) self.layer_scalar = nn.Buffer(torch.ones(1)) self.ple_block = EmbeddingGemma2TextPLEBlock(config) def forward( self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor, attention_mask: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_embeddings=position_embeddings, **kwargs, ) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.pre_feedforward_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = self.post_feedforward_layernorm(hidden_states) hidden_states = residual + hidden_states hidden_states = self.ple_block(hidden_states, per_layer_input) hidden_states *= self.layer_scalar return hidden_states @auto_docstring class EmbeddingGemma2PreTrainedModel(PreTrainedModel): config: EmbeddingGemma2Config base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["EmbeddingGemma2EncoderLayer"] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True _can_compile_fullgraph = True _supports_attention_backend = True _can_record_outputs = None # override input_modalities = ("image", "text", "video", "audio") @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, EmbeddingGemma2RotaryEmbedding): for layer_type, rope_init_fn in module.rope_init_fns.items(): rope_config = module.config.per_layer_config[layer_type] curr_inv_freq, _ = rope_init_fn(rope_config, layer_type=layer_type) init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq) init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq) elif isinstance(module, EmbeddingGemma2TextScaledWordEmbedding): init.constant_(module.embed_scale, module.scalar_embed_scale) elif isinstance(module, EmbeddingGemma2EncoderLayer): init.ones_(module.layer_scalar) @auto_docstring( custom_intro=""" The EmbeddingGemma 2 text backbone. It owns the `embedding_projection` that maps the final hidden states down to `config.embedding_dim`, so that the composite model needs no `forward` override. """ ) class EmbeddingGemma2TextModel(EmbeddingGemma2PreTrainedModel): config: EmbeddingGemma2TextConfig input_modalities = ("text",) _can_record_outputs = { "hidden_states": EmbeddingGemma2EncoderLayer, "attentions": EmbeddingGemma2Attention, } def __init__(self, config: EmbeddingGemma2TextConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size # bfloat16 rounding turns sqrt(512)=22.6274 into 22.625; see https://github.com/huggingface/transformers/pull/29402 self.embed_tokens = EmbeddingGemma2TextScaledWordEmbedding( config.vocab_size, config.hidden_size, self.padding_idx, embed_scale=self.config.hidden_size**0.5 ) self.layers = nn.ModuleList( [EmbeddingGemma2EncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = EmbeddingGemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = EmbeddingGemma2RotaryEmbedding(config) self.gradient_checkpointing = False self.unique_layer_types = set(self.config.layer_types) self.ple = EmbeddingGemma2TextPLE(config) # Projecting per token is equivalent to projecting after mean pooling self.embedding_projection = nn.Linear(config.hidden_size, config.embedding_dim, bias=False) # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutput: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if input_ids is not None: inputs_embeds = self.embed_tokens(input_ids) per_layer_inputs = self.ple(inputs_embeds) if position_ids is None: position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) # It may already have been prepared as a mapping by the caller if not isinstance(attention_mask_mapping := attention_mask, dict): mask_kwargs = { "config": self.config, "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, } attention_mask_mapping = { "full_attention": create_bidirectional_mask(**mask_kwargs), "sliding_attention": create_bidirectional_sliding_window_mask(**mask_kwargs), } # embed positions hidden_states = inputs_embeds position_embeddings = {} for layer_type in self.unique_layer_types: position_embeddings[layer_type] = self.rotary_emb(hidden_states, position_ids, layer_type) for i, encoder_layer in enumerate(self.layers): hidden_states = encoder_layer( hidden_states, per_layer_inputs[:, :, i, :], attention_mask=attention_mask_mapping[self.config.layer_types[i]], position_embeddings=position_embeddings[self.config.layer_types[i]], **kwargs, ) hidden_states = self.norm(hidden_states) hidden_states = self.embedding_projection(hidden_states) return BaseModelOutput(last_hidden_state=hidden_states) class EmbeddingGemma2MultimodalEmbedder(nn.Module): """Embeds token ids or soft tokens for multimodal content into language model space.""" def __init__( self, multimodal_config: PreTrainedConfig, text_config: EmbeddingGemma2TextConfig, ): super().__init__() self.multimodal_hidden_size = getattr(multimodal_config, "output_proj_dims", multimodal_config.hidden_size) self.eps = multimodal_config.rms_norm_eps self.text_hidden_size = text_config.hidden_size self.embedding_projection = nn.Linear(self.multimodal_hidden_size, self.text_hidden_size, bias=False) self.embedding_pre_projection_norm = EmbeddingGemma2RMSNorm( self.multimodal_hidden_size, eps=self.eps, with_scale=False ) def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: """Embeds token ids or soft tokens for multimodal content into language model space. Args: 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]`. """ embs_normed = self.embedding_pre_projection_norm(inputs_embeds) return self.embedding_projection(embs_normed) @auto_docstring @dataclass class EmbeddingGemma2AudioModelOutput(BaseModelOutputWithPooling): r""" attention_mask (`torch.BoolTensor`, *optional*): A torch.BoolTensor of shape `(batch_size, num_frames)`. True for valid positions, False for padding. """ attention_mask: torch.BoolTensor | None = None @auto_docstring( custom_intro=""" Base class for EmbeddingGemma2 outputs, with hidden states and attentions. """ ) @dataclass class EmbeddingGemma2ModelOutput(BaseModelOutput): r""" image_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. audio_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. audio_hidden_states of the model produced by the audio encoder and after projecting the last hidden state. """ image_hidden_states: torch.FloatTensor | None = None audio_hidden_states: torch.FloatTensor | None = None @auto_docstring( custom_intro=""" The EmbeddingGemma 2 model: a vision backbone, an audio backbone and a text backbone whose final hidden states are projected to `config.text_config.embedding_dim`. Intended to be wrapped by SentenceTransformers' mean pooling and normalization. """ ) class EmbeddingGemma2Model(EmbeddingGemma2PreTrainedModel): # we are filtering the logits/labels so we shouldn't divide the loss based on num_items_in_batch accepts_loss_kwargs = False config: EmbeddingGemma2Config _keys_to_ignore_on_load_unexpected = [ r"(^|\.)vision_tower\.", r"(^|\.)embed_vision\.", r"(^|\.)audio_tower\.", r"(^|\.)embed_audio\.", ] def __init__(self, config: EmbeddingGemma2Config): super().__init__(config) self.vision_tower = AutoModel.from_config(config.vision_config) if config.vision_config is not None else None self.vocab_size = config.text_config.vocab_size self.language_model = AutoModel.from_config(config.text_config) self.audio_tower = AutoModel.from_config(config.audio_config) if config.audio_config is not None else None self.embed_vision = ( EmbeddingGemma2MultimodalEmbedder(config.vision_config, config.text_config) if config.vision_config is not None else None ) self.embed_audio = ( EmbeddingGemma2MultimodalEmbedder(config.audio_config, config.text_config) if config.audio_config is not None else None ) self.post_init() @can_return_tuple @auto_docstring(custom_intro="Projects the last hidden state from the vision model into language model space.") def get_image_features( self, pixel_values: torch.FloatTensor, image_position_ids: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*): The patch positions as (x, y) coordinates in the image. Padding patches are indicated by (-1, -1). """ if self.vision_tower is None: raise ValueError( "Image features were requested, but the model was initialized without a vision_config. " "Cannot process images without a vision tower and vision embedder." ) vision_outputs = self.vision_tower( pixel_values=pixel_values, pixel_position_ids=image_position_ids, **kwargs, ) last_hidden_state = vision_outputs.last_hidden_state pooler_output = self.embed_vision(inputs_embeds=last_hidden_state) output_length = pixel_values.shape[-2] // (self.config.vision_config.pooling_kernel_size**2) k_squared = int((image_position_ids.shape[1] // output_length) ** 0.5) ** 2 non_pad_mask = (image_position_ids != -1).all(dim=-1) split_sizes = (non_pad_mask.sum(dim=-1) // k_squared).tolist() vision_outputs.pooler_output = torch.split(pooler_output, split_sizes) return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, ) -> tuple[torch.BoolTensor, torch.BoolTensor, torch.BoolTensor]: """ Obtains mask for multimodal placeholders (replaced by soft tokens) and hard text tokens. Masks will be obtained from `mm_token_type_ids`, `input_ids`, or `inputs_embeds` as available and in that precedence order. If passing `input_ids` or `inputs_embeds`, the image mask will be derived using `config.image_token_id`. Same goes for audio and video masks Args: input_ids: A tensor containing the hard token IDs from the text tokenizer. inputs_embeds: A tensor containing the embeddings for all hard text tokens. Returns: image_mask, video_mask, audio_mask """ if input_ids is not None: special_image_mask = input_ids == self.config.image_token_id special_video_mask = input_ids == self.config.video_token_id special_audio_mask = input_ids == self.config.audio_token_id else: special_image_mask = ( inputs_embeds == self.get_input_embeddings()( torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) ) ).all(-1) special_video_mask = ( inputs_embeds == self.get_input_embeddings()( torch.full((), self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device) ) ).all(-1) special_audio_mask = ( inputs_embeds == self.get_input_embeddings()( torch.full((), self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device) ) ).all(-1) return special_image_mask, special_video_mask, special_audio_mask @merge_with_config_defaults @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, pixel_values: torch.FloatTensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, input_features: torch.FloatTensor | None = None, attention_mask: torch.Tensor | None = None, input_features_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, image_position_ids: torch.LongTensor | None = None, video_position_ids: torch.LongTensor | None = None, num_frames_per_video: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> EmbeddingGemma2ModelOutput: r""" input_features_mask (`torch.FloatTensor` of shape `(num_images, seq_length)`): The attention mask for the input audio. image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*): 2D patch position coordinates from the image processor, with `(-1, -1)` indicating padding. Passed through to the vision encoder for positional embedding computation. video_position_ids (`torch.LongTensor` of shape `(total_num_frames, max_patches, 2)`, *optional*): 2D patch position coordinates from the video processor, with `(-1, -1)` indicating padding. Passed through to the vision encoder for positional embedding computation. num_frames_per_video (`torch.LongTensor` of shape `(num_videos,)`, *optional*): Number of frames belonging to each video. Required whenever `pixel_values_videos` is passed, since the frames of all videos are concatenated along a single axis. """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") image_mask, video_mask, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds) multimodal_mask = image_mask | video_mask | audio_mask # Replace image id with PAD if the image token if OOV, to avoid index-errors if inputs_embeds is None: llm_input_ids = torch.where(multimodal_mask, self.config.text_config.pad_token_id, input_ids) inputs_embeds = self.get_input_embeddings()(llm_input_ids) # Merge text and images if pixel_values is not None: image_features = self.get_image_features(pixel_values, image_position_ids, return_dict=True).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) # Confirm the number of soft tokens from the vision tower matches the number of slots in the embeddings. n_image_tokens = image_mask.sum() image_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) torch_compilable_check( inputs_embeds[image_mask].numel() == image_features.numel(), f"Image features and image tokens do not match, tokens: {n_image_tokens}, features:" f" {image_features.shape[0]}", ) inputs_embeds = inputs_embeds.masked_scatter( image_mask.to(inputs_embeds.device), image_features.to(inputs_embeds.device) ) if pixel_values_videos is not None: video_features = self.get_video_features( pixel_values_videos, video_position_ids, num_frames_per_video, return_dict=True ).pooler_output video_features = torch.cat(video_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) # Confirm the number of soft tokens from the vision tower matches the number of slots in the embeddings. n_video_tokens = video_mask.sum() video_mask = video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) torch_compilable_check( inputs_embeds[video_mask].numel() == video_features.numel(), f"Video features and video tokens do not match, tokens: {n_video_tokens}, features:" f" {video_features.shape[0]}", ) inputs_embeds = inputs_embeds.masked_scatter( video_mask.to(inputs_embeds.device), video_features.to(inputs_embeds.device) ) # Merge text and audio if input_features is not None and input_features_mask is not None: audio_output = self.get_audio_features(input_features, input_features_mask, return_dict=True) audio_features = audio_output.pooler_output audio_mask_from_encoder = audio_output.attention_mask # True = valid # Keep only real audio soft tokens, mirroring the vision encoder's padding stripping. audio_features = audio_features[audio_mask_from_encoder.to(audio_features.device)] n_audio_tokens = audio_mask.sum() audio_mask = audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) torch_compilable_check( inputs_embeds[audio_mask].numel() == audio_features.numel(), f"Audio features and audio tokens do not match, tokens: {n_audio_tokens}, features:" f" {audio_features.shape[0] * audio_features.shape[1]}", ) inputs_embeds = inputs_embeds.masked_scatter( audio_mask.to(inputs_embeds.device), audio_features.to(inputs_embeds.device, inputs_embeds.dtype) ) outputs = self.language_model( attention_mask=attention_mask, position_ids=position_ids, inputs_embeds=inputs_embeds, return_dict=True, **kwargs, ) return EmbeddingGemma2ModelOutput( last_hidden_state=outputs.last_hidden_state, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=image_features if pixel_values is not None else None, audio_hidden_states=audio_features if input_features is not None else None, ) @can_return_tuple @auto_docstring(custom_intro="Projects the last hidden state from the audio encoder into language model space.") def get_audio_features( self, input_features: torch.Tensor, input_features_mask: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> tuple | EmbeddingGemma2AudioModelOutput: r""" input_features (`torch.FloatTensor` of shape `(num_images, seq_length, num_features)`): The tensors corresponding to the input audio. input_features_mask (`torch.FloatTensor` of shape `(num_images, seq_length)`): The attention mask for the input audio. """ if self.audio_tower is None: raise ValueError( "Audio features were requested, but the model was initialized without an audio_config. " "Cannot process audio without an audio tower and audio embedder." ) audio_outputs = self.audio_tower(input_features, input_features_mask, return_dict=True, **kwargs) audio_outputs.pooler_output = self.embed_audio(inputs_embeds=audio_outputs.last_hidden_state) return audio_outputs @can_return_tuple @auto_docstring(custom_intro="Projects the last hidden state from the vision encoder into language model space.") def get_video_features( self, pixel_values_videos: torch.FloatTensor, video_position_ids: torch.LongTensor | None = None, num_frames_per_video: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" pixel_values_videos (`torch.FloatTensor` of shape `(total_num_frames, max_patches, patch_pixels)`): The frames of every video in the batch, concatenated along the frame axis rather than stacked on a separate video axis, so that videos of different lengths can be batched together. video_position_ids (`torch.LongTensor` of shape `(total_num_frames, max_patches, 2)`, *optional*): 2D patch position coordinates from the video processor, with `(-1, -1)` indicating padding. Passed through to the vision encoder for positional embedding computation. num_frames_per_video (`torch.LongTensor` of shape `(num_videos,)`): Number of frames belonging to each video, used to split the flat frame sequence back per video. """ if self.vision_tower is None: raise ValueError( "Video features were requested, but the model was initialized without a vision_config. " "Cannot process video without a vision tower and vision embedder." ) if num_frames_per_video is None: raise ValueError( "`num_frames_per_video` is required when passing `pixel_values_videos`. The frames of all " "videos are concatenated along a single axis, so the per-video frame counts are the only " "way to split the result back per video. The video processor returns them." ) vision_outputs = self.vision_tower( pixel_values=pixel_values_videos, pixel_position_ids=video_position_ids, **kwargs, ) last_hidden_state = vision_outputs.last_hidden_state pooler_output = self.embed_vision(inputs_embeds=last_hidden_state) k_squared = self.config.vision_config.pooling_kernel_size**2 non_pad_mask = (video_position_ids != -1).all(dim=-1) # (total_num_frames, max_patches) num_tokens_per_frame = non_pad_mask.sum(dim=-1) // k_squared split_sizes = [ int(video_frames.sum()) for video_frames in torch.split(num_tokens_per_frame, num_frames_per_video.tolist()) ] vision_outputs.pooler_output = torch.split(pooler_output, split_sizes) return vision_outputs __all__ = ["EmbeddingGemma2Model", "EmbeddingGemma2PreTrainedModel", "EmbeddingGemma2TextModel"]