#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/gemma3n/modular_gemma3n.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_gemma3n.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2025 Google Inc. HuggingFace Inc. 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 Sequence
from typing import Any

from huggingface_hub.dataclasses import strict

from ...configuration_utils import PreTrainedConfig
from ...utils import auto_docstring, is_timm_available, logging, requires_backends


if is_timm_available():
    from timm.data import ImageNetInfo, infer_imagenet_subset

logger = logging.get_logger(__name__)


@auto_docstring(checkpoint="google/gemma-3n-E4B")
@strict
class Gemma3nTextConfig(PreTrainedConfig):
    r"""
    vocab_size_per_layer_input (`int`, *optional*, defaults to 262144):
        Vocabulary size of the per-layer text embeddings that augment the standard embeddings.
    hidden_size_per_layer_input (`int`, *optional*, defaults to 256):
        Dimension of the hidden representations for per-layer embeddings.
    altup_active_idx (`int`, *optional*, defaults to 0):
        The index of the prediction from which AltUp will compute additional predictions or correct the active prediction.
    altup_coef_clip (`float`, *optional*, defaults to 120.0):
        The maximum amplitude of an AltUp prediction or correction coefficient weight.
    altup_correct_scale (`bool`, *optional*, defaults to `True`):
        If True, apply the `AltUp.correct_output_scale` to the corrected prediction at `altup_active_idx`.
    altup_num_inputs (`int`, *optional*, defaults to 4):
        The number of predictions that AltUp should make given the input sequence.
    num_kv_shared_layers (`int`, *optional*, defaults to 15):
        The number of layers that share KV cache values. During the forward pass, the last `num_kv_shared_layers`
        layers in the model "share" the KV values in that each local and global layer in this range uses the KV
        cache values computed for the last local or global layer, respectively, before entering this range. The
        value should be a multiple of the attention pattern size (see `layer_types` parameter).
    laurel_rank (`int`, *optional*, defaults to 64):
        The intermediate size for the linear projections in the Learned Augmented Residual Layer.
    activation_sparsity_pattern (`Sequence[float]`, *optional*):
        The sparsity factor used to extract the top-k activations for a given layer. The provided Sequence must
        explicitly provide a sparsity value for each layer in the model. By default, the first 10 layers are
        sparse with a sparsity factor of 0.95 and the rest are dense.

    ```python
    >>> from transformers import Gemma3nTextModel, Gemma3nTextConfig

    >>> # Initializing a Gemma3nText gemma3n_text-E4B style configuration
    >>> configuration = Gemma3nTextConfig()

    >>> # Initializing a model from the gemma3n_text-E4B style configuration
    >>> model = Gemma3nTextModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    """

    model_type = "gemma3n_text"
    keys_to_ignore_at_inference = ["past_key_values"]
    base_model_tp_plan = {
        "layers.*.self_attn.q_proj": "colwise",
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
        "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
        "layers.*.self_attn.v_norm": "replicated_with_grad_allreduce",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }
    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }

    vocab_size: int = 262_400
    hidden_size: int = 2048
    intermediate_size: int | list[int] = 16_384
    num_hidden_layers: int = 35
    num_attention_heads: int = 8
    num_key_value_heads: int = 2
    head_dim: int = 256
    hidden_activation: str = "gelu_pytorch_tanh"
    max_position_embeddings: int = 32_768
    initializer_range: float = 0.02
    rms_norm_eps: float = 1e-6
    use_cache: bool = True
    pad_token_id: int | None = 0
    eos_token_id: int | list[int] | None = 1
    bos_token_id: int | None = 2
    tie_word_embeddings: bool = True
    rope_parameters: dict | None = None
    attention_bias: bool = False
    attention_dropout: int | float | None = 0.0
    sliding_window: int = 512
    layer_types: list[str] | None = None
    final_logit_softcapping: float = 30.0
    default_theta = {"global": 1_000_000.0, "local": 10_000.0}
    vocab_size_per_layer_input: int = 262_144
    hidden_size_per_layer_input: int = 256
    altup_active_idx: int = 0
    altup_coef_clip: float = 120.0
    altup_correct_scale: bool = True
    altup_num_inputs: int = 4
    num_kv_shared_layers: int = 15
    laurel_rank: int = 64
    activation_sparsity_pattern: float | list[float] | None = None

    def __post_init__(self, **kwargs):
        if (
            isinstance(self.intermediate_size, Sequence)
            and (intsize_len := len(self.intermediate_size)) != self.num_hidden_layers
        ):
            raise ValueError(
                "intermediate_size must have an explicit intermediate size for every layer or one for all layers. "
                f"Expected {self.num_hidden_layers} values but got {intsize_len}."
            )
        elif not isinstance(self.intermediate_size, Sequence):
            self.intermediate_size = [self.intermediate_size] * self.num_hidden_layers

        if self.layer_types is None:
            self.layer_types = [
                "full_attention" if (i + 1) % 5 == 0 else "sliding_attention" for i in range(self.num_hidden_layers)
            ]

        if self.activation_sparsity_pattern is None:
            num_sparse_layers = 10 if self.num_hidden_layers > 10 else 0
            self.activation_sparsity_pattern = [0.95] * num_sparse_layers + [0.0] * (
                self.num_hidden_layers - num_sparse_layers
            )

        if (len_asp := len(self.activation_sparsity_pattern)) != self.num_hidden_layers:
            raise ValueError(
                "activation_sparsity_pattern must have an explicit activation sparsity value for every layer."
                f"Expected {self.num_hidden_layers} values but got {len_asp}."
            )

        super().__post_init__(**kwargs)

    def validate_architecture(self):
        """Part of `@strict`-powered validation. Validates the architecture of the config."""
        if self.hidden_size % self.num_attention_heads != 0:
            raise ValueError(
                f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
                f"heads ({self.num_attention_heads})."
            )

    def convert_rope_params_to_dict(self, **kwargs):
        rope_scaling = kwargs.pop("rope_scaling", None)

        # Try to set `rope_scaling` if available, otherwise use `rope_parameters`. If we find `rope_parameters`
        # as arg in the inputs, we can safely assume that it is in the new format. New naming used -> new format
        default_rope_params = {
            "sliding_attention": {"rope_type": "default"},
            "full_attention": {"rope_type": "default"},
        }
        self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else default_rope_params
        if rope_scaling is not None:
            self.rope_parameters["full_attention"].update(rope_scaling)

        # Set default values if not present
        if self.rope_parameters.get("full_attention") is None:
            self.rope_parameters["full_attention"] = {"rope_type": "default"}
        self.rope_parameters["full_attention"].setdefault(
            "rope_theta", kwargs.pop("rope_theta", self.default_theta["global"])
        )
        if self.rope_parameters.get("sliding_attention") is None:
            self.rope_parameters["sliding_attention"] = {"rope_type": "default"}
        self.rope_parameters["sliding_attention"].setdefault(
            "rope_theta", kwargs.pop("rope_local_base_freq", self.default_theta["local"])
        )

        # Standardize and validate the correctness of rotary position embeddings parameters
        self.standardize_rope_params()
        return kwargs


@auto_docstring(checkpoint="google/gemma-3n-E4B")
@strict
class Gemma3nAudioConfig(PreTrainedConfig):
    r"""
    vocab_offset (`int`, *optional*, defaults to 262272):
        Offset between the tokenizer vocab index for the token ids embedded by `Gemma3nMultimodalEmbedder` and the
        0-indexed `Gemma3nMultimodalEmbedder.embedding` table.
    input_feat_size (`int`, *optional*, defaults to 128):
        The number of channels in each mel-spectrogram frame.
    gradient_clipping (`float`, *optional*, defaults to 10000000000.0):
        Clipping value used to stabilize extremely large gradient values.
    conf_attention_chunk_size (`int`, *optional*, defaults to 12):
        The sub-sequence size for local attention processing inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_attention_context_left (`int`, *optional*, defaults to 13):
        The left context size of the local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_attention_context_right (`int`, *optional*, defaults to 0):
        The right context size of the local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_attention_logit_cap (`float`, *optional*, defaults to 50.0):
        Logit cap applied during local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_num_attention_heads (`int`, *optional*, defaults to 8):
        The number of attention heads in local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_num_hidden_layers (`int`, *optional*, defaults to 12):
        The number of layers that use local attention inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_conv_kernel_size (`int`, *optional*, defaults to 5):
        Convolution kernel size for the conformer block inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_reduction_factor (`int`, *optional*, defaults to 4):
        Reduction factor used in the conformer block inside the Conformer ("conf") section of the
        Universal Speech Model.
    conf_residual_weight (`float`, *optional*, defaults to 0.5):
        Residual connection weight inside the Conformer ("conf") section of the
        Universal Speech Model.
    sscp_conv_channel_size (`tuple(int, int)`, *optional*, defaults to `(128, 32)`):
        The channel sizes for the first and second convolutional layers in the Sub-sample Convolution Projection
        ("sscp") section of the Universal Speech Model.
    sscp_conv_group_norm_eps (`float`, *optional*, defaults to 0.001):
        Epsilon used in group normalization in the subsample convolution projection in the Sub-sample Convolution
        Projection ("sscp") section of the Universal Speech Model.
    sscp_conv_kernel_size (`tuple(tuple(int, int), tuple(int, int))`, *optional*, defaults to `((3, 3), (3, 3))`):
        Kernel sizes of the two convolutional layers in the subsample convolution projection  in the Sub-sample
        Convolution Projection ("sscp") section of the Universal Speech Model. The kernel sizes are specified as a
        tuple of height and width for each layer, where the height corresponds to the time dimension and the width
        corresponds to the frequency dimension.
    sscp_conv_stride_size (`tuple(tuple(int, int), tuple(int, int))`, *optional*, defaults to `((2, 2), (2, 2))`):
        Stride sizes of the two convolutional layers in the subsample convolution projection in the Sub-sample
        Convolution Projection ("sscp") section of the Universal Speech Model. The stride sizes are specified as a
        tuple of height and width for each layer, where the height corresponds to the time dimension and the width
        corresponds to the frequency dimension.

    Example:

    ```python
    >>> from transformers import Gemma3nAudioConfig, Gemma3nAudioEncoder

    >>> # Initializing a Gemma3nAudioEncoder gemma3n_audio-E4B-style configuration
    >>> configuration = Gemma3nAudioConfig()

    >>> # Initializing a model from the gemma3n_audio-E4B style configuration
    >>> model = Gemma3nAudioEncoder(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    """

    model_type = "gemma3n_audio"

    vocab_size: int = 128
    vocab_offset: int = 262_144 + 128  # text vocab size + vision vocab size
    input_feat_size: int = 128
    hidden_size: int = 1536
    rms_norm_eps: float = 1e-6
    gradient_clipping: float = 10_000_000_000.0
    conf_attention_chunk_size: int = 12
    conf_attention_context_left: int = 13
    conf_attention_context_right: int = 0
    conf_attention_logit_cap: float = 50.0
    conf_num_attention_heads: int = 8
    conf_num_hidden_layers: int = 12
    conf_conv_kernel_size: int = 5
    conf_reduction_factor: int = 4
    conf_residual_weight: float = 0.5
    sscp_conv_channel_size: list[int] | tuple[int, int] = (128, 32)
    sscp_conv_group_norm_eps: float = 1e-3
    sscp_conv_kernel_size: list | tuple[tuple[int, int], tuple[int, int]] = (
        (3, 3),
        (3, 3),
    )
    sscp_conv_stride_size: list | tuple[tuple[int, int], tuple[int, int]] = (
        (2, 2),
        (2, 2),
    )


@auto_docstring(checkpoint="google/gemma-3n-E4B")
@strict
class Gemma3nVisionConfig(PreTrainedConfig):
    r"""
    architecture (`str`, *optional*, defaults to `"resnet50"`):
        The timm architecture to load.
    do_pooling (`bool`, *optional*, defaults to `True`):
        Whether to do pooling for the last_hidden_state in `TimmWrapperModel` or not.
    model_args (`dict[str, Any]`, *optional*):
        Additional keyword arguments to pass to the `timm.create_model` function. e.g. `model_args={"depth": 3}`
        for `timm/vit_base_patch32_clip_448.laion2b_ft_in12k_in1k` to create a model with 3 blocks. Defaults to `None`.
    vocab_offset (`int`, *optional*, defaults to 262144):
        Offset between the tokenizer vocab index for the token ids embedded by `Gemma3nMultimodalEmbedder` and the
        0-indexed `Gemma3nMultimodalEmbedder.embedding` table.

    Example:
    ```python
    >>> from transformers import Gemma3nVisionConfig, TimmWrapper

    >>> # Initializing a TimmWrapper gemma3n_vision-E4B-style configuration
    >>> configuration = Gemma3nVisionConfig()

    >>> # Initializing a gemma3n_vision-E4B-style TimmWrapper from the configuration
    >>> model = TimmWrapper(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```
    """

    model_type = "gemma3n_vision"
    architecture: str = "mobilenetv5_300m_enc"

    initializer_range: float = 0.02
    do_pooling: bool = False
    model_args: dict | None = None
    hidden_size: int = 2048
    vocab_size: int = 128
    vocab_offset: int = 262_144
    rms_norm_eps: float = 1e-06

    @classmethod
    def from_dict(cls, config_dict: dict[str, Any], **kwargs):
        # Create a copy to avoid mutating the original dict
        config_dict = config_dict.copy()

        label_names = config_dict.get("label_names")
        is_custom_model = "num_labels" in kwargs or "id2label" in kwargs

        # if no labels added to config, use imagenet labeller in timm
        if label_names is None and not is_custom_model:
            requires_backends(cls, ["timm"])
            imagenet_subset = infer_imagenet_subset(config_dict)
            if imagenet_subset:
                dataset_info = ImageNetInfo(imagenet_subset)
                synsets = dataset_info.label_names()
                label_descriptions = dataset_info.label_descriptions(as_dict=True)
                label_names = [label_descriptions[synset] for synset in synsets]

        if label_names is not None and not is_custom_model:
            kwargs["id2label"] = dict(enumerate(label_names))

            # if all label names are unique, create label2id mapping as well
            if len(set(label_names)) == len(label_names):
                kwargs["label2id"] = {name: i for i, name in enumerate(label_names)}
            else:
                kwargs["label2id"] = None

        # timm config stores the `num_classes` attribute in both the root of config and in the "pretrained_cfg" dict.
        # We are removing these attributes in order to have the native `transformers` num_labels attribute in config
        # and to avoid duplicate attributes
        num_labels_in_kwargs = kwargs.pop("num_labels", None)
        num_labels_in_dict = config_dict.pop("num_classes", None)

        # passed num_labels has priority over num_classes in config_dict
        kwargs["num_labels"] = num_labels_in_kwargs or num_labels_in_dict

        # pop num_classes from "pretrained_cfg",
        # it is not necessary to have it, only root one is used in timm
        if "pretrained_cfg" in config_dict and "num_classes" in config_dict["pretrained_cfg"]:
            config_dict["pretrained_cfg"].pop("num_classes", None)

        return super().from_dict(config_dict, **kwargs)

    def to_dict(self) -> dict[str, Any]:
        output = super().to_dict()
        output.setdefault("num_classes", self.num_labels)
        output.setdefault("label_names", list(self.id2label.values()))
        output.pop("id2label", None)
        output.pop("label2id", None)
        return output


@auto_docstring(checkpoint="google/gemma-3n-E4B")
@strict
class Gemma3nConfig(PreTrainedConfig):
    r"""
    audio_soft_tokens_per_image (`int`, *optional*, defaults to 188):
        The number of soft tokens per audio clip.
    vision_soft_tokens_per_image (`int`, *optional*, defaults to 256):
        The number of soft tokens per image.
    boi_token_id (`int`, *optional*, defaults to 255999):
        The begin-of-image token index to wrap the image prompt.
    eoi_token_id (`int`, *optional*, defaults to 262144):
        The end-of-image token index to wrap the image prompt.
    boa_token_id (`int`, *optional*, defaults to 256000):
        The begin-of-audio token index to wrap the audio prompt.
    eoa_token_id (`int`, *optional*, defaults to 262272):
        The end-of-audio token index to wrap the audio prompt.

    Example:

    ```python
    >>> from transformers import Gemma3nForConditionalGeneration, Gemma3nConfig, Gemma3nTextConfig

    >>> # Initializing a MobileNet vision config, which is loaded from TIMM
    >>> vision_config = Gemma3nVisionConfig()

    >>> # Initializing a Gemma3n Audio config
    >>> audio_config = Gemma3nAudioConfig()

    >>> # Initializing a Gemma3n Text config
    >>> text_config = Gemma3nTextConfig()

    >>> # Initializing a Gemma3n gemma-3-4b style configuration
    >>> configuration = Gemma3nConfig(text_config, vision_config, audio_config)

    >>> # Initializing a model from the gemma-3-4b style configuration
    >>> model = Gemma3nTextConfig(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```"""

    model_type = "gemma3n"
    sub_configs = {
        "text_config": Gemma3nTextConfig,
        "vision_config": Gemma3nVisionConfig,
        "audio_config": Gemma3nAudioConfig,
    }

    text_config: Gemma3nTextConfig | dict[str, Any] | None = None
    vision_config: Gemma3nVisionConfig | dict[str, Any] | None = None
    audio_config: Gemma3nAudioConfig | dict[str, Any] | None = None
    audio_soft_tokens_per_image: int | None = 188
    vision_soft_tokens_per_image: int | None = 256
    boi_token_id: int | None = 255_999
    eoi_token_id: int | None = 262_144
    image_token_id: int | None = 262_145
    boa_token_id: int | None = 256_000
    eoa_token_id: int | None = 262_272
    audio_token_id: int | None = 262_273
    initializer_range: float | None = 0.02
    tie_word_embeddings: bool | None = True

    def __post_init__(self, **kwargs):
        if self.text_config is None:
            self.text_config = Gemma3nTextConfig()
            logger.info("text_config is None, using default Gemma3nTextConfig text config.")
        elif isinstance(self.text_config, dict):
            self.text_config = Gemma3nTextConfig(**self.text_config)

        if isinstance(self.vision_config, dict):
            self.vision_config = Gemma3nVisionConfig(**self.vision_config)
        elif self.vision_config is None:
            self.vision_config = Gemma3nVisionConfig()
            logger.info("vision_config is None, using default Gemma3nVisionConfig vision config.")

        if isinstance(self.audio_config, dict):
            self.audio_config = Gemma3nAudioConfig(**self.audio_config)
        elif self.audio_config is None:
            self.audio_config = Gemma3nAudioConfig()
            logger.info("audio_config is None. Using default Gemma3nAudioConfig.")

        super().__post_init__(**kwargs)


__all__ = ["Gemma3nAudioConfig", "Gemma3nConfig", "Gemma3nTextConfig", "Gemma3nVisionConfig"]
