Upload 15 files

config.json

@@ -0,0 +1,221 @@
+{
+  "_name_or_path": "C:\\Users\\teres\\source\\repos\\ModelFineTune\\phi4-model",
+  "architectures": [
+    "Phi4MMForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "audio_processor": {
+    "config": {
+      "activation": "swish",
+      "activation_checkpointing": {
+        "interval": 1,
+        "module": "transformer",
+        "offload": false
+      },
+      "attention_dim": 1024,
+      "attention_heads": 16,
+      "batch_norm": false,
+      "bias_in_glu": true,
+      "causal": true,
+      "chunk_size": -1,
+      "cnn_layer_norm": true,
+      "conv_activation": "swish",
+      "conv_glu_type": "swish",
+      "depthwise_multiplier": 1,
+      "depthwise_seperable_out_channel": 1024,
+      "dropout_rate": 0.0,
+      "encoder_embedding_config": {
+        "input_size": 80
+      },
+      "ext_pw_kernel_size": 1,
+      "ext_pw_out_channel": 1024,
+      "input_layer": "nemo_conv",
+      "input_size": 80,
+      "kernel_size": 3,
+      "left_chunk": 18,
+      "linear_units": 1536,
+      "nemo_conv_settings": {
+        "conv_channels": 1024
+      },
+      "num_blocks": 24,
+      "relative_attention_bias_args": {
+        "t5_bias_max_distance": 500,
+        "type": "t5"
+      },
+      "time_reduction": 8
+    },
+    "name": "cascades"
+  },
+  "auto_map": {
+    "AutoConfig": "configuration_phi4mm.Phi4MMConfig",
+    "AutoModelForCausalLM": "modeling_phi4mm.Phi4MMForCausalLM",
+    "AutoTokenizer": "Xenova/gpt-4o"
+  },
+  "bos_token_id": 199999,
+  "embd_layer": {
+    "audio_embd_layer": {
+      "compression_rate": 8,
+      "downsample_rate": 1,
+      "embedding_cls": "audio",
+      "enable_gradient_checkpointing": true,
+      "projection_cls": "mlp",
+      "use_conv_downsample": false,
+      "use_qformer": false
+    },
+    "embedding_cls": "image_audio",
+    "image_embd_layer": {
+      "crop_size": 448,
+      "embedding_cls": "tune_image",
+      "enable_gradient_checkpointing": true,
+      "hd_transform_order": "sub_glb",
+      "image_token_compression_cls": "avg_pool_2d",
+      "projection_cls": "mlp",
+      "use_hd_transform": true,
+      "with_learnable_separator": true
+    }
+  },
+  "embd_pdrop": 0.0,
+  "eos_token_id": 199999,
+  "full_attn_mod": 1,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "img_processor": null,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "interpolate_factor": 1,
+  "lm_head_bias": false,
+  "max_position_embeddings": 131072,
+  "mlp_bias": false,
+  "model_type": "phi4mm",
+  "num_attention_heads": 24,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 8,
+  "original_max_position_embeddings": 4096,
+  "pad_token_id": 199999,
+  "partial_rotary_factor": 0.75,
+  "resid_pdrop": 0.0,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "long_factor": [
+      1,
+      1.118320672,
+      1.250641126,
+      1.398617824,
+      1.564103225,
+      1.74916897,
+      1.956131817,
+      2.187582649,
+      2.446418898,
+      2.735880826,
+      3.059592084,
+      3.421605075,
+      3.826451687,
+      4.279200023,
+      4.785517845,
+      5.351743533,
+      5.984965424,
+      6.693110555,
+      7.485043894,
+      8.370679318,
+      9.36110372,
+      10.4687158,
+      11.70738129,
+      13.09260651,
+      14.64173252,
+      16.37415215,
+      18.31155283,
+      20.47818807,
+      22.90118105,
+      25.61086418,
+      28.64115884,
+      32.03,
+      32.1,
+      32.13,
+      32.23,
+      32.6,
+      32.61,
+      32.64,
+      32.66,
+      32.7,
+      32.71,
+      32.93,
+      32.97,
+      33.28,
+      33.49,
+      33.5,
+      44.16,
+      47.77
+    ],
+    "short_factor": [
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0,
+      1.0
+    ],
+    "type": "longrope"
+  },
+  "rope_theta": 10000.0,
+  "sliding_window": 262144,
+  "speech_lora": {
+    "dp": 0.01,
+    "layer": "((layers.*self_attn\\.(qkv|o)_proj)|(layers.*mlp\\.(gate_up|down)_proj))",
+    "lora_alpha": 640,
+    "r": 320
+  },
+  "tie_word_embeddings": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.1",
+  "use_cache": true,
+  "vision_lora": {
+    "dp": 0.0,
+    "layer": "layers.*((self_attn\\.(qkv_proj|o_proj))|(mlp\\.(gate_up|down)_proj))",
+    "lora_alpha": 512,
+    "r": 256
+  },
+  "vocab_size": 200064
+}

configuration_phi4mm.py

@@ -0,0 +1,235 @@
+# coding=utf-8
+# Copyright 2024 Microsoft and the 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.
+
+""" Phi-4-MM model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Phi4MMConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Phi4MMModel`]. It is used to instantiate a Phi-4-MM
+    model according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 200064):
+            Vocabulary size of the Phi-4-MM model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Phi4MMModel`].
+        hidden_size (`int`, *optional*, defaults to 3072):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 8192):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        resid_pdrop (`float`, *optional*, defaults to 0.0):
+            Dropout probability for mlp outputs.
+        embd_pdrop (`int`, *optional*, defaults to 0.0):
+            The dropout ratio for the embeddings.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio after computing the attention scores.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 4096):
+            The maximum sequence length that this model might ever be used with.
+        original_max_position_embeddings (`int`, *optional*, defaults to 4096):
+            The maximum sequence length that this model was trained with. This is used to determine the size of the
+            original RoPE embeddings when using long scaling.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon value used for the RMSNorm.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`. Whether to tie weight embeddings or not.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`dict`, *optional*):
+            The scaling strategy for the RoPE embeddings. If `None`, no scaling is applied. If a dictionary, it must
+            contain the following keys: `type`, `short_factor` and `long_factor`. The `type` must be `longrope` and
+            the `short_factor` and `long_factor` must be lists of numbers with the same length as the hidden size
+            divided by the number of attention heads divided by 2.
+        partial_rotary_factor (`float`, *optional*, defaults to 0.5):
+            Percentage of the query and keys which will have rotary embedding.
+        bos_token_id (`int`, *optional*, defaults to 199999):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 199999):
+            The id of the "end-of-sequence" token.
+        pad_token_id (`int`, *optional*, defaults to 199999):
+            The id of the padding token.
+        sliding_window (`int`, *optional*):
+            Sliding window attention window size. If `None`, no sliding window is applied.
+
+    Example:
+
+    ```python
+    >>> from transformers import Phi4MMModel, Phi4MMConfig
+
+    >>> # Initializing a Phi-4-MM style configuration
+    >>> configuration = Phi4MMConfig.from_pretrained("TBA")
+
+    >>> # Initializing a model from the configuration
+    >>> model = Phi4MMModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "phi4mm"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=200064,
+        hidden_size=3072,
+        intermediate_size=8192,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        resid_pdrop=0.0,
+        embd_pdrop=0.0,
+        attention_dropout=0.0,
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        original_max_position_embeddings=4096,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        partial_rotary_factor=1,
+        bos_token_id=199999,
+        eos_token_id=199999,
+        pad_token_id=199999,
+        sliding_window=None,
+        embd_layer: str = "default",
+        img_processor=None,
+        audio_processor=None,
+        vision_lora=None,
+        speech_lora=None,
+        **kwargs,
+    ):
+        self.embd_layer = embd_layer
+        self.img_processor = img_processor
+        self.audio_processor = audio_processor
+        self.vision_lora = vision_lora
+        self.speech_lora = speech_lora
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attention_dropout = attention_dropout
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.partial_rotary_factor = partial_rotary_factor
+        self._rope_scaling_adjustment()
+        self._rope_scaling_validation()
+        self.sliding_window = sliding_window
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def _rope_scaling_adjustment(self):
+        """
+        Adjust the `type` of the `rope_scaling` configuration for backward compatibility.
+        """
+        if self.rope_scaling is None:
+            return
+
+        rope_scaling_type = self.rope_scaling.get("type", None)
+
+        # For backward compatibility if previous version used "su" or "yarn"
+        if rope_scaling_type is not None and rope_scaling_type in ["su", "yarn"]:
+            self.rope_scaling["type"] = "longrope"
+
+    def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
+        if self.rope_scaling is None:
+            return
+
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 3:
+            raise ValueError(
+                "`rope_scaling` must be a dictionary with three fields, `type`, `short_factor` and `long_factor`, "
+                f"got {self.rope_scaling}"
+            )
+        rope_scaling_type = self.rope_scaling.get("type", None)
+        rope_scaling_short_factor = self.rope_scaling.get("short_factor", None)
+        rope_scaling_long_factor = self.rope_scaling.get("long_factor", None)
+        if rope_scaling_type is None or rope_scaling_type not in ["longrope"]:
+            raise ValueError(f"`rope_scaling`'s type field must be one of ['longrope'], got {rope_scaling_type}")
+        if not (
+            isinstance(rope_scaling_short_factor, list)
+            and all(isinstance(x, (int, float)) for x in rope_scaling_short_factor)
+        ):
+            raise ValueError(
+                f"`rope_scaling`'s short_factor field must be a list of numbers, got {rope_scaling_short_factor}"
+            )
+        rotary_ndims = int(self.hidden_size // self.num_attention_heads * self.partial_rotary_factor)
+        if not len(rope_scaling_short_factor) == rotary_ndims // 2:
+            raise ValueError(
+                f"`rope_scaling`'s short_factor field must have length {rotary_ndims // 2}, got {len(rope_scaling_short_factor)}"
+            )
+        if not (
+            isinstance(rope_scaling_long_factor, list)
+            and all(isinstance(x, (int, float)) for x in rope_scaling_long_factor)
+        ):
+            raise ValueError(
+                f"`rope_scaling`'s long_factor field must be a list of numbers, got {rope_scaling_long_factor}"
+            )
+        if not len(rope_scaling_long_factor) == rotary_ndims // 2:
+            raise ValueError(
+                f"`rope_scaling`'s long_factor field must have length {rotary_ndims // 2}, got {len(rope_scaling_long_factor)}"
+            )

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 199999,
+  "eos_token_id": [
+    200020,
+    199999
+  ],
+  "pad_token_id": 199999,
+  "transformers_version": "4.46.1"
+}

model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b32e9f73edc116823abe74dd6b7f13b5bdd7a5262e7710835772697d988a2aa
+size 4993293944

model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:293a733b7d151acf7b3803632124a67a532145f5a7386ba6f99717e3cdabbc8a
+size 4953762936

model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:502a0e423ba01769c3e292512706f14d21b3d1b91d9fc7092765fa53d7ba6794
+size 4936985712

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9716ae50890580ed9f6b6d77508ec0e0b5bf843ada3ede1c5dd7a7141bf2cea
+size 3702770248

processing_phi4mm.py

@@ -0,0 +1,733 @@
+# Copyright 2024 Microsoft and the 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.
+
+"""
+Processor class for Phi4MM
+"""
+import re
+from typing import List, Optional, Tuple, Union
+import math
+from enum import Enum
+
+import numpy as np
+import scipy
+import torch
+import torchvision
+
+from transformers import AutoFeatureExtractor, AutoImageProcessor
+from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.image_utils import (
+    ImageInput,
+    make_list_of_images,
+    valid_images,
+)
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils_base import PaddingStrategy, TextInput, TruncationStrategy
+from transformers.utils import TensorType, logging
+from torch.nn.utils.rnn import pad_sequence
+
+
+logger = logging.get_logger(__name__)
+
+# Special tokens
+_COMPATIBLE_IMAGE_SPECIAL_TOKEN_PATTERN = r'<\|image_\d+\|>'  # For backward compatibility
+_COMPATIBLE_AUDIO_SPECIAL_TOKEN_PATTERN = r'<\|audio_\d+\|>'  # For backward compatibility
+_IMAGE_SPECIAL_TOKEN = '<|endoftext10|>'
+_AUDIO_SPECIAL_TOKEN = '<|endoftext11|>'
+_IMAGE_SPECIAL_TOKEN_ID = 200010  # '<|endoftext10|>', or we can better name it (in `tokenizer_config.json`)
+_AUDIO_SPECIAL_TOKEN_ID = 200011  # '<|endoftext11|>'
+
+
+class InputMode(Enum):
+    LANGUAGE = 0
+    VISION = 1
+    SPEECH = 2
+    VISION_SPEECH = 3
+
+
+class Phi4MMImageProcessor(BaseImageProcessor):
+    r"""
+    Constructs a Phi4MM image processor.
+    """
+    model_input_names = ["input_image_embeds", "image_sizes", "image_attention_mask"]
+
+    def __init__(
+        self,
+        dynamic_hd,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.dynamic_hd = dynamic_hd
+
+    def find_closest_aspect_ratio(self, aspect_ratio, target_ratios, width, height, image_size):
+        best_ratio_diff = float('inf')
+        best_ratio = (1, 1)
+        area = width * height
+        for ratio in target_ratios:
+            target_aspect_ratio = ratio[0] / ratio[1]
+            ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+            if ratio_diff < best_ratio_diff:
+                best_ratio_diff = ratio_diff
+                best_ratio = ratio
+            elif ratio_diff == best_ratio_diff:
+                if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                    best_ratio = ratio
+        return best_ratio
+
+    def dynamic_preprocess(self, image, min_num=1, max_num=12, image_size=384, mask_size=27, use_thumbnail=True):
+        orig_width, orig_height = image.size
+
+        w_crop_num = math.ceil(orig_width/float(image_size))
+        h_crop_num = math.ceil(orig_height/float(image_size))
+        if w_crop_num * h_crop_num > max_num:
+
+            aspect_ratio = orig_width / orig_height
+
+            # calculate the existing image aspect ratio
+            target_ratios = set(
+                (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
+                i * j <= max_num and i * j >= min_num)
+            target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+            # find the closest aspect ratio to the target
+            target_aspect_ratio = self.find_closest_aspect_ratio(
+                aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+            # calculate the target width and height
+            target_width = image_size * target_aspect_ratio[0]
+            target_height = image_size * target_aspect_ratio[1]
+        else:
+            target_width = image_size * w_crop_num
+            target_height = image_size * h_crop_num
+            target_aspect_ratio = (w_crop_num, h_crop_num)
+
+        # Calculate the ratio
+        ratio_width = target_width / orig_width
+        ratio_height = target_height / orig_height
+        if ratio_width < ratio_height:
+            new_size = (target_width, int(orig_height * ratio_width))
+            padding_width = 0
+            padding_height = target_height - int(orig_height * ratio_width)
+        else:
+            new_size = (int(orig_width * ratio_height), target_height)
+            padding_width = target_width - int(orig_width * ratio_height)
+            padding_height = 0
+
+        attention_mask = torch.ones((int(mask_size*target_aspect_ratio[1]), int(mask_size*target_aspect_ratio[0])))
+        if padding_width >= 14:
+            attention_mask[:, -math.floor(padding_width/14):] = 0
+        if padding_height >= 14:
+            attention_mask[-math.floor(padding_height/14):,:] = 0
+        assert attention_mask.sum() > 0
+
+        if min(new_size[1], target_height) < 10 or min(new_size[0], target_width) < 10:
+            raise ValueError(f'the aspect ratio is very extreme {new_size}')
+
+        image = torchvision.transforms.functional.resize(image, [new_size[1], new_size[0]],)
+
+        resized_img = torchvision.transforms.functional.pad(image, [0, 0, padding_width, padding_height], fill=[255,255,255])
+
+        return resized_img, attention_mask
+
+    def pad_to_max_num_crops(self, images, max_crops=5):
+        """
+        images: B x 3 x H x W, B<=max_crops
+        """
+        B, _, H, W = images.shape
+        if B < max_crops:
+            pad = torch.zeros(max_crops - B, 3, H, W, dtype=images.dtype, device=images.device)
+            images = torch.cat([images, pad], dim=0)
+        return images
+
+    def pad_mask_to_max_num_crops(self, masks, max_crops=5):
+        B, H, W = masks.shape
+        if B < max_crops:
+            pad = torch.ones(max_crops - B, H, W, dtype=masks.dtype, device=masks.device)
+            masks = torch.cat([masks, pad], dim=0)
+        return masks
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+    ):
+        """
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+        """
+        images = make_list_of_images(images)
+
+        if not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+
+        # Basic settings.
+        img_processor = torchvision.transforms.Compose([
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize(
+                (0.5, 0.5, 0.5),
+                (0.5, 0.5, 0.5)
+            ),
+        ])
+        dyhd_base_resolution = 448
+
+        # Dynamic HD
+        base_resolution = dyhd_base_resolution
+        images = [image.convert('RGB') for image in images]
+        # cover 384 and 448 resolution
+        mask_resolution = base_resolution // 14
+        elems, image_attention_masks = [], []
+        for im in images:
+            elem, attention_mask = self.dynamic_preprocess(im, max_num=self.dynamic_hd, image_size=base_resolution, mask_size=mask_resolution)
+            elems.append(elem)
+            image_attention_masks.append(attention_mask)
+        hd_images = [img_processor(im) for im in elems]
+        global_image = [torch.nn.functional.interpolate(im.unsqueeze(0).float(), size=(base_resolution, base_resolution), mode='bicubic',).to(im.dtype) for im in hd_images]
+        shapes = [[im.size(1), im.size(2)] for im in hd_images]
+        mask_shapes = [[mask.size(0), mask.size(1)] for mask in image_attention_masks]
+        global_attention_mask = [torch.ones((1, mask_resolution, mask_resolution)) for _ in hd_images]
+        hd_images_reshape = [im.reshape(1, 3,
+                                            h//base_resolution,
+                                            base_resolution,
+                                            w//base_resolution,
+                                            base_resolution
+                                            ).permute(0,2,4,1,3,5).reshape(-1, 3, base_resolution, base_resolution).contiguous() for im, (h, w) in zip(hd_images, shapes)]
+        attention_masks_reshape = [mask.reshape(1,
+                                            h//mask_resolution,
+                                            mask_resolution,
+                                            w//mask_resolution,
+                                            mask_resolution
+                                            ).permute(0,1,3,2,4).reshape(-1, mask_resolution, mask_resolution).contiguous() for mask, (h, w) in zip(image_attention_masks, mask_shapes)]
+        downsample_attention_masks = [mask[:,0::2,0::2].reshape(1,
+                                            h//mask_resolution,
+                                            w//mask_resolution,
+                                            mask_resolution//2+mask_resolution%2,
+                                            mask_resolution//2+mask_resolution%2
+                                            ).permute(0,1,3,2,4) for mask, (h,w) in zip(attention_masks_reshape, mask_shapes)]
+        downsample_attention_masks = [mask.reshape(mask.size(1)*mask.size(2), mask.size(3)*mask.size(4))for mask in downsample_attention_masks]
+        num_img_tokens = [256 + 1 + int(mask.sum().item()) + int(mask[:,0].sum().item()) + 16 for mask in downsample_attention_masks]
+
+        hd_images_reshape = [torch.cat([_global_image] + [_im], dim=0) for _global_image, _im in zip(global_image, hd_images_reshape)]
+        hd_masks_reshape = [torch.cat([_global_mask] + [_mask], dim=0) for _global_mask, _mask in zip(global_attention_mask, attention_masks_reshape)]
+        max_crops = max([img.size(0) for img in hd_images_reshape])
+        image_transformed = [self.pad_to_max_num_crops(im, max_crops) for im in hd_images_reshape]
+        image_transformed = torch.stack(image_transformed, dim=0)
+        mask_transformed = [self.pad_mask_to_max_num_crops(mask, max_crops) for mask in hd_masks_reshape]
+        mask_transformed = torch.stack(mask_transformed, dim=0)
+
+        returned_input_image_embeds = image_transformed
+        returned_image_sizes = torch.tensor(shapes, dtype=torch.long)
+        returned_image_attention_mask = mask_transformed
+        returned_num_img_tokens = num_img_tokens
+
+        data = {
+            "input_image_embeds": returned_input_image_embeds,
+            "image_sizes": returned_image_sizes,
+            "image_attention_mask": returned_image_attention_mask,
+            "num_img_tokens": returned_num_img_tokens,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+AudioInput = Tuple[Union[np.ndarray, torch.Tensor], int]
+AudioInputs = List[AudioInput]
+
+
+def speechlib_mel(sample_rate, n_fft, n_mels, fmin=None, fmax=None):
+    """Create a Mel filter-bank the same as SpeechLib FbankFC.
+
+    Args:
+        sample_rate (int): Sample rate in Hz. number > 0 [scalar]
+        n_fft (int): FFT size. int > 0 [scalar]
+        n_mel (int): Mel filter size. int > 0 [scalar]
+        fmin (float): lowest frequency (in Hz). If None use 0.0.
+            float >= 0 [scalar]
+        fmax: highest frequency (in Hz). If None use sample_rate / 2.
+            float >= 0 [scalar]
+
+    Returns
+        out (numpy.ndarray): Mel transform matrix
+            [shape=(n_mels, 1 + n_fft/2)]
+    """
+
+    bank_width = int(n_fft // 2 + 1)
+    if fmax is None:
+        fmax = sample_rate / 2
+    if fmin is None:
+        fmin = 0
+    assert fmin >= 0, "fmin cannot be negtive"
+    assert fmin < fmax <= sample_rate / 2, "fmax must be between (fmin, samplerate / 2]"
+
+    def mel(f):
+        return 1127.0 * np.log(1.0 + f / 700.0)
+
+    def bin2mel(fft_bin):
+        return 1127.0 * np.log(1.0 + fft_bin * sample_rate / (n_fft * 700.0))
+
+    def f2bin(f):
+        return int((f * n_fft / sample_rate) + 0.5)
+
+    # Spec 1: FFT bin range [f2bin(fmin) + 1, f2bin(fmax) - 1]
+    klo = f2bin(fmin) + 1
+    khi = f2bin(fmax)
+
+    khi = max(khi, klo)
+
+    # Spec 2: SpeechLib uses trianges in Mel space
+    mlo = mel(fmin)
+    mhi = mel(fmax)
+    m_centers = np.linspace(mlo, mhi, n_mels + 2)
+    ms = (mhi - mlo) / (n_mels + 1)
+
+    matrix = np.zeros((n_mels, bank_width), dtype=np.float32)
+    for m in range(0, n_mels):
+        left = m_centers[m]
+        center = m_centers[m + 1]
+        right = m_centers[m + 2]
+        for fft_bin in range(klo, khi):
+            mbin = bin2mel(fft_bin)
+            if left < mbin < right:
+                matrix[m, fft_bin] = 1.0 - abs(center - mbin) / ms
+
+    return matrix
+
+
+class Phi4MMAudioFeatureExtractor(SequenceFeatureExtractor):
+    model_input_names = ["input_audio_embeds", "audio_embed_sizes", "audio_attention_mask"]
+
+    def __init__(self, audio_compression_rate, audio_downsample_rate, audio_feat_stride, **kwargs):
+        feature_size = 80
+        sampling_rate = 16000
+        padding_value = 0.0
+        super().__init__(feature_size, sampling_rate, padding_value, **kwargs)
+
+        self.compression_rate = audio_compression_rate
+        self.qformer_compression_rate = audio_downsample_rate
+        self.feat_stride = audio_feat_stride
+
+        self._eightk_method = "fillzero"
+        self._mel = speechlib_mel(16000, 512, 80, fmin=None, fmax=7690).T
+
+        self._hamming400 = np.hamming(400)  # for 16k audio
+        self._hamming200 = np.hamming(200)  # for 8k audio
+
+    def duration_to_frames(self, duration):
+        """duration in s, estimated frames"""
+        frame_rate = 10
+
+        num_frames = duration * 1000 // frame_rate
+        return num_frames
+
+    def __call__(
+        self,
+        audios: List[AudioInput],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+    ):
+        # Ref: https://github.com/huggingface/transformers/blob/v4.47.0/src/transformers/models/audio_spectrogram_transformer/feature_extraction_audio_spectrogram_transformer.py#L161
+        returned_input_audio_embeds = []
+        returned_audio_embed_sizes = []
+        audio_frames_list = []
+
+        for audio_data, sample_rate in audios:
+            audio_embeds = self._extract_features(audio_data, sample_rate)
+            audio_frames = len(audio_embeds) * self.feat_stride
+            audio_embed_size = self._compute_audio_embed_size(audio_frames)
+
+            returned_input_audio_embeds.append(torch.tensor(audio_embeds))
+            returned_audio_embed_sizes.append(torch.tensor(audio_embed_size).long())
+            audio_frames_list.append(audio_frames)
+
+        returned_input_audio_embeds = pad_sequence(
+            returned_input_audio_embeds, batch_first=True
+        )
+        returned_audio_embed_sizes = torch.stack(returned_audio_embed_sizes, dim=0)
+        audio_frames = torch.tensor(audio_frames_list)
+        returned_audio_attention_mask = torch.arange(0, audio_frames.max()).unsqueeze(0) < audio_frames.unsqueeze(1) if len(audios) > 1 else None
+
+        data = {
+            "input_audio_embeds": returned_input_audio_embeds,
+            "audio_embed_sizes": returned_audio_embed_sizes,
+        }
+        if returned_audio_attention_mask is not None:
+            data["audio_attention_mask"] = returned_audio_attention_mask
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def _extract_spectrogram(self, wav, fs):
+        """Extract spectrogram features from waveform.
+        Args:
+            wav (1D array): waveform of the input
+            fs (int): sampling rate of the waveform, 16000 or 8000.
+                If fs=8000, the waveform will be resampled to 16000Hz.
+        Output:
+            log_fbank (2D array): a TxD matrix of log Mel filterbank features.
+                D=80, and T is the number of frames.
+        """
+        if wav.ndim > 1:
+            wav = np.squeeze(wav)
+
+        # by default, we extract the mean if stereo
+        if len(wav.shape) == 2:
+            wav = wav.mean(1)
+
+        # Resample to 16000 or 8000 if needed
+        if fs > 16000:
+            wav = scipy.signal.resample_poly(wav, 1, fs // 16000)
+            fs = 16000
+        elif 8000 < fs < 16000:
+            wav = scipy.signal.resample_poly(wav, 1, fs // 8000)
+            fs = 8000
+        elif fs < 8000:
+            raise RuntimeError(f"Unsupported sample rate {fs}")
+
+        if fs == 8000:
+            if self._eightk_method == "resample":
+                # Input audio is 8 kHz. Convert to 16 kHz before feature
+                # extraction
+                wav = scipy.signal.resample_poly(wav, 2, 1)
+                fs = 16000
+            # Do nothing here for fillzero method
+        elif fs != 16000:
+            # Input audio is not a supported sample rate.
+            raise RuntimeError(f"Input data using an unsupported sample rate: {fs}")
+
+        preemphasis = 0.97
+
+        if fs == 8000:
+            n_fft = 256
+            win_length = 200
+            hop_length = 80
+            fft_window = self._hamming200
+        elif fs == 16000:
+            n_fft = 512
+            win_length = 400
+            hop_length = 160
+            fft_window = self._hamming400
+
+        # Spec 1: SpeechLib cut remaining sample insufficient for a hop
+        n_batch = (wav.shape[0] - win_length) // hop_length + 1
+        # Here we don't use stride_tricks since the input array may not satisfy
+        # memory layout requirement and we need writeable output
+        # Here we only use list of views before copy to desination
+        # so it is more efficient than broadcasting
+        y_frames = np.array(
+            [wav[_stride : _stride + win_length] for _stride in range(0, hop_length * n_batch, hop_length)],
+            dtype=np.float32,
+        )
+
+        # Spec 2: SpeechLib applies preemphasis within each batch
+        y_frames_prev = np.roll(y_frames, 1, axis=1)
+        y_frames_prev[:, 0] = y_frames_prev[:, 1]
+        y_frames = (y_frames - preemphasis * y_frames_prev) * 32768
+
+        S = np.fft.rfft(fft_window * y_frames, n=n_fft, axis=1).astype(np.complex64)
+
+        if fs == 8000:
+            # Need to pad the output to look like 16 kHz data but with zeros in
+            # the 4 to 8 kHz bins.
+            frames, bins = S.shape
+            padarray = np.zeros((frames, bins))
+            S = np.concatenate((S[:, 0:-1], padarray), axis=1)  # Nyquist bin gets set to zero
+
+        spec = np.abs(S).astype(np.float32)
+        return spec
+
+    def _extract_features(self, wav, fs):
+        """Extract log filterbank features from waveform.
+        Args:
+            wav (1D array): waveform of the input
+            fs (int): sampling rate of the waveform, 16000 or 8000.
+                If fs=8000, the waveform will be resampled to 16000Hz.
+        Output:
+            log_fbank (2D array): a TxD matrix of log Mel filterbank features.
+                D=80, and T is the number of frames.
+        """
+        spec = self._extract_spectrogram(wav, fs)
+        spec_power = spec**2
+
+        fbank_power = np.clip(spec_power.dot(self._mel), 1.0, None)
+        log_fbank = np.log(fbank_power).astype(np.float32)
+
+        return log_fbank
+
+    def _compute_audio_embed_size(self, audio_frames):
+        integer = audio_frames // self.compression_rate
+        remainder = audio_frames % self.compression_rate
+
+        result = integer if remainder == 0 else integer + 1
+
+        integer = result // self.qformer_compression_rate
+        remainder = result % self.qformer_compression_rate
+        result = integer if remainder == 0 else integer + 1  # qformer compression
+
+        return result
+
+
+class Phi4MMProcessor(ProcessorMixin):
+    r"""
+    Constructs a Phi4MM processor which raps an image processor, a audio processor, and a GPT tokenizer into a single processor.
+
+    [`Phi4MMProcessor`] offers all the functionalities of [`Phi4MMImageProcessor`] and [`GPT2Tokenizer`]. See the
+    [`~Phi4MMProcessor.__call__`] and [`~Phi4MMProcessor.decode`] for more information.
+
+    Args:
+        image_processor ([`Phi4MMImageProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`GPT2Tokenizer`], *optional*):
+            The tokenizer is a required input.
+    """
+
+    attributes = ["image_processor", "audio_processor", "tokenizer"]
+    tokenizer_class = "GPT2TokenizerFast"
+    image_processor_class = "AutoImageProcessor"  # Phi4MMImageProcessor will be registered later
+    audio_processor_class = "AutoFeatureExtractor"  # Phi4MMAudioFeatureExtractor will be registered later
+
+    def __init__(self, image_processor, audio_processor, tokenizer):
+        self.image_processor = image_processor
+        self.audio_processor = audio_processor
+        self.tokenizer = tokenizer
+
+    def __call__(
+        self,
+        text: Union[TextInput, List[TextInput]],
+        images: Optional[ImageInput] = None,
+        audios: Optional[AudioInputs] = None,
+        padding: Union[bool, str, PaddingStrategy] = False,
+        truncation: Optional[Union[bool, str, TruncationStrategy]] = None,
+        max_length=None,
+        return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
+    ) -> BatchFeature:
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forards the `text`
+        and `kwargs` arguments to GPT2Tokenizer's [`~GPT2Tokenizer.__call__`] if `text` is not `None` to encode
+        the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
+        Phi4MMImageProcessor's [`~Phi4MMImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
+        of the above two methods for more information.
+
+        Args:
+            text (`str`, `List[str]`, `List[List[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
+                Select a strategy to pad the returned sequences (according to the model's padding side and padding
+                index) among:
+                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+                  sequence if provided).
+                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
+                  acceptable input length for the model if that argument is not provided.
+                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
+                  lengths).
+            max_length (`int`, *optional*):
+                Maximum length of the returned list and optionally padding length (see above).
+            truncation (`bool`, *optional*):
+                Activates truncation to cut input sequences longer than `max_length` to `max_length`.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model.
+            - **input_image_embeds** -- Pixel values to be fed to a model.
+            - **image_sizes** -- List of tuples specifying the size of each image in `input_image_embeds`.
+            - **image_attention_mask** -- List of attention masks for each image in `input_image_embeds`.
+            - **input_audio_embeds** -- Audio embeddings to be fed to a model.
+            - **audio_embed_sizes** -- List of integers specifying the size of each audio in `input_audio_embeds`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model.
+        """
+        image_inputs = self.image_processor(images, return_tensors=return_tensors) if images is not None else {}
+        audio_inputs = self.audio_processor(audios, return_tensors=return_tensors) if audios is not None else {}
+        inputs = self._convert_images_audios_text_to_inputs(
+            image_inputs,
+            audio_inputs,
+            text,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+        )
+
+        # idenfity the input mode
+        if len(image_inputs) > 0 and len(audio_inputs) > 0:
+            input_mode = InputMode.VISION_SPEECH
+        elif len(image_inputs) > 0:
+            input_mode = InputMode.VISION
+        elif len(audio_inputs) > 0:
+            input_mode = InputMode.SPEECH
+        else:
+            input_mode = InputMode.LANGUAGE
+        inputs["input_mode"] = torch.tensor([input_mode.value], dtype=torch.long)
+
+        return inputs
+
+    @property
+    def special_image_token_id(self):
+        return self.tokenizer.convert_tokens_to_ids(self.special_image_token)
+
+    def get_special_image_token_id(self):
+        return self.tokenizer.convert_tokens_to_ids(self.special_image_token)
+
+    @property
+    def chat_template(self):
+        return self.tokenizer.chat_template
+
+    def _convert_images_audios_text_to_inputs(
+        self, images, audios, text, padding=False, truncation=None, max_length=None, return_tensors=None
+    ):
+        # prepare image id to image input ids
+        if len(images) > 0:
+            input_image_embeds = images["input_image_embeds"]
+            image_sizes = images["image_sizes"]
+            image_attention_mask = images["image_attention_mask"]
+            num_img_tokens = images['num_img_tokens']
+        else:
+            input_image_embeds = torch.tensor([])
+            image_sizes = torch.tensor([])
+            image_attention_mask = torch.tensor([])
+            num_img_tokens = []
+
+        # prepare audio id to audio input ids
+        if len(audios) > 0:
+            input_audio_embeds = audios["input_audio_embeds"]
+            audio_embed_sizes = audios["audio_embed_sizes"]
+            audio_attention_mask = audios.get("audio_attention_mask", None)
+        else:
+            input_audio_embeds = torch.tensor([])
+            audio_embed_sizes = torch.tensor([])
+            audio_attention_mask = None
+
+        # Replace certain special tokens for compatibility
+        # Ref: https://stackoverflow.com/questions/11475885/python-replace-regex
+        if isinstance(text, str):
+            text = [text]
+        assert isinstance(text, list)
+        processed_text = [re.sub(_COMPATIBLE_IMAGE_SPECIAL_TOKEN_PATTERN, _IMAGE_SPECIAL_TOKEN, t) for t in text]
+        processed_text = [re.sub(_COMPATIBLE_AUDIO_SPECIAL_TOKEN_PATTERN, _AUDIO_SPECIAL_TOKEN, t) for t in processed_text]
+
+        input_ids_list = [self.tokenizer(t).input_ids for t in processed_text]
+
+        img_cnt, audio_cnt = 0, 0  # only needed for later assertion
+        image_token_count_iter = iter(num_img_tokens)
+        audio_embed_size_iter = iter(audio_embed_sizes.tolist())
+        new_input_ids_list = []
+        for input_ids in input_ids_list:
+            i = 0
+            while i < len(input_ids):
+                token_id = input_ids[i]
+                if token_id == _AUDIO_SPECIAL_TOKEN_ID:
+                    token_count = next(audio_embed_size_iter)
+                    audio_cnt += 1
+                elif token_id == _IMAGE_SPECIAL_TOKEN_ID:
+                    token_count = next(image_token_count_iter)
+                    img_cnt += 1
+                else:
+                    i += 1
+                    continue
+                tokens = [token_id] * token_count
+                input_ids = input_ids[:i] + tokens + input_ids[i + 1:]
+                i += token_count
+            input_ids = torch.tensor(input_ids, dtype=torch.long)
+            new_input_ids_list.append(input_ids)
+        lengths = torch.tensor([len(input_ids) for input_ids in new_input_ids_list])
+        max_len = lengths.max()
+        input_ids = input_ids.new_full((len(new_input_ids_list), max_len), self.tokenizer.pad_token_id)
+        # batched inference requires left padding
+        for i in range(len(new_input_ids_list)):
+            input_ids[i, max_len - len(new_input_ids_list[i]):] = new_input_ids_list[i]
+
+        # If the below assertion fails, it might be that input pure-text
+        # messages contain image/audio special tokens literally
+        # (<|endoftext10|>, <|endoftext11|>).
+        assert (
+            img_cnt == len(num_img_tokens)
+        ), (
+            f"Number of image tokens in prompt_token_ids ({img_cnt}) "
+            f"does not match number of images ({len(num_img_tokens)})"
+        )
+        assert (
+            audio_cnt == len(audio_embed_sizes)
+        ), (
+            f"Number of audio tokens in prompt_token_ids ({audio_cnt}) "
+            f"does not match number of audios ({len(audio_embed_sizes)})"
+        )
+
+        # prepare attention mask
+        seq_range = torch.arange(max_len - 1, -1, -1)
+        attention_mask = seq_range.unsqueeze(0) < lengths.unsqueeze(1)
+
+        # prepare batch feature
+        data = {
+            "input_ids": input_ids,
+            "input_image_embeds": input_image_embeds,
+            "image_sizes": image_sizes,
+            "image_attention_mask": image_attention_mask,
+            "input_audio_embeds": input_audio_embeds,
+            "audio_embed_sizes": audio_embed_sizes,
+            "audio_attention_mask": audio_attention_mask,
+            "attention_mask": attention_mask,
+        }
+
+        return BatchFeature(
+            data=data
+        )
+
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Llama
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to GPT2Tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Llama
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to GPT2Tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.model_input_names
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        audio_processor_input_names = self.audio_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names + audio_processor_input_names))
+
+
+AutoImageProcessor.register("Phi4MMImageProcessor", Phi4MMImageProcessor)
+AutoFeatureExtractor.register("Phi4MMAudioFeatureExtractor", Phi4MMAudioFeatureExtractor)

processor_config.json

@@ -0,0 +1,6 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_phi4mm.Phi4MMProcessor"
+  },
+  "processor_class": "Phi4MMProcessor"
+}

trainer_state.json

@@ -0,0 +1,2329 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 252.32,
+  "eval_steps": 500,
+  "global_step": 3280,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.8,
+      "grad_norm": 79.95787811279297,
+      "learning_rate": 1.2e-05,
+      "loss": 10.8172,
+      "step": 10
+    },
+    {
+      "epoch": 1.56,
+      "grad_norm": 15.179234504699707,
+      "learning_rate": 3.2000000000000005e-05,
+      "loss": 1.4707,
+      "step": 20
+    },
+    {
+      "epoch": 2.32,
+      "grad_norm": 5.771359443664551,
+      "learning_rate": 5.2000000000000004e-05,
+      "loss": 0.8206,
+      "step": 30
+    },
+    {
+      "epoch": 3.08,
+      "grad_norm": 5.559211730957031,
+      "learning_rate": 7.2e-05,
+      "loss": 0.9515,
+      "step": 40
+    },
+    {
+      "epoch": 3.88,
+      "grad_norm": 2.7674646377563477,
+      "learning_rate": 9.200000000000001e-05,
+      "loss": 0.7728,
+      "step": 50
+    },
+    {
+      "epoch": 4.64,
+      "grad_norm": 3.8835511207580566,
+      "learning_rate": 9.994979079497908e-05,
+      "loss": 0.8743,
+      "step": 60
+    },
+    {
+      "epoch": 5.4,
+      "grad_norm": 2.758652925491333,
+      "learning_rate": 9.986610878661087e-05,
+      "loss": 0.6475,
+      "step": 70
+    },
+    {
+      "epoch": 6.16,
+      "grad_norm": 1.3047797679901123,
+      "learning_rate": 9.978242677824268e-05,
+      "loss": 0.7883,
+      "step": 80
+    },
+    {
+      "epoch": 6.96,
+      "grad_norm": 3.255369186401367,
+      "learning_rate": 9.969874476987448e-05,
+      "loss": 2.3198,
+      "step": 90
+    },
+    {
+      "epoch": 7.72,
+      "grad_norm": 2.8563857078552246,
+      "learning_rate": 9.961506276150628e-05,
+      "loss": 0.6541,
+      "step": 100
+    },
+    {
+      "epoch": 8.48,
+      "grad_norm": 1.9255300760269165,
+      "learning_rate": 9.953138075313808e-05,
+      "loss": 0.7187,
+      "step": 110
+    },
+    {
+      "epoch": 9.24,
+      "grad_norm": 2.496511697769165,
+      "learning_rate": 9.944769874476987e-05,
+      "loss": 0.6086,
+      "step": 120
+    },
+    {
+      "epoch": 10.0,
+      "grad_norm": 1.0335757732391357,
+      "learning_rate": 9.936401673640167e-05,
+      "loss": 0.5804,
+      "step": 130
+    },
+    {
+      "epoch": 10.8,
+      "grad_norm": 2.37092661857605,
+      "learning_rate": 9.928033472803347e-05,
+      "loss": 0.4609,
+      "step": 140
+    },
+    {
+      "epoch": 11.56,
+      "grad_norm": 5.012497901916504,
+      "learning_rate": 9.919665271966527e-05,
+      "loss": 0.4665,
+      "step": 150
+    },
+    {
+      "epoch": 12.32,
+      "grad_norm": 1.2357938289642334,
+      "learning_rate": 9.911297071129707e-05,
+      "loss": 0.541,
+      "step": 160
+    },
+    {
+      "epoch": 13.08,
+      "grad_norm": 3.381162405014038,
+      "learning_rate": 9.902928870292888e-05,
+      "loss": 0.5019,
+      "step": 170
+    },
+    {
+      "epoch": 13.88,
+      "grad_norm": 1.603922724723816,
+      "learning_rate": 9.894560669456067e-05,
+      "loss": 0.408,
+      "step": 180
+    },
+    {
+      "epoch": 14.64,
+      "grad_norm": 1.6547144651412964,
+      "learning_rate": 9.886192468619247e-05,
+      "loss": 0.3947,
+      "step": 190
+    },
+    {
+      "epoch": 15.4,
+      "grad_norm": 0.8650034070014954,
+      "learning_rate": 9.877824267782427e-05,
+      "loss": 0.593,
+      "step": 200
+    },
+    {
+      "epoch": 16.16,
+      "grad_norm": 0.9181660413742065,
+      "learning_rate": 9.869456066945607e-05,
+      "loss": 0.3337,
+      "step": 210
+    },
+    {
+      "epoch": 16.96,
+      "grad_norm": 1.2371879816055298,
+      "learning_rate": 9.861087866108786e-05,
+      "loss": 0.4019,
+      "step": 220
+    },
+    {
+      "epoch": 17.72,
+      "grad_norm": 1.0304924249649048,
+      "learning_rate": 9.852719665271966e-05,
+      "loss": 0.3883,
+      "step": 230
+    },
+    {
+      "epoch": 18.48,
+      "grad_norm": 1.2461388111114502,
+      "learning_rate": 9.844351464435146e-05,
+      "loss": 0.6623,
+      "step": 240
+    },
+    {
+      "epoch": 19.24,
+      "grad_norm": 0.6923120617866516,
+      "learning_rate": 9.835983263598327e-05,
+      "loss": 0.3604,
+      "step": 250
+    },
+    {
+      "epoch": 20.0,
+      "grad_norm": 4.839296817779541,
+      "learning_rate": 9.827615062761507e-05,
+      "loss": 0.387,
+      "step": 260
+    },
+    {
+      "epoch": 20.8,
+      "grad_norm": 1.1753798723220825,
+      "learning_rate": 9.819246861924687e-05,
+      "loss": 0.3514,
+      "step": 270
+    },
+    {
+      "epoch": 21.56,
+      "grad_norm": 0.9725382924079895,
+      "learning_rate": 9.810878661087866e-05,
+      "loss": 0.3922,
+      "step": 280
+    },
+    {
+      "epoch": 22.32,
+      "grad_norm": 0.645702600479126,
+      "learning_rate": 9.802510460251046e-05,
+      "loss": 0.3331,
+      "step": 290
+    },
+    {
+      "epoch": 23.08,
+      "grad_norm": 1.4330626726150513,
+      "learning_rate": 9.794142259414226e-05,
+      "loss": 0.3946,
+      "step": 300
+    },
+    {
+      "epoch": 23.88,
+      "grad_norm": 1.2405084371566772,
+      "learning_rate": 9.785774058577406e-05,
+      "loss": 0.354,
+      "step": 310
+    },
+    {
+      "epoch": 24.64,
+      "grad_norm": 0.9036368131637573,
+      "learning_rate": 9.777405857740585e-05,
+      "loss": 0.3856,
+      "step": 320
+    },
+    {
+      "epoch": 25.4,
+      "grad_norm": 0.7258665561676025,
+      "learning_rate": 9.769037656903767e-05,
+      "loss": 0.2859,
+      "step": 330
+    },
+    {
+      "epoch": 26.16,
+      "grad_norm": 0.4196911156177521,
+      "learning_rate": 9.760669456066946e-05,
+      "loss": 0.2962,
+      "step": 340
+    },
+    {
+      "epoch": 26.96,
+      "grad_norm": 1.2342430353164673,
+      "learning_rate": 9.752301255230126e-05,
+      "loss": 0.3721,
+      "step": 350
+    },
+    {
+      "epoch": 27.72,
+      "grad_norm": 0.6905648112297058,
+      "learning_rate": 9.743933054393306e-05,
+      "loss": 0.2722,
+      "step": 360
+    },
+    {
+      "epoch": 28.48,
+      "grad_norm": 0.851387083530426,
+      "learning_rate": 9.735564853556486e-05,
+      "loss": 0.3564,
+      "step": 370
+    },
+    {
+      "epoch": 29.24,
+      "grad_norm": 0.5951876640319824,
+      "learning_rate": 9.727196652719665e-05,
+      "loss": 0.2892,
+      "step": 380
+    },
+    {
+      "epoch": 30.0,
+      "grad_norm": 0.5764455199241638,
+      "learning_rate": 9.718828451882845e-05,
+      "loss": 0.3402,
+      "step": 390
+    },
+    {
+      "epoch": 30.8,
+      "grad_norm": 7.990225315093994,
+      "learning_rate": 9.710460251046025e-05,
+      "loss": 0.3439,
+      "step": 400
+    },
+    {
+      "epoch": 31.56,
+      "grad_norm": 1.119363784790039,
+      "learning_rate": 9.702092050209205e-05,
+      "loss": 0.2595,
+      "step": 410
+    },
+    {
+      "epoch": 32.32,
+      "grad_norm": 27.885967254638672,
+      "learning_rate": 9.693723849372386e-05,
+      "loss": 0.9194,
+      "step": 420
+    },
+    {
+      "epoch": 33.08,
+      "grad_norm": 5.194603443145752,
+      "learning_rate": 9.685355648535566e-05,
+      "loss": 0.3321,
+      "step": 430
+    },
+    {
+      "epoch": 33.88,
+      "grad_norm": 2.1785478591918945,
+      "learning_rate": 9.676987447698745e-05,
+      "loss": 0.3875,
+      "step": 440
+    },
+    {
+      "epoch": 34.64,
+      "grad_norm": 0.844071626663208,
+      "learning_rate": 9.668619246861925e-05,
+      "loss": 0.2458,
+      "step": 450
+    },
+    {
+      "epoch": 35.4,
+      "grad_norm": 2.1278724670410156,
+      "learning_rate": 9.660251046025105e-05,
+      "loss": 0.358,
+      "step": 460
+    },
+    {
+      "epoch": 36.16,
+      "grad_norm": 0.6194890141487122,
+      "learning_rate": 9.651882845188285e-05,
+      "loss": 0.3334,
+      "step": 470
+    },
+    {
+      "epoch": 36.96,
+      "grad_norm": 0.7329260110855103,
+      "learning_rate": 9.643514644351464e-05,
+      "loss": 0.2981,
+      "step": 480
+    },
+    {
+      "epoch": 37.72,
+      "grad_norm": 0.8790725469589233,
+      "learning_rate": 9.635146443514644e-05,
+      "loss": 0.2903,
+      "step": 490
+    },
+    {
+      "epoch": 38.48,
+      "grad_norm": 0.5892160534858704,
+      "learning_rate": 9.626778242677825e-05,
+      "loss": 0.362,
+      "step": 500
+    },
+    {
+      "epoch": 39.24,
+      "grad_norm": 1.6055381298065186,
+      "learning_rate": 9.618410041841005e-05,
+      "loss": 0.3838,
+      "step": 510
+    },
+    {
+      "epoch": 40.0,
+      "grad_norm": 11.784494400024414,
+      "learning_rate": 9.610041841004185e-05,
+      "loss": 0.3082,
+      "step": 520
+    },
+    {
+      "epoch": 40.8,
+      "grad_norm": 0.5743730664253235,
+      "learning_rate": 9.601673640167365e-05,
+      "loss": 0.31,
+      "step": 530
+    },
+    {
+      "epoch": 41.56,
+      "grad_norm": 0.592658519744873,
+      "learning_rate": 9.593305439330544e-05,
+      "loss": 0.5208,
+      "step": 540
+    },
+    {
+      "epoch": 42.32,
+      "grad_norm": 5.394102573394775,
+      "learning_rate": 9.584937238493724e-05,
+      "loss": 0.4108,
+      "step": 550
+    },
+    {
+      "epoch": 43.08,
+      "grad_norm": 0.7347181439399719,
+      "learning_rate": 9.576569037656904e-05,
+      "loss": 0.3476,
+      "step": 560
+    },
+    {
+      "epoch": 43.88,
+      "grad_norm": 0.9459385871887207,
+      "learning_rate": 9.568200836820084e-05,
+      "loss": 0.3008,
+      "step": 570
+    },
+    {
+      "epoch": 44.64,
+      "grad_norm": 5.524572849273682,
+      "learning_rate": 9.559832635983263e-05,
+      "loss": 0.266,
+      "step": 580
+    },
+    {
+      "epoch": 45.4,
+      "grad_norm": 0.705575704574585,
+      "learning_rate": 9.551464435146445e-05,
+      "loss": 0.3705,
+      "step": 590
+    },
+    {
+      "epoch": 46.16,
+      "grad_norm": 0.5305906534194946,
+      "learning_rate": 9.543096234309624e-05,
+      "loss": 0.2619,
+      "step": 600
+    },
+    {
+      "epoch": 46.96,
+      "grad_norm": 3.7031350135803223,
+      "learning_rate": 9.534728033472804e-05,
+      "loss": 0.4041,
+      "step": 610
+    },
+    {
+      "epoch": 47.72,
+      "grad_norm": 0.9455975294113159,
+      "learning_rate": 9.526359832635984e-05,
+      "loss": 0.2948,
+      "step": 620
+    },
+    {
+      "epoch": 48.48,
+      "grad_norm": 0.13628531992435455,
+      "learning_rate": 9.517991631799164e-05,
+      "loss": 0.2962,
+      "step": 630
+    },
+    {
+      "epoch": 49.24,
+      "grad_norm": 0.5263031721115112,
+      "learning_rate": 9.509623430962343e-05,
+      "loss": 0.3395,
+      "step": 640
+    },
+    {
+      "epoch": 50.0,
+      "grad_norm": 0.47290968894958496,
+      "learning_rate": 9.501255230125523e-05,
+      "loss": 0.3163,
+      "step": 650
+    },
+    {
+      "epoch": 50.8,
+      "grad_norm": 0.7216980457305908,
+      "learning_rate": 9.492887029288703e-05,
+      "loss": 0.303,
+      "step": 660
+    },
+    {
+      "epoch": 51.56,
+      "grad_norm": 0.44606828689575195,
+      "learning_rate": 9.484518828451884e-05,
+      "loss": 0.3086,
+      "step": 670
+    },
+    {
+      "epoch": 52.32,
+      "grad_norm": 0.6489527821540833,
+      "learning_rate": 9.476150627615064e-05,
+      "loss": 0.3226,
+      "step": 680
+    },
+    {
+      "epoch": 53.08,
+      "grad_norm": 0.4259757101535797,
+      "learning_rate": 9.467782426778243e-05,
+      "loss": 0.2513,
+      "step": 690
+    },
+    {
+      "epoch": 53.88,
+      "grad_norm": 0.5952382683753967,
+      "learning_rate": 9.459414225941423e-05,
+      "loss": 0.3113,
+      "step": 700
+    },
+    {
+      "epoch": 54.64,
+      "grad_norm": 1.0894474983215332,
+      "learning_rate": 9.451046025104603e-05,
+      "loss": 0.2583,
+      "step": 710
+    },
+    {
+      "epoch": 55.4,
+      "grad_norm": 0.5647149085998535,
+      "learning_rate": 9.442677824267783e-05,
+      "loss": 0.3042,
+      "step": 720
+    },
+    {
+      "epoch": 56.16,
+      "grad_norm": 0.9288455843925476,
+      "learning_rate": 9.434309623430963e-05,
+      "loss": 0.3281,
+      "step": 730
+    },
+    {
+      "epoch": 56.96,
+      "grad_norm": 0.4492562711238861,
+      "learning_rate": 9.425941422594142e-05,
+      "loss": 0.3379,
+      "step": 740
+    },
+    {
+      "epoch": 57.72,
+      "grad_norm": 1.0189441442489624,
+      "learning_rate": 9.417573221757323e-05,
+      "loss": 0.3194,
+      "step": 750
+    },
+    {
+      "epoch": 58.48,
+      "grad_norm": 0.40981024503707886,
+      "learning_rate": 9.409205020920503e-05,
+      "loss": 0.2423,
+      "step": 760
+    },
+    {
+      "epoch": 59.24,
+      "grad_norm": 0.354190468788147,
+      "learning_rate": 9.400836820083683e-05,
+      "loss": 0.3266,
+      "step": 770
+    },
+    {
+      "epoch": 60.0,
+      "grad_norm": 0.4593465328216553,
+      "learning_rate": 9.392468619246863e-05,
+      "loss": 0.3068,
+      "step": 780
+    },
+    {
+      "epoch": 60.8,
+      "grad_norm": 0.759248673915863,
+      "learning_rate": 9.384100418410042e-05,
+      "loss": 0.3078,
+      "step": 790
+    },
+    {
+      "epoch": 61.56,
+      "grad_norm": 0.8159565925598145,
+      "learning_rate": 9.375732217573222e-05,
+      "loss": 0.2668,
+      "step": 800
+    },
+    {
+      "epoch": 62.32,
+      "grad_norm": 0.7888874411582947,
+      "learning_rate": 9.367364016736402e-05,
+      "loss": 0.2936,
+      "step": 810
+    },
+    {
+      "epoch": 63.08,
+      "grad_norm": 0.8634017705917358,
+      "learning_rate": 9.358995815899582e-05,
+      "loss": 0.3182,
+      "step": 820
+    },
+    {
+      "epoch": 63.88,
+      "grad_norm": 0.5454868078231812,
+      "learning_rate": 9.350627615062762e-05,
+      "loss": 0.3117,
+      "step": 830
+    },
+    {
+      "epoch": 64.64,
+      "grad_norm": 0.9895418286323547,
+      "learning_rate": 9.342259414225943e-05,
+      "loss": 0.3014,
+      "step": 840
+    },
+    {
+      "epoch": 65.4,
+      "grad_norm": 1.4120994806289673,
+      "learning_rate": 9.333891213389122e-05,
+      "loss": 0.3001,
+      "step": 850
+    },
+    {
+      "epoch": 66.16,
+      "grad_norm": 0.6670629978179932,
+      "learning_rate": 9.325523012552302e-05,
+      "loss": 0.2651,
+      "step": 860
+    },
+    {
+      "epoch": 66.96,
+      "grad_norm": 0.6429235339164734,
+      "learning_rate": 9.317154811715482e-05,
+      "loss": 0.2935,
+      "step": 870
+    },
+    {
+      "epoch": 67.72,
+      "grad_norm": 0.6127046942710876,
+      "learning_rate": 9.308786610878662e-05,
+      "loss": 0.2936,
+      "step": 880
+    },
+    {
+      "epoch": 68.48,
+      "grad_norm": 0.6747534275054932,
+      "learning_rate": 9.300418410041841e-05,
+      "loss": 0.2557,
+      "step": 890
+    },
+    {
+      "epoch": 69.24,
+      "grad_norm": 0.8579817414283752,
+      "learning_rate": 9.292050209205021e-05,
+      "loss": 0.3045,
+      "step": 900
+    },
+    {
+      "epoch": 70.0,
+      "grad_norm": 0.24220693111419678,
+      "learning_rate": 9.283682008368201e-05,
+      "loss": 0.2818,
+      "step": 910
+    },
+    {
+      "epoch": 70.8,
+      "grad_norm": 0.5585035681724548,
+      "learning_rate": 9.275313807531382e-05,
+      "loss": 0.2579,
+      "step": 920
+    },
+    {
+      "epoch": 71.56,
+      "grad_norm": 0.5965076684951782,
+      "learning_rate": 9.266945606694562e-05,
+      "loss": 0.3582,
+      "step": 930
+    },
+    {
+      "epoch": 72.32,
+      "grad_norm": 0.4647338092327118,
+      "learning_rate": 9.258577405857742e-05,
+      "loss": 0.2384,
+      "step": 940
+    },
+    {
+      "epoch": 73.08,
+      "grad_norm": 0.4868353605270386,
+      "learning_rate": 9.250209205020921e-05,
+      "loss": 0.3097,
+      "step": 950
+    },
+    {
+      "epoch": 73.88,
+      "grad_norm": 0.4390128552913666,
+      "learning_rate": 9.241841004184101e-05,
+      "loss": 0.2722,
+      "step": 960
+    },
+    {
+      "epoch": 74.64,
+      "grad_norm": 1.2558406591415405,
+      "learning_rate": 9.233472803347281e-05,
+      "loss": 0.2546,
+      "step": 970
+    },
+    {
+      "epoch": 75.4,
+      "grad_norm": 1.468024492263794,
+      "learning_rate": 9.225104602510461e-05,
+      "loss": 0.3539,
+      "step": 980
+    },
+    {
+      "epoch": 76.16,
+      "grad_norm": 1.3816922903060913,
+      "learning_rate": 9.21673640167364e-05,
+      "loss": 0.3131,
+      "step": 990
+    },
+    {
+      "epoch": 76.96,
+      "grad_norm": 1.0851378440856934,
+      "learning_rate": 9.208368200836822e-05,
+      "loss": 0.3418,
+      "step": 1000
+    },
+    {
+      "epoch": 77.72,
+      "grad_norm": 0.6077636480331421,
+      "learning_rate": 9.200000000000001e-05,
+      "loss": 0.301,
+      "step": 1010
+    },
+    {
+      "epoch": 78.48,
+      "grad_norm": 0.39989063143730164,
+      "learning_rate": 9.191631799163181e-05,
+      "loss": 0.2551,
+      "step": 1020
+    },
+    {
+      "epoch": 79.24,
+      "grad_norm": 0.5581826567649841,
+      "learning_rate": 9.183263598326361e-05,
+      "loss": 0.2955,
+      "step": 1030
+    },
+    {
+      "epoch": 80.0,
+      "grad_norm": 0.42415767908096313,
+      "learning_rate": 9.17489539748954e-05,
+      "loss": 0.3024,
+      "step": 1040
+    },
+    {
+      "epoch": 80.8,
+      "grad_norm": 0.431111603975296,
+      "learning_rate": 9.16652719665272e-05,
+      "loss": 0.3108,
+      "step": 1050
+    },
+    {
+      "epoch": 81.56,
+      "grad_norm": 0.45482027530670166,
+      "learning_rate": 9.1581589958159e-05,
+      "loss": 0.25,
+      "step": 1060
+    },
+    {
+      "epoch": 82.32,
+      "grad_norm": 0.6043058633804321,
+      "learning_rate": 9.14979079497908e-05,
+      "loss": 0.3466,
+      "step": 1070
+    },
+    {
+      "epoch": 83.08,
+      "grad_norm": 0.7600142955780029,
+      "learning_rate": 9.14142259414226e-05,
+      "loss": 0.2888,
+      "step": 1080
+    },
+    {
+      "epoch": 83.88,
+      "grad_norm": 0.36540111899375916,
+      "learning_rate": 9.133054393305441e-05,
+      "loss": 0.2896,
+      "step": 1090
+    },
+    {
+      "epoch": 84.64,
+      "grad_norm": 1.1366398334503174,
+      "learning_rate": 9.12468619246862e-05,
+      "loss": 0.3023,
+      "step": 1100
+    },
+    {
+      "epoch": 85.4,
+      "grad_norm": 0.646086573600769,
+      "learning_rate": 9.1163179916318e-05,
+      "loss": 0.2462,
+      "step": 1110
+    },
+    {
+      "epoch": 86.16,
+      "grad_norm": 0.6224349141120911,
+      "learning_rate": 9.10794979079498e-05,
+      "loss": 0.2741,
+      "step": 1120
+    },
+    {
+      "epoch": 86.96,
+      "grad_norm": 0.8657971024513245,
+      "learning_rate": 9.09958158995816e-05,
+      "loss": 0.3016,
+      "step": 1130
+    },
+    {
+      "epoch": 87.72,
+      "grad_norm": 0.86732017993927,
+      "learning_rate": 9.09121338912134e-05,
+      "loss": 0.263,
+      "step": 1140
+    },
+    {
+      "epoch": 88.48,
+      "grad_norm": 0.8562549948692322,
+      "learning_rate": 9.08284518828452e-05,
+      "loss": 0.2726,
+      "step": 1150
+    },
+    {
+      "epoch": 89.24,
+      "grad_norm": 0.5194992423057556,
+      "learning_rate": 9.074476987447699e-05,
+      "loss": 0.3161,
+      "step": 1160
+    },
+    {
+      "epoch": 90.0,
+      "grad_norm": 0.3380357027053833,
+      "learning_rate": 9.066108786610879e-05,
+      "loss": 0.26,
+      "step": 1170
+    },
+    {
+      "epoch": 90.8,
+      "grad_norm": 0.4834354519844055,
+      "learning_rate": 9.057740585774059e-05,
+      "loss": 0.2901,
+      "step": 1180
+    },
+    {
+      "epoch": 91.56,
+      "grad_norm": 0.7634447813034058,
+      "learning_rate": 9.04937238493724e-05,
+      "loss": 0.2471,
+      "step": 1190
+    },
+    {
+      "epoch": 92.32,
+      "grad_norm": 0.5605065822601318,
+      "learning_rate": 9.04100418410042e-05,
+      "loss": 0.3091,
+      "step": 1200
+    },
+    {
+      "epoch": 93.08,
+      "grad_norm": 0.6867684721946716,
+      "learning_rate": 9.0326359832636e-05,
+      "loss": 0.292,
+      "step": 1210
+    },
+    {
+      "epoch": 93.88,
+      "grad_norm": 0.5395390391349792,
+      "learning_rate": 9.024267782426779e-05,
+      "loss": 0.2727,
+      "step": 1220
+    },
+    {
+      "epoch": 94.64,
+      "grad_norm": 0.8648020029067993,
+      "learning_rate": 9.015899581589959e-05,
+      "loss": 0.3273,
+      "step": 1230
+    },
+    {
+      "epoch": 95.4,
+      "grad_norm": 0.5256586074829102,
+      "learning_rate": 9.007531380753139e-05,
+      "loss": 0.2611,
+      "step": 1240
+    },
+    {
+      "epoch": 96.16,
+      "grad_norm": 0.726409375667572,
+      "learning_rate": 8.999163179916318e-05,
+      "loss": 0.2993,
+      "step": 1250
+    },
+    {
+      "epoch": 96.96,
+      "grad_norm": 0.5897384285926819,
+      "learning_rate": 8.990794979079498e-05,
+      "loss": 0.2628,
+      "step": 1260
+    },
+    {
+      "epoch": 97.72,
+      "grad_norm": 0.3650963306427002,
+      "learning_rate": 8.982426778242678e-05,
+      "loss": 0.2975,
+      "step": 1270
+    },
+    {
+      "epoch": 98.48,
+      "grad_norm": 0.6548069715499878,
+      "learning_rate": 8.974058577405858e-05,
+      "loss": 0.2732,
+      "step": 1280
+    },
+    {
+      "epoch": 99.24,
+      "grad_norm": 0.6239974498748779,
+      "learning_rate": 8.965690376569037e-05,
+      "loss": 0.3471,
+      "step": 1290
+    },
+    {
+      "epoch": 100.0,
+      "grad_norm": 0.00010844325879588723,
+      "learning_rate": 8.957322175732217e-05,
+      "loss": 0.2865,
+      "step": 1300
+    },
+    {
+      "epoch": 100.8,
+      "grad_norm": 55.954044342041016,
+      "learning_rate": 8.94979079497908e-05,
+      "loss": 2.0083,
+      "step": 1310
+    },
+    {
+      "epoch": 101.56,
+      "grad_norm": 1.1408623456954956,
+      "learning_rate": 8.94142259414226e-05,
+      "loss": 0.7388,
+      "step": 1320
+    },
+    {
+      "epoch": 102.32,
+      "grad_norm": 0.7127572298049927,
+      "learning_rate": 8.93305439330544e-05,
+      "loss": 0.2341,
+      "step": 1330
+    },
+    {
+      "epoch": 103.08,
+      "grad_norm": 188.80079650878906,
+      "learning_rate": 8.92468619246862e-05,
+      "loss": 0.8824,
+      "step": 1340
+    },
+    {
+      "epoch": 103.88,
+      "grad_norm": 22.78482437133789,
+      "learning_rate": 8.9163179916318e-05,
+      "loss": 0.6129,
+      "step": 1350
+    },
+    {
+      "epoch": 104.64,
+      "grad_norm": 1.2376023530960083,
+      "learning_rate": 8.90794979079498e-05,
+      "loss": 0.4303,
+      "step": 1360
+    },
+    {
+      "epoch": 105.4,
+      "grad_norm": 1.032791018486023,
+      "learning_rate": 8.899581589958159e-05,
+      "loss": 0.3326,
+      "step": 1370
+    },
+    {
+      "epoch": 106.16,
+      "grad_norm": 0.6220753192901611,
+      "learning_rate": 8.891213389121339e-05,
+      "loss": 0.3165,
+      "step": 1380
+    },
+    {
+      "epoch": 106.96,
+      "grad_norm": 0.9835271835327148,
+      "learning_rate": 8.882845188284519e-05,
+      "loss": 0.2899,
+      "step": 1390
+    },
+    {
+      "epoch": 107.72,
+      "grad_norm": 0.4846683144569397,
+      "learning_rate": 8.8744769874477e-05,
+      "loss": 0.3234,
+      "step": 1400
+    },
+    {
+      "epoch": 108.48,
+      "grad_norm": 42.94027328491211,
+      "learning_rate": 8.86610878661088e-05,
+      "loss": 0.2739,
+      "step": 1410
+    },
+    {
+      "epoch": 109.24,
+      "grad_norm": 0.42791783809661865,
+      "learning_rate": 8.857740585774059e-05,
+      "loss": 0.3374,
+      "step": 1420
+    },
+    {
+      "epoch": 110.0,
+      "grad_norm": 0.5355058312416077,
+      "learning_rate": 8.849372384937239e-05,
+      "loss": 0.2941,
+      "step": 1430
+    },
+    {
+      "epoch": 110.8,
+      "grad_norm": 0.43855488300323486,
+      "learning_rate": 8.841004184100419e-05,
+      "loss": 0.254,
+      "step": 1440
+    },
+    {
+      "epoch": 111.56,
+      "grad_norm": 0.6513474583625793,
+      "learning_rate": 8.832635983263599e-05,
+      "loss": 0.2992,
+      "step": 1450
+    },
+    {
+      "epoch": 112.32,
+      "grad_norm": 0.5990163087844849,
+      "learning_rate": 8.824267782426778e-05,
+      "loss": 0.2838,
+      "step": 1460
+    },
+    {
+      "epoch": 113.08,
+      "grad_norm": 0.5194154381752014,
+      "learning_rate": 8.815899581589958e-05,
+      "loss": 0.2896,
+      "step": 1470
+    },
+    {
+      "epoch": 113.88,
+      "grad_norm": 0.90041583776474,
+      "learning_rate": 8.807531380753139e-05,
+      "loss": 0.2853,
+      "step": 1480
+    },
+    {
+      "epoch": 114.64,
+      "grad_norm": 0.46713006496429443,
+      "learning_rate": 8.799163179916319e-05,
+      "loss": 0.341,
+      "step": 1490
+    },
+    {
+      "epoch": 115.4,
+      "grad_norm": 1.0562089681625366,
+      "learning_rate": 8.790794979079499e-05,
+      "loss": 0.2485,
+      "step": 1500
+    },
+    {
+      "epoch": 116.16,
+      "grad_norm": 0.7914270162582397,
+      "learning_rate": 8.782426778242678e-05,
+      "loss": 0.2931,
+      "step": 1510
+    },
+    {
+      "epoch": 116.96,
+      "grad_norm": 0.7104222178459167,
+      "learning_rate": 8.774058577405858e-05,
+      "loss": 0.3166,
+      "step": 1520
+    },
+    {
+      "epoch": 117.72,
+      "grad_norm": 0.6477789878845215,
+      "learning_rate": 8.765690376569038e-05,
+      "loss": 0.2712,
+      "step": 1530
+    },
+    {
+      "epoch": 118.48,
+      "grad_norm": 0.2977544069290161,
+      "learning_rate": 8.757322175732218e-05,
+      "loss": 0.2537,
+      "step": 1540
+    },
+    {
+      "epoch": 119.24,
+      "grad_norm": 0.6447045803070068,
+      "learning_rate": 8.748953974895398e-05,
+      "loss": 0.2848,
+      "step": 1550
+    },
+    {
+      "epoch": 120.0,
+      "grad_norm": 0.4993550479412079,
+      "learning_rate": 8.740585774058579e-05,
+      "loss": 0.2712,
+      "step": 1560
+    },
+    {
+      "epoch": 120.8,
+      "grad_norm": 0.28479063510894775,
+      "learning_rate": 8.732217573221758e-05,
+      "loss": 0.2969,
+      "step": 1570
+    },
+    {
+      "epoch": 121.56,
+      "grad_norm": 0.7489855885505676,
+      "learning_rate": 8.723849372384938e-05,
+      "loss": 0.2512,
+      "step": 1580
+    },
+    {
+      "epoch": 122.32,
+      "grad_norm": 0.6503575444221497,
+      "learning_rate": 8.715481171548118e-05,
+      "loss": 0.268,
+      "step": 1590
+    },
+    {
+      "epoch": 123.08,
+      "grad_norm": 0.5870686769485474,
+      "learning_rate": 8.707112970711298e-05,
+      "loss": 0.302,
+      "step": 1600
+    },
+    {
+      "epoch": 123.88,
+      "grad_norm": 0.8388033509254456,
+      "learning_rate": 8.698744769874477e-05,
+      "loss": 0.2784,
+      "step": 1610
+    },
+    {
+      "epoch": 124.64,
+      "grad_norm": 0.7110853791236877,
+      "learning_rate": 8.690376569037657e-05,
+      "loss": 0.2576,
+      "step": 1620
+    },
+    {
+      "epoch": 125.4,
+      "grad_norm": 0.6697489619255066,
+      "learning_rate": 8.682008368200837e-05,
+      "loss": 0.2863,
+      "step": 1630
+    },
+    {
+      "epoch": 126.16,
+      "grad_norm": 0.6678580045700073,
+      "learning_rate": 8.673640167364017e-05,
+      "loss": 0.2945,
+      "step": 1640
+    },
+    {
+      "epoch": 126.96,
+      "grad_norm": 0.5099469423294067,
+      "learning_rate": 8.665271966527198e-05,
+      "loss": 0.2744,
+      "step": 1650
+    },
+    {
+      "epoch": 127.72,
+      "grad_norm": 0.8461157083511353,
+      "learning_rate": 8.656903765690378e-05,
+      "loss": 0.2881,
+      "step": 1660
+    },
+    {
+      "epoch": 128.48,
+      "grad_norm": 0.5801335573196411,
+      "learning_rate": 8.648535564853557e-05,
+      "loss": 0.2743,
+      "step": 1670
+    },
+    {
+      "epoch": 129.24,
+      "grad_norm": 0.00013900638441555202,
+      "learning_rate": 8.640167364016737e-05,
+      "loss": 0.2389,
+      "step": 1680
+    },
+    {
+      "epoch": 130.0,
+      "grad_norm": 0.8893792033195496,
+      "learning_rate": 8.631799163179917e-05,
+      "loss": 0.3059,
+      "step": 1690
+    },
+    {
+      "epoch": 130.8,
+      "grad_norm": 0.5562736392021179,
+      "learning_rate": 8.623430962343097e-05,
+      "loss": 0.2586,
+      "step": 1700
+    },
+    {
+      "epoch": 131.56,
+      "grad_norm": 0.41472992300987244,
+      "learning_rate": 8.615062761506276e-05,
+      "loss": 0.2632,
+      "step": 1710
+    },
+    {
+      "epoch": 132.32,
+      "grad_norm": 0.23705990612506866,
+      "learning_rate": 8.606694560669456e-05,
+      "loss": 0.2678,
+      "step": 1720
+    },
+    {
+      "epoch": 133.08,
+      "grad_norm": 0.6223066449165344,
+      "learning_rate": 8.598326359832637e-05,
+      "loss": 0.3278,
+      "step": 1730
+    },
+    {
+      "epoch": 133.88,
+      "grad_norm": 0.7489154934883118,
+      "learning_rate": 8.589958158995817e-05,
+      "loss": 0.3094,
+      "step": 1740
+    },
+    {
+      "epoch": 134.64,
+      "grad_norm": 0.0001321160380030051,
+      "learning_rate": 8.581589958158997e-05,
+      "loss": 0.2386,
+      "step": 1750
+    },
+    {
+      "epoch": 135.4,
+      "grad_norm": 0.6472476720809937,
+      "learning_rate": 8.573221757322177e-05,
+      "loss": 0.319,
+      "step": 1760
+    },
+    {
+      "epoch": 136.16,
+      "grad_norm": 0.6706916689872742,
+      "learning_rate": 8.564853556485356e-05,
+      "loss": 0.2565,
+      "step": 1770
+    },
+    {
+      "epoch": 136.96,
+      "grad_norm": 0.6396327614784241,
+      "learning_rate": 8.556485355648536e-05,
+      "loss": 0.3257,
+      "step": 1780
+    },
+    {
+      "epoch": 137.72,
+      "grad_norm": 0.945931613445282,
+      "learning_rate": 8.548117154811716e-05,
+      "loss": 0.2912,
+      "step": 1790
+    },
+    {
+      "epoch": 138.48,
+      "grad_norm": 0.4821135699748993,
+      "learning_rate": 8.539748953974896e-05,
+      "loss": 0.239,
+      "step": 1800
+    },
+    {
+      "epoch": 139.24,
+      "grad_norm": 0.5348692536354065,
+      "learning_rate": 8.531380753138077e-05,
+      "loss": 0.2632,
+      "step": 1810
+    },
+    {
+      "epoch": 140.0,
+      "grad_norm": 0.3767673671245575,
+      "learning_rate": 8.523012552301257e-05,
+      "loss": 0.2752,
+      "step": 1820
+    },
+    {
+      "epoch": 140.8,
+      "grad_norm": 0.5291851758956909,
+      "learning_rate": 8.514644351464436e-05,
+      "loss": 0.2997,
+      "step": 1830
+    },
+    {
+      "epoch": 141.56,
+      "grad_norm": 0.6979625225067139,
+      "learning_rate": 8.506276150627616e-05,
+      "loss": 0.2731,
+      "step": 1840
+    },
+    {
+      "epoch": 142.32,
+      "grad_norm": 0.2552458941936493,
+      "learning_rate": 8.497907949790796e-05,
+      "loss": 0.2046,
+      "step": 1850
+    },
+    {
+      "epoch": 143.08,
+      "grad_norm": 0.7138965725898743,
+      "learning_rate": 8.489539748953976e-05,
+      "loss": 0.3448,
+      "step": 1860
+    },
+    {
+      "epoch": 143.88,
+      "grad_norm": 0.602065920829773,
+      "learning_rate": 8.481171548117155e-05,
+      "loss": 0.3201,
+      "step": 1870
+    },
+    {
+      "epoch": 144.64,
+      "grad_norm": 0.8118859529495239,
+      "learning_rate": 8.472803347280335e-05,
+      "loss": 0.2622,
+      "step": 1880
+    },
+    {
+      "epoch": 145.4,
+      "grad_norm": 0.00017894129268825054,
+      "learning_rate": 8.464435146443515e-05,
+      "loss": 0.2388,
+      "step": 1890
+    },
+    {
+      "epoch": 146.16,
+      "grad_norm": 0.5059804320335388,
+      "learning_rate": 8.456066945606696e-05,
+      "loss": 0.3429,
+      "step": 1900
+    },
+    {
+      "epoch": 146.96,
+      "grad_norm": 0.47752124071121216,
+      "learning_rate": 8.447698744769876e-05,
+      "loss": 0.2575,
+      "step": 1910
+    },
+    {
+      "epoch": 147.72,
+      "grad_norm": 0.6666246056556702,
+      "learning_rate": 8.439330543933056e-05,
+      "loss": 0.284,
+      "step": 1920
+    },
+    {
+      "epoch": 148.48,
+      "grad_norm": 0.6572575569152832,
+      "learning_rate": 8.430962343096235e-05,
+      "loss": 0.2432,
+      "step": 1930
+    },
+    {
+      "epoch": 149.24,
+      "grad_norm": 0.644818902015686,
+      "learning_rate": 8.422594142259415e-05,
+      "loss": 0.2844,
+      "step": 1940
+    },
+    {
+      "epoch": 150.0,
+      "grad_norm": 1.5558511018753052,
+      "learning_rate": 8.414225941422595e-05,
+      "loss": 0.2891,
+      "step": 1950
+    },
+    {
+      "epoch": 150.8,
+      "grad_norm": 0.9465558528900146,
+      "learning_rate": 8.405857740585775e-05,
+      "loss": 0.2765,
+      "step": 1960
+    },
+    {
+      "epoch": 151.56,
+      "grad_norm": 0.8076108694076538,
+      "learning_rate": 8.397489539748954e-05,
+      "loss": 0.2718,
+      "step": 1970
+    },
+    {
+      "epoch": 152.32,
+      "grad_norm": 0.6895241141319275,
+      "learning_rate": 8.389121338912134e-05,
+      "loss": 0.2662,
+      "step": 1980
+    },
+    {
+      "epoch": 153.08,
+      "grad_norm": 0.42998993396759033,
+      "learning_rate": 8.380753138075314e-05,
+      "loss": 0.2715,
+      "step": 1990
+    },
+    {
+      "epoch": 153.88,
+      "grad_norm": 0.5157560706138611,
+      "learning_rate": 8.372384937238494e-05,
+      "loss": 0.2677,
+      "step": 2000
+    },
+    {
+      "epoch": 154.64,
+      "grad_norm": 0.5690245628356934,
+      "learning_rate": 8.364016736401675e-05,
+      "loss": 0.2882,
+      "step": 2010
+    },
+    {
+      "epoch": 155.4,
+      "grad_norm": 0.5015438199043274,
+      "learning_rate": 8.355648535564855e-05,
+      "loss": 0.2652,
+      "step": 2020
+    },
+    {
+      "epoch": 156.16,
+      "grad_norm": 0.3558717370033264,
+      "learning_rate": 8.347280334728034e-05,
+      "loss": 0.2938,
+      "step": 2030
+    },
+    {
+      "epoch": 156.96,
+      "grad_norm": 0.8188201189041138,
+      "learning_rate": 8.338912133891214e-05,
+      "loss": 0.2977,
+      "step": 2040
+    },
+    {
+      "epoch": 157.72,
+      "grad_norm": 0.8057423830032349,
+      "learning_rate": 8.330543933054394e-05,
+      "loss": 0.2342,
+      "step": 2050
+    },
+    {
+      "epoch": 158.48,
+      "grad_norm": 0.4321524500846863,
+      "learning_rate": 8.322175732217574e-05,
+      "loss": 0.2955,
+      "step": 2060
+    },
+    {
+      "epoch": 159.24,
+      "grad_norm": 0.5147210955619812,
+      "learning_rate": 8.313807531380753e-05,
+      "loss": 0.2526,
+      "step": 2070
+    },
+    {
+      "epoch": 160.0,
+      "grad_norm": 0.771668016910553,
+      "learning_rate": 8.305439330543933e-05,
+      "loss": 0.3067,
+      "step": 2080
+    },
+    {
+      "epoch": 160.8,
+      "grad_norm": 0.6141147017478943,
+      "learning_rate": 8.297071129707113e-05,
+      "loss": 0.2733,
+      "step": 2090
+    },
+    {
+      "epoch": 161.56,
+      "grad_norm": 1.0377253293991089,
+      "learning_rate": 8.288702928870293e-05,
+      "loss": 0.2909,
+      "step": 2100
+    },
+    {
+      "epoch": 162.32,
+      "grad_norm": 0.3286207318305969,
+      "learning_rate": 8.280334728033472e-05,
+      "loss": 0.2491,
+      "step": 2110
+    },
+    {
+      "epoch": 163.08,
+      "grad_norm": 0.5966292023658752,
+      "learning_rate": 8.271966527196652e-05,
+      "loss": 0.2761,
+      "step": 2120
+    },
+    {
+      "epoch": 163.88,
+      "grad_norm": 0.7708263993263245,
+      "learning_rate": 8.263598326359832e-05,
+      "loss": 0.2874,
+      "step": 2130
+    },
+    {
+      "epoch": 164.64,
+      "grad_norm": 0.6249658465385437,
+      "learning_rate": 8.255230125523013e-05,
+      "loss": 0.2681,
+      "step": 2140
+    },
+    {
+      "epoch": 165.4,
+      "grad_norm": 0.4472973048686981,
+      "learning_rate": 8.246861924686193e-05,
+      "loss": 0.2753,
+      "step": 2150
+    },
+    {
+      "epoch": 166.16,
+      "grad_norm": 0.7961902022361755,
+      "learning_rate": 8.238493723849373e-05,
+      "loss": 0.2798,
+      "step": 2160
+    },
+    {
+      "epoch": 166.96,
+      "grad_norm": 0.48364782333374023,
+      "learning_rate": 8.230125523012552e-05,
+      "loss": 0.2924,
+      "step": 2170
+    },
+    {
+      "epoch": 167.72,
+      "grad_norm": 0.5379577279090881,
+      "learning_rate": 8.221757322175732e-05,
+      "loss": 0.2509,
+      "step": 2180
+    },
+    {
+      "epoch": 168.48,
+      "grad_norm": 0.7065776586532593,
+      "learning_rate": 8.213389121338912e-05,
+      "loss": 0.2763,
+      "step": 2190
+    },
+    {
+      "epoch": 169.24,
+      "grad_norm": 0.664776086807251,
+      "learning_rate": 8.205020920502092e-05,
+      "loss": 0.2726,
+      "step": 2200
+    },
+    {
+      "epoch": 170.0,
+      "grad_norm": 0.47445446252822876,
+      "learning_rate": 8.196652719665271e-05,
+      "loss": 0.2696,
+      "step": 2210
+    },
+    {
+      "epoch": 170.8,
+      "grad_norm": 0.5877684354782104,
+      "learning_rate": 8.188284518828451e-05,
+      "loss": 0.3092,
+      "step": 2220
+    },
+    {
+      "epoch": 171.56,
+      "grad_norm": 0.8296095728874207,
+      "learning_rate": 8.179916317991632e-05,
+      "loss": 0.2458,
+      "step": 2230
+    },
+    {
+      "epoch": 172.32,
+      "grad_norm": 0.2947175204753876,
+      "learning_rate": 8.171548117154812e-05,
+      "loss": 0.245,
+      "step": 2240
+    },
+    {
+      "epoch": 173.08,
+      "grad_norm": 0.6972278952598572,
+      "learning_rate": 8.163179916317992e-05,
+      "loss": 0.3018,
+      "step": 2250
+    },
+    {
+      "epoch": 173.88,
+      "grad_norm": 0.36442530155181885,
+      "learning_rate": 8.154811715481172e-05,
+      "loss": 0.2875,
+      "step": 2260
+    },
+    {
+      "epoch": 174.64,
+      "grad_norm": 0.38456207513809204,
+      "learning_rate": 8.146443514644351e-05,
+      "loss": 0.2328,
+      "step": 2270
+    },
+    {
+      "epoch": 175.4,
+      "grad_norm": 0.573784589767456,
+      "learning_rate": 8.138075313807531e-05,
+      "loss": 0.2851,
+      "step": 2280
+    },
+    {
+      "epoch": 176.16,
+      "grad_norm": 0.5698201656341553,
+      "learning_rate": 8.129707112970711e-05,
+      "loss": 0.2807,
+      "step": 2290
+    },
+    {
+      "epoch": 176.96,
+      "grad_norm": 0.7705904841423035,
+      "learning_rate": 8.121338912133891e-05,
+      "loss": 0.2982,
+      "step": 2300
+    },
+    {
+      "epoch": 177.72,
+      "grad_norm": 0.5071648359298706,
+      "learning_rate": 8.11297071129707e-05,
+      "loss": 0.2127,
+      "step": 2310
+    },
+    {
+      "epoch": 178.48,
+      "grad_norm": 1.354844331741333,
+      "learning_rate": 8.104602510460252e-05,
+      "loss": 0.2938,
+      "step": 2320
+    },
+    {
+      "epoch": 179.24,
+      "grad_norm": 0.4220433533191681,
+      "learning_rate": 8.096234309623431e-05,
+      "loss": 0.241,
+      "step": 2330
+    },
+    {
+      "epoch": 180.0,
+      "grad_norm": 0.4945012331008911,
+      "learning_rate": 8.087866108786611e-05,
+      "loss": 0.2482,
+      "step": 2340
+    },
+    {
+      "epoch": 180.8,
+      "grad_norm": 0.6846901774406433,
+      "learning_rate": 8.079497907949791e-05,
+      "loss": 0.2146,
+      "step": 2350
+    },
+    {
+      "epoch": 181.56,
+      "grad_norm": 0.6438813209533691,
+      "learning_rate": 8.07112970711297e-05,
+      "loss": 0.252,
+      "step": 2360
+    },
+    {
+      "epoch": 182.32,
+      "grad_norm": 0.486453652381897,
+      "learning_rate": 8.06276150627615e-05,
+      "loss": 0.2417,
+      "step": 2370
+    },
+    {
+      "epoch": 183.08,
+      "grad_norm": 0.5143964290618896,
+      "learning_rate": 8.05439330543933e-05,
+      "loss": 0.2882,
+      "step": 2380
+    },
+    {
+      "epoch": 183.88,
+      "grad_norm": 1.5306233167648315,
+      "learning_rate": 8.04602510460251e-05,
+      "loss": 0.2624,
+      "step": 2390
+    },
+    {
+      "epoch": 184.64,
+      "grad_norm": 0.47094181180000305,
+      "learning_rate": 8.037656903765691e-05,
+      "loss": 0.2009,
+      "step": 2400
+    },
+    {
+      "epoch": 185.4,
+      "grad_norm": 0.510780930519104,
+      "learning_rate": 8.029288702928871e-05,
+      "loss": 0.3383,
+      "step": 2410
+    },
+    {
+      "epoch": 186.16,
+      "grad_norm": 1.601507306098938,
+      "learning_rate": 8.02092050209205e-05,
+      "loss": 0.2143,
+      "step": 2420
+    },
+    {
+      "epoch": 186.96,
+      "grad_norm": 0.5631161332130432,
+      "learning_rate": 8.01255230125523e-05,
+      "loss": 0.2093,
+      "step": 2430
+    },
+    {
+      "epoch": 187.72,
+      "grad_norm": 0.2900833189487457,
+      "learning_rate": 8.00418410041841e-05,
+      "loss": 0.1871,
+      "step": 2440
+    },
+    {
+      "epoch": 188.48,
+      "grad_norm": 0.6894598007202148,
+      "learning_rate": 7.99581589958159e-05,
+      "loss": 0.2654,
+      "step": 2450
+    },
+    {
+      "epoch": 189.24,
+      "grad_norm": 8.970075607299805,
+      "learning_rate": 7.98744769874477e-05,
+      "loss": 0.178,
+      "step": 2460
+    },
+    {
+      "epoch": 190.0,
+      "grad_norm": 0.2863423228263855,
+      "learning_rate": 7.97907949790795e-05,
+      "loss": 0.1958,
+      "step": 2470
+    },
+    {
+      "epoch": 190.8,
+      "grad_norm": 0.3206811547279358,
+      "learning_rate": 7.97071129707113e-05,
+      "loss": 0.2548,
+      "step": 2480
+    },
+    {
+      "epoch": 191.56,
+      "grad_norm": 0.6876190900802612,
+      "learning_rate": 7.96234309623431e-05,
+      "loss": 0.1843,
+      "step": 2490
+    },
+    {
+      "epoch": 192.32,
+      "grad_norm": 1.09211266040802,
+      "learning_rate": 7.95397489539749e-05,
+      "loss": 0.1721,
+      "step": 2500
+    },
+    {
+      "epoch": 193.08,
+      "grad_norm": 14.786331176757812,
+      "learning_rate": 7.94560669456067e-05,
+      "loss": 0.327,
+      "step": 2510
+    },
+    {
+      "epoch": 193.88,
+      "grad_norm": 43.85841369628906,
+      "learning_rate": 7.93723849372385e-05,
+      "loss": 0.3656,
+      "step": 2520
+    },
+    {
+      "epoch": 194.64,
+      "grad_norm": 0.32373273372650146,
+      "learning_rate": 7.92887029288703e-05,
+      "loss": 0.2476,
+      "step": 2530
+    },
+    {
+      "epoch": 195.4,
+      "grad_norm": 0.00038791695260442793,
+      "learning_rate": 7.920502092050209e-05,
+      "loss": 0.1761,
+      "step": 2540
+    },
+    {
+      "epoch": 196.16,
+      "grad_norm": 1.776816487312317,
+      "learning_rate": 7.912133891213389e-05,
+      "loss": 0.1828,
+      "step": 2550
+    },
+    {
+      "epoch": 196.96,
+      "grad_norm": 0.5436874628067017,
+      "learning_rate": 7.903765690376569e-05,
+      "loss": 0.2261,
+      "step": 2560
+    },
+    {
+      "epoch": 197.72,
+      "grad_norm": 0.45118626952171326,
+      "learning_rate": 7.89539748953975e-05,
+      "loss": 0.1712,
+      "step": 2570
+    },
+    {
+      "epoch": 198.48,
+      "grad_norm": 3.249994993209839,
+      "learning_rate": 7.88702928870293e-05,
+      "loss": 0.247,
+      "step": 2580
+    },
+    {
+      "epoch": 199.24,
+      "grad_norm": 1.1368451118469238,
+      "learning_rate": 7.878661087866109e-05,
+      "loss": 0.2208,
+      "step": 2590
+    },
+    {
+      "epoch": 200.0,
+      "grad_norm": 0.6388035416603088,
+      "learning_rate": 7.870292887029289e-05,
+      "loss": 0.1518,
+      "step": 2600
+    },
+    {
+      "epoch": 200.8,
+      "grad_norm": 3.905496597290039,
+      "learning_rate": 7.861924686192469e-05,
+      "loss": 0.1279,
+      "step": 2610
+    },
+    {
+      "epoch": 201.56,
+      "grad_norm": 27.926372528076172,
+      "learning_rate": 7.853556485355649e-05,
+      "loss": 0.4156,
+      "step": 2620
+    },
+    {
+      "epoch": 202.32,
+      "grad_norm": 0.19806312024593353,
+      "learning_rate": 7.845188284518828e-05,
+      "loss": 0.0903,
+      "step": 2630
+    },
+    {
+      "epoch": 203.08,
+      "grad_norm": 0.0014255548594519496,
+      "learning_rate": 7.836820083682008e-05,
+      "loss": 0.1962,
+      "step": 2640
+    },
+    {
+      "epoch": 203.88,
+      "grad_norm": 2.5122647285461426,
+      "learning_rate": 7.828451882845189e-05,
+      "loss": 0.1382,
+      "step": 2650
+    },
+    {
+      "epoch": 204.64,
+      "grad_norm": 5.140716075897217,
+      "learning_rate": 7.820083682008369e-05,
+      "loss": 0.2123,
+      "step": 2660
+    },
+    {
+      "epoch": 205.4,
+      "grad_norm": 1.3244779109954834,
+      "learning_rate": 7.811715481171549e-05,
+      "loss": 0.1756,
+      "step": 2670
+    },
+    {
+      "epoch": 206.16,
+      "grad_norm": 26.31413459777832,
+      "learning_rate": 7.803347280334728e-05,
+      "loss": 0.1425,
+      "step": 2680
+    },
+    {
+      "epoch": 206.96,
+      "grad_norm": 1.5801438093185425,
+      "learning_rate": 7.794979079497908e-05,
+      "loss": 0.1822,
+      "step": 2690
+    },
+    {
+      "epoch": 207.72,
+      "grad_norm": 1.6209617853164673,
+      "learning_rate": 7.786610878661088e-05,
+      "loss": 0.1909,
+      "step": 2700
+    },
+    {
+      "epoch": 208.48,
+      "grad_norm": 0.6468285918235779,
+      "learning_rate": 7.778242677824268e-05,
+      "loss": 0.1741,
+      "step": 2710
+    },
+    {
+      "epoch": 209.24,
+      "grad_norm": 0.04982404410839081,
+      "learning_rate": 7.769874476987448e-05,
+      "loss": 0.1027,
+      "step": 2720
+    },
+    {
+      "epoch": 210.0,
+      "grad_norm": 0.0010119529906660318,
+      "learning_rate": 7.761506276150629e-05,
+      "loss": 0.1457,
+      "step": 2730
+    },
+    {
+      "epoch": 210.8,
+      "grad_norm": 38.17147445678711,
+      "learning_rate": 7.753138075313808e-05,
+      "loss": 0.164,
+      "step": 2740
+    },
+    {
+      "epoch": 211.56,
+      "grad_norm": 1.5547058582305908,
+      "learning_rate": 7.744769874476988e-05,
+      "loss": 0.1102,
+      "step": 2750
+    },
+    {
+      "epoch": 212.32,
+      "grad_norm": 0.011813919991254807,
+      "learning_rate": 7.736401673640168e-05,
+      "loss": 0.1455,
+      "step": 2760
+    },
+    {
+      "epoch": 213.08,
+      "grad_norm": 0.007001452147960663,
+      "learning_rate": 7.728033472803348e-05,
+      "loss": 0.1694,
+      "step": 2770
+    },
+    {
+      "epoch": 213.88,
+      "grad_norm": 0.7564431428909302,
+      "learning_rate": 7.719665271966527e-05,
+      "loss": 0.2061,
+      "step": 2780
+    },
+    {
+      "epoch": 214.64,
+      "grad_norm": 1.1208900213241577,
+      "learning_rate": 7.711297071129707e-05,
+      "loss": 0.2951,
+      "step": 2790
+    },
+    {
+      "epoch": 215.4,
+      "grad_norm": 0.4398239552974701,
+      "learning_rate": 7.702928870292887e-05,
+      "loss": 0.1161,
+      "step": 2800
+    },
+    {
+      "epoch": 216.16,
+      "grad_norm": 0.34080252051353455,
+      "learning_rate": 7.694560669456067e-05,
+      "loss": 0.1177,
+      "step": 2810
+    },
+    {
+      "epoch": 216.96,
+      "grad_norm": 0.27304738759994507,
+      "learning_rate": 7.686192468619248e-05,
+      "loss": 0.117,
+      "step": 2820
+    },
+    {
+      "epoch": 217.72,
+      "grad_norm": 0.00020999423577450216,
+      "learning_rate": 7.677824267782428e-05,
+      "loss": 0.1518,
+      "step": 2830
+    },
+    {
+      "epoch": 218.48,
+      "grad_norm": 0.3551616072654724,
+      "learning_rate": 7.669456066945607e-05,
+      "loss": 0.0911,
+      "step": 2840
+    },
+    {
+      "epoch": 219.24,
+      "grad_norm": 18.045923233032227,
+      "learning_rate": 7.661087866108787e-05,
+      "loss": 0.1732,
+      "step": 2850
+    },
+    {
+      "epoch": 220.0,
+      "grad_norm": 0.0011837411439046264,
+      "learning_rate": 7.652719665271967e-05,
+      "loss": 0.3062,
+      "step": 2860
+    },
+    {
+      "epoch": 220.8,
+      "grad_norm": 7.931725025177002,
+      "learning_rate": 7.644351464435147e-05,
+      "loss": 0.1608,
+      "step": 2870
+    },
+    {
+      "epoch": 221.56,
+      "grad_norm": 0.7454537749290466,
+      "learning_rate": 7.635983263598326e-05,
+      "loss": 0.163,
+      "step": 2880
+    },
+    {
+      "epoch": 222.32,
+      "grad_norm": 6.1156840324401855,
+      "learning_rate": 7.627615062761506e-05,
+      "loss": 0.0693,
+      "step": 2890
+    },
+    {
+      "epoch": 223.08,
+      "grad_norm": 1.552299976348877,
+      "learning_rate": 7.619246861924687e-05,
+      "loss": 0.1688,
+      "step": 2900
+    },
+    {
+      "epoch": 223.88,
+      "grad_norm": 11.509157180786133,
+      "learning_rate": 7.610878661087867e-05,
+      "loss": 0.1652,
+      "step": 2910
+    },
+    {
+      "epoch": 224.64,
+      "grad_norm": 0.007240507751703262,
+      "learning_rate": 7.602510460251047e-05,
+      "loss": 0.1246,
+      "step": 2920
+    },
+    {
+      "epoch": 225.4,
+      "grad_norm": 0.2530362010002136,
+      "learning_rate": 7.594142259414227e-05,
+      "loss": 0.0838,
+      "step": 2930
+    },
+    {
+      "epoch": 226.16,
+      "grad_norm": 0.5732694864273071,
+      "learning_rate": 7.585774058577406e-05,
+      "loss": 0.1891,
+      "step": 2940
+    },
+    {
+      "epoch": 226.96,
+      "grad_norm": 3.3265058994293213,
+      "learning_rate": 7.577405857740586e-05,
+      "loss": 0.219,
+      "step": 2950
+    },
+    {
+      "epoch": 227.72,
+      "grad_norm": 19.179094314575195,
+      "learning_rate": 7.569037656903766e-05,
+      "loss": 0.1078,
+      "step": 2960
+    },
+    {
+      "epoch": 228.48,
+      "grad_norm": 0.5755670666694641,
+      "learning_rate": 7.560669456066946e-05,
+      "loss": 0.1628,
+      "step": 2970
+    },
+    {
+      "epoch": 229.24,
+      "grad_norm": 13.489233016967773,
+      "learning_rate": 7.552301255230127e-05,
+      "loss": 0.1031,
+      "step": 2980
+    },
+    {
+      "epoch": 230.0,
+      "grad_norm": 0.42450758814811707,
+      "learning_rate": 7.543933054393307e-05,
+      "loss": 0.1062,
+      "step": 2990
+    },
+    {
+      "epoch": 230.8,
+      "grad_norm": 74.2090835571289,
+      "learning_rate": 7.535564853556486e-05,
+      "loss": 0.2763,
+      "step": 3000
+    },
+    {
+      "epoch": 231.56,
+      "grad_norm": 0.0013025101507082582,
+      "learning_rate": 7.527196652719666e-05,
+      "loss": 0.0475,
+      "step": 3010
+    },
+    {
+      "epoch": 232.32,
+      "grad_norm": 0.36847984790802,
+      "learning_rate": 7.518828451882846e-05,
+      "loss": 0.1262,
+      "step": 3020
+    },
+    {
+      "epoch": 233.08,
+      "grad_norm": 2.1009180545806885,
+      "learning_rate": 7.510460251046026e-05,
+      "loss": 0.2701,
+      "step": 3030
+    },
+    {
+      "epoch": 233.88,
+      "grad_norm": 3.9676990509033203,
+      "learning_rate": 7.502092050209205e-05,
+      "loss": 0.0567,
+      "step": 3040
+    },
+    {
+      "epoch": 234.64,
+      "grad_norm": 2.426058769226074,
+      "learning_rate": 7.493723849372385e-05,
+      "loss": 0.1485,
+      "step": 3050
+    },
+    {
+      "epoch": 235.4,
+      "grad_norm": 0.45143410563468933,
+      "learning_rate": 7.485355648535565e-05,
+      "loss": 0.0808,
+      "step": 3060
+    },
+    {
+      "epoch": 236.16,
+      "grad_norm": 0.001772599876858294,
+      "learning_rate": 7.476987447698746e-05,
+      "loss": 0.159,
+      "step": 3070
+    },
+    {
+      "epoch": 236.96,
+      "grad_norm": 0.6228379011154175,
+      "learning_rate": 7.468619246861926e-05,
+      "loss": 0.3161,
+      "step": 3080
+    },
+    {
+      "epoch": 237.72,
+      "grad_norm": 1.1277211904525757,
+      "learning_rate": 7.460251046025106e-05,
+      "loss": 0.0895,
+      "step": 3090
+    },
+    {
+      "epoch": 238.48,
+      "grad_norm": 0.5020686388015747,
+      "learning_rate": 7.451882845188285e-05,
+      "loss": 0.066,
+      "step": 3100
+    },
+    {
+      "epoch": 239.24,
+      "grad_norm": 0.4364977777004242,
+      "learning_rate": 7.443514644351465e-05,
+      "loss": 0.0973,
+      "step": 3110
+    },
+    {
+      "epoch": 240.0,
+      "grad_norm": 0.331858366727829,
+      "learning_rate": 7.435146443514645e-05,
+      "loss": 0.1385,
+      "step": 3120
+    },
+    {
+      "epoch": 240.8,
+      "grad_norm": 3.562025308609009,
+      "learning_rate": 7.426778242677825e-05,
+      "loss": 0.1554,
+      "step": 3130
+    },
+    {
+      "epoch": 241.56,
+      "grad_norm": 0.4786972403526306,
+      "learning_rate": 7.418410041841004e-05,
+      "loss": 0.1127,
+      "step": 3140
+    },
+    {
+      "epoch": 242.32,
+      "grad_norm": 0.22106488049030304,
+      "learning_rate": 7.410041841004186e-05,
+      "loss": 0.1212,
+      "step": 3150
+    },
+    {
+      "epoch": 243.08,
+      "grad_norm": 0.3226916193962097,
+      "learning_rate": 7.401673640167365e-05,
+      "loss": 0.0622,
+      "step": 3160
+    },
+    {
+      "epoch": 243.88,
+      "grad_norm": 0.18059372901916504,
+      "learning_rate": 7.393305439330545e-05,
+      "loss": 0.0578,
+      "step": 3170
+    },
+    {
+      "epoch": 244.64,
+      "grad_norm": 0.18617404997348785,
+      "learning_rate": 7.384937238493725e-05,
+      "loss": 0.1017,
+      "step": 3180
+    },
+    {
+      "epoch": 245.4,
+      "grad_norm": 0.4568879008293152,
+      "learning_rate": 7.376569037656905e-05,
+      "loss": 0.0456,
+      "step": 3190
+    },
+    {
+      "epoch": 246.16,
+      "grad_norm": 0.00673318887129426,
+      "learning_rate": 7.368200836820084e-05,
+      "loss": 0.0578,
+      "step": 3200
+    },
+    {
+      "epoch": 246.96,
+      "grad_norm": 0.0004931857693009079,
+      "learning_rate": 7.359832635983264e-05,
+      "loss": 0.0363,
+      "step": 3210
+    },
+    {
+      "epoch": 247.72,
+      "grad_norm": 0.18300887942314148,
+      "learning_rate": 7.351464435146444e-05,
+      "loss": 0.0854,
+      "step": 3220
+    },
+    {
+      "epoch": 248.48,
+      "grad_norm": 0.003890759777277708,
+      "learning_rate": 7.343096234309624e-05,
+      "loss": 0.1226,
+      "step": 3230
+    },
+    {
+      "epoch": 249.24,
+      "grad_norm": 1.134564757347107,
+      "learning_rate": 7.334728033472805e-05,
+      "loss": 0.0956,
+      "step": 3240
+    },
+    {
+      "epoch": 250.0,
+      "grad_norm": 0.00022713415091857314,
+      "learning_rate": 7.326359832635985e-05,
+      "loss": 0.1,
+      "step": 3250
+    },
+    {
+      "epoch": 250.8,
+      "grad_norm": 7.879546165466309,
+      "learning_rate": 7.317991631799164e-05,
+      "loss": 0.0877,
+      "step": 3260
+    },
+    {
+      "epoch": 251.56,
+      "grad_norm": 2.57737135887146,
+      "learning_rate": 7.309623430962344e-05,
+      "loss": 0.09,
+      "step": 3270
+    },
+    {
+      "epoch": 252.32,
+      "grad_norm": 0.4551700949668884,
+      "learning_rate": 7.301255230125524e-05,
+      "loss": 0.1288,
+      "step": 3280
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 12000,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1000,
+  "save_steps": 10,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 1.8125347399735104e+17,
+  "train_batch_size": 1,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8cef66ef127754c375a8e9b549e2244a7cc084c0396c35d5511163281be21d7
+size 5507

vision_siglip_navit.py

@@ -0,0 +1,1717 @@
+# coding=utf-8
+# Copyright 2024 The 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.
+""" Siglip model configuration"""
+
+import os
+from typing import Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+SIGLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "google/siglip-base-patch16-224": "https://huggingface.co/google/siglip-base-patch16-224/resolve/main/config.json",
+}
+
+
+class SiglipTextConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`SiglipTextModel`]. It is used to instantiate a
+    Siglip text encoder according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the text encoder of the Siglip
+    [google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) architecture.
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the Siglip text model. Defines the number of different tokens that can be represented by
+            the `inputs_ids` passed when calling [`SiglipModel`].
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        max_position_embeddings (`int`, *optional*, defaults to 64):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        pad_token_id (`int`, *optional*, defaults to 1):
+            The id of the padding token in the vocabulary.
+        bos_token_id (`int`, *optional*, defaults to 49406):
+            The id of the beginning-of-sequence token in the vocabulary.
+        eos_token_id (`int`, *optional*, defaults to 49407):
+            The id of the end-of-sequence token in the vocabulary.
+    Example:
+    ```python
+    >>> from transformers import SiglipTextConfig, SiglipTextModel
+    >>> # Initializing a SiglipTextConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipTextConfig()
+    >>> # Initializing a SiglipTextModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipTextModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "siglip_text_model"
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        max_position_embeddings=64,
+        hidden_act="gelu_pytorch_tanh",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        # This differs from `CLIPTokenizer`'s default and from openai/siglip
+        # See https://github.com/huggingface/transformers/pull/24773#issuecomment-1632287538
+        pad_token_id=1,
+        bos_token_id=49406,
+        eos_token_id=49407,
+        _flash_attn_2_enabled=True,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.attention_dropout = attention_dropout
+        self._flash_attn_2_enabled = _flash_attn_2_enabled
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the text config dict if we are loading from SiglipConfig
+        if config_dict.get("model_type") == "siglip":
+            config_dict = config_dict["text_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class SiglipVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`SiglipVisionModel`]. It is used to instantiate a
+    Siglip vision encoder according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the vision encoder of the Siglip
+    [google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) architecture.
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_channels (`int`, *optional*, defaults to 3):
+            Number of channels in the input images.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        patch_size (`int`, *optional*, defaults to 16):
+            The size (resolution) of each patch.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+    Example:
+    ```python
+    >>> from transformers import SiglipVisionConfig, SiglipVisionModel
+    >>> # Initializing a SiglipVisionConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipVisionConfig()
+    >>> # Initializing a SiglipVisionModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipVisionModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "siglip_vision_model"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        image_size=224,
+        patch_size=16,
+        hidden_act="gelu_pytorch_tanh",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        _flash_attn_2_enabled=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self._flash_attn_2_enabled = _flash_attn_2_enabled
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        cls._set_token_in_kwargs(kwargs)
+
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the vision config dict if we are loading from SiglipConfig
+        if config_dict.get("model_type") == "siglip":
+            config_dict = config_dict["vision_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class SiglipConfig(PretrainedConfig):
+    r"""
+    [`SiglipConfig`] is the configuration class to store the configuration of a [`SiglipModel`]. It is used to
+    instantiate a Siglip model according to the specified arguments, defining the text model and vision model configs.
+    Instantiating a configuration with the defaults will yield a similar configuration to that of the Siglip
+    [google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) architecture.
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        text_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`SiglipTextConfig`].
+        vision_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`SiglipVisionConfig`].
+        kwargs (*optional*):
+            Dictionary of keyword arguments.
+    Example:
+    ```python
+    >>> from transformers import SiglipConfig, SiglipModel
+    >>> # Initializing a SiglipConfig with google/siglip-base-patch16-224 style configuration
+    >>> configuration = SiglipConfig()
+    >>> # Initializing a SiglipModel (with random weights) from the google/siglip-base-patch16-224 style configuration
+    >>> model = SiglipModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    >>> # We can also initialize a SiglipConfig from a SiglipTextConfig and a SiglipVisionConfig
+    >>> from transformers import SiglipTextConfig, SiglipVisionConfig
+    >>> # Initializing a SiglipText and SiglipVision configuration
+    >>> config_text = SiglipTextConfig()
+    >>> config_vision = SiglipVisionConfig()
+    >>> config = SiglipConfig.from_text_vision_configs(config_text, config_vision)
+    ```"""
+
+    model_type = "siglip"
+
+    def __init__(self, text_config=None, vision_config=None, **kwargs):
+        super().__init__(**kwargs)
+
+        if text_config is None:
+            text_config = {}
+            logger.info("`text_config` is `None`. Initializing the `SiglipTextConfig` with default values.")
+
+        if vision_config is None:
+            vision_config = {}
+            logger.info("`vision_config` is `None`. initializing the `SiglipVisionConfig` with default values.")
+
+        self.text_config = SiglipTextConfig(**text_config)
+        self.vision_config = SiglipVisionConfig(**vision_config)
+
+        self.initializer_factor = 1.0
+
+    @classmethod
+    def from_text_vision_configs(cls, text_config: SiglipTextConfig, vision_config: SiglipVisionConfig, **kwargs):
+        r"""
+        Instantiate a [`SiglipConfig`] (or a derived class) from siglip text model configuration and siglip vision
+        model configuration.
+        Returns:
+            [`SiglipConfig`]: An instance of a configuration object
+        """
+
+        return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)
+
+# coding=utf-8
+# Copyright 2024 Google AI and 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.
+""" PyTorch Siglip model."""
+
+
+import math
+import warnings
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn.init import _calculate_fan_in_and_fan_out
+
+from transformers.activations import ACT2FN
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    ModelOutput,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    is_flash_attn_2_available,
+    logging,
+    replace_return_docstrings,
+)
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "google/siglip-base-patch16-224"
+
+SIGLIP_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "google/siglip-base-patch16-224",
+    # See all SigLIP models at https://huggingface.co/models?filter=siglip
+]
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+def _trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    # Values are generated by using a truncated uniform distribution and
+    # then using the inverse CDF for the normal distribution.
+    # Get upper and lower cdf values
+    l = norm_cdf((a - mean) / std)
+    u = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [l, u], then translate to
+    # [2l-1, 2u-1].
+    tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated
+    # standard normal
+    if tensor.dtype in [torch.float16, torch.bfloat16]:
+        # The `erfinv_` op is not (yet?) defined in float16+cpu, bfloat16+gpu
+        og_dtype = tensor.dtype
+        tensor = tensor.to(torch.float32)
+        tensor.erfinv_()
+        tensor = tensor.to(og_dtype)
+    else:
+        tensor.erfinv_()
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.0))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    if tensor.dtype == torch.float16:
+        # The `clamp_` op is not (yet?) defined in float16+cpu
+        tensor = tensor.to(torch.float32)
+        tensor.clamp_(min=a, max=b)
+        tensor = tensor.to(torch.float16)
+    else:
+        tensor.clamp_(min=a, max=b)
+
+
+def trunc_normal_tf_(
+    tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0
+) -> torch.Tensor:
+    """Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \\leq \text{mean} \\leq b`.
+    NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
+    bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
+    and the result is subsquently scaled and shifted by the mean and std args.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    """
+    with torch.no_grad():
+        _trunc_normal_(tensor, 0, 1.0, a, b)
+        tensor.mul_(std).add_(mean)
+
+
+def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    if mode == "fan_in":
+        denom = fan_in
+    elif mode == "fan_out":
+        denom = fan_out
+    elif mode == "fan_avg":
+        denom = (fan_in + fan_out) / 2
+
+    variance = scale / denom
+
+    if distribution == "truncated_normal":
+        # constant is stddev of standard normal truncated to (-2, 2)
+        trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
+    elif distribution == "normal":
+        with torch.no_grad():
+            tensor.normal_(std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        with torch.no_grad():
+            tensor.uniform_(-bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")
+
+
+def default_flax_embed_init(tensor):
+    variance_scaling_(tensor, mode="fan_in", distribution="normal")
+
+
+@dataclass
+# Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with CLIP->Siglip
+class SiglipVisionModelOutput(ModelOutput):
+    """
+    Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.
+    Args:
+        image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
+            The image embeddings obtained by applying the projection layer to the pooler_output.
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    image_embeds: Optional[torch.FloatTensor] = None
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+# Copied from transformers.models.clip.modeling_clip.CLIPTextModelOutput with CLIP->Siglip
+class SiglipTextModelOutput(ModelOutput):
+    """
+    Base class for text model's outputs that also contains a pooling of the last hidden states.
+    Args:
+        text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
+            The text embeddings obtained by applying the projection layer to the pooler_output.
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    text_embeds: Optional[torch.FloatTensor] = None
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+# Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->Siglip
+class SiglipOutput(ModelOutput):
+    """
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
+            Contrastive loss for image-text similarity.
+        logits_per_image:(`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
+            The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text
+            similarity scores.
+        logits_per_text:(`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
+            The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image
+            similarity scores.
+        text_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`):
+            The text embeddings obtained by applying the projection layer to the pooled output of [`SiglipTextModel`].
+        image_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`):
+            The image embeddings obtained by applying the projection layer to the pooled output of [`SiglipVisionModel`].
+        text_model_output(`BaseModelOutputWithPooling`):
+            The output of the [`SiglipTextModel`].
+        vision_model_output(`BaseModelOutputWithPooling`):
+            The output of the [`SiglipVisionModel`].
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits_per_image: torch.FloatTensor = None
+    logits_per_text: torch.FloatTensor = None
+    text_embeds: torch.FloatTensor = None
+    image_embeds: torch.FloatTensor = None
+    text_model_output: BaseModelOutputWithPooling = None
+    vision_model_output: BaseModelOutputWithPooling = None
+
+    def to_tuple(self) -> Tuple[Any]:
+        return tuple(
+            self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
+            for k in self.keys()
+        )
+
+
+class SiglipVisionEmbeddings(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",
+        )
+
+        self.num_patches_per_side = self.image_size // self.patch_size
+        self.num_patches = self.num_patches_per_side**2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+
+    def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.BoolTensor) -> torch.Tensor:
+        batch_size = pixel_values.size(0)
+
+        patch_embeds = self.patch_embedding(pixel_values)
+        embeddings = patch_embeds.flatten(2).transpose(1, 2)
+
+        max_im_h, max_im_w = pixel_values.size(2), pixel_values.size(3)
+        max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size
+        boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side)
+        position_ids = torch.full(
+            size=(
+                batch_size,
+                max_nb_patches_h * max_nb_patches_w,
+            ),
+            fill_value=0,
+        )
+
+        for batch_idx, p_attn_mask in enumerate(patch_attention_mask):
+            nb_patches_h = p_attn_mask[:, 0].sum()
+            nb_patches_w = p_attn_mask[0].sum()
+
+            fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h)
+            fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w)
+
+            bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)
+            bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)
+
+            pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()
+            position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids
+
+        position_ids = position_ids.to(self.position_embedding.weight.device)
+
+        embeddings = embeddings + self.position_embedding(position_ids)
+        return embeddings
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPTextEmbeddings with CLIP->Siglip
+class SiglipTextEmbeddings(nn.Module):
+    def __init__(self, config: SiglipTextConfig):
+        super().__init__()
+        embed_dim = config.hidden_size
+
+        self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
+        self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer(
+            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+    ) -> torch.Tensor:
+        seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.token_embedding(input_ids)
+
+        position_embeddings = self.position_embedding(position_ids)
+        embeddings = inputs_embeds + position_embeddings
+
+        return embeddings
+
+
+class SiglipAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        batch_size, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        k_v_seq_len = key_states.shape[-2]
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
+
+        if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class SiglipFlashAttention2(SiglipAttention):
+    """
+    Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.is_causal = False  # Hack to make sure we don't use a causal mask
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        # cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        # query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        # if past_key_value is not None:
+        #     cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+        #     key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        dropout_rate = self.dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                "The input hidden states seems to be silently casted in float32, this might be related to the fact"
+                " you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.embed_dim).contiguous()
+        attn_output = self.out_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights
+
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`int`, *optional*):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+
+        # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+        causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip
+class SiglipMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->Siglip
+class SiglipEncoderLayer(nn.Module):
+    def __init__(self, config: SiglipConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = (
+            SiglipAttention(config)
+            if not getattr(config, "_flash_attn_2_enabled", False)
+            else SiglipFlashAttention2(config)
+        )
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            attention_mask (`torch.FloatTensor`):
+                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class SiglipPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = SiglipConfig
+    base_model_prefix = "siglip"
+    supports_gradient_checkpointing = True
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+
+        if isinstance(module, SiglipVisionEmbeddings):
+            width = (
+                self.config.vision_config.hidden_size
+                if isinstance(self.config, SiglipConfig)
+                else self.config.hidden_size
+            )
+            nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width))
+        elif isinstance(module, nn.Embedding):
+            default_flax_embed_init(module.weight)
+        elif isinstance(module, SiglipAttention):
+            nn.init.normal_(module.q_proj.weight)
+            nn.init.normal_(module.k_proj.weight)
+            nn.init.normal_(module.v_proj.weight)
+            nn.init.normal_(module.out_proj.weight)
+            nn.init.zeros_(module.q_proj.bias)
+            nn.init.zeros_(module.k_proj.bias)
+            nn.init.zeros_(module.v_proj.bias)
+            nn.init.zeros_(module.out_proj.bias)
+        elif isinstance(module, SiglipMLP):
+            nn.init.normal_(module.fc1.weight)
+            nn.init.normal_(module.fc2.weight)
+            nn.init.normal_(module.fc1.bias, std=1e-6)
+            nn.init.normal_(module.fc2.bias, std=1e-6)
+        elif isinstance(module, SiglipMultiheadAttentionPoolingHead):
+            nn.init.normal_(module.probe.data)
+            nn.init.normal_(module.attention.in_proj_weight.data)
+            nn.init.zeros_(module.attention.in_proj_bias.data)
+        elif isinstance(module, SiglipModel):
+            logit_scale_init = torch.tensor(0.0)
+            module.logit_scale.data.fill_(logit_scale_init)
+            module.logit_bias.data.zero_()
+        elif isinstance(module, (nn.Linear, nn.Conv2d)):
+            lecun_normal_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+SIGLIP_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+    Parameters:
+        config ([`SiglipConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+SIGLIP_TEXT_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+            [What are attention masks?](../glossary#attention-mask)
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.max_position_embeddings - 1]`.
+            [What are position IDs?](../glossary#position-ids)
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+SIGLIP_VISION_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
+            [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+SIGLIP_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+            [What are attention masks?](../glossary#attention-mask)
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.max_position_embeddings - 1]`.
+            [What are position IDs?](../glossary#position-ids)
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
+            [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
+        return_loss (`bool`, *optional*):
+            Whether or not to return the contrastive loss.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->Siglip
+class SiglipEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`SiglipEncoderLayer`].
+    Args:
+        config: SiglipConfig
+    """
+
+    def __init__(self, config: SiglipConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    # Ignore copy
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+                [What are attention masks?](../glossary#attention-mask)
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        hidden_states = inputs_embeds
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    encoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    output_attentions,
+                )
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                    attention_mask,
+                    output_attentions=output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
+        )
+
+
+class SiglipTextTransformer(nn.Module):
+    def __init__(self, config: SiglipTextConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+        self.embeddings = SiglipTextEmbeddings(config)
+        self.encoder = SiglipEncoder(config)
+        self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+        self.head = nn.Linear(embed_dim, embed_dim)
+
+    @add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipTextConfig)
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is None:
+            raise ValueError("You have to specify input_ids")
+
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
+
+        # note: SigLIP's text model does not use a causal mask, unlike the original CLIP model.
+        # expand attention_mask
+        if attention_mask is not None:
+            # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+        # Assuming "sticky" EOS tokenization, last token is always EOS.
+        pooled_output = last_hidden_state[:, -1, :]
+        pooled_output = self.head(pooled_output)
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """The text model from SigLIP without any head or projection on top.""",
+    SIGLIP_START_DOCSTRING,
+)
+class SiglipTextModel(SiglipPreTrainedModel):
+    config_class = SiglipTextConfig
+
+    _no_split_modules = ["SiglipTextEmbeddings", "SiglipEncoderLayer"]
+
+    def __init__(self, config: SiglipTextConfig):
+        super().__init__(config)
+        self.text_model = SiglipTextTransformer(config)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.text_model.embeddings.token_embedding
+
+    def set_input_embeddings(self, value):
+        self.text_model.embeddings.token_embedding = value
+
+    @add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipTextConfig)
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+        Examples:
+        ```python
+        >>> from transformers import AutoTokenizer, SiglipTextModel
+        >>> model = SiglipTextModel.from_pretrained("google/siglip-base-patch16-224")
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224")
+        >>> # important: make sure to set padding="max_length" as that's how the model was trained
+        >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt")
+        >>> outputs = model(**inputs)
+        >>> last_hidden_state = outputs.last_hidden_state
+        >>> pooled_output = outputs.pooler_output  # pooled (EOS token) states
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        return self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+
+class SiglipVisionTransformer(nn.Module):
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = SiglipVisionEmbeddings(config)
+        self.encoder = SiglipEncoder(config)
+        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self.head = SiglipMultiheadAttentionPoolingHead(config)
+
+    @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipVisionConfig)
+    def forward(
+        self,
+        pixel_values,
+        patch_attention_mask: Optional[torch.BoolTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = pixel_values.size(0)
+        if patch_attention_mask is None:
+            patch_attention_mask = torch.ones(
+                size=(
+                    batch_size,
+                    pixel_values.size(2) // self.config.patch_size,
+                    pixel_values.size(3) // self.config.patch_size,
+                ),
+                dtype=torch.bool,
+                device=pixel_values.device,
+            )
+
+        hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)
+
+        patch_attention_mask = patch_attention_mask.view(batch_size, -1)
+        # The call to `_upad_input` in `_flash_attention_forward` is expensive
+        # So when the `patch_attention_mask` is full of 1s (i.e. attending to the whole sequence),
+        # avoiding passing the attention_mask, which is equivalent to attending to the full sequence
+        if not torch.any(~patch_attention_mask):
+            attention_mask=None
+        else:
+            attention_mask = (
+                _prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype)
+                if not self.config._flash_attn_2_enabled
+                else patch_attention_mask
+            )
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.post_layernorm(last_hidden_state)
+
+        pooled_output = self.head(
+            hidden_state=last_hidden_state,
+            attention_mask=patch_attention_mask,
+        )
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+class SiglipMultiheadAttentionPoolingHead(nn.Module):
+    """Multihead Attention Pooling."""
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__()
+
+        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+
+    def forward(self, hidden_state, attention_mask):
+        batch_size = hidden_state.shape[0]
+        probe = self.probe.repeat(batch_size, 1, 1)
+
+        hidden_state = self.attention(
+            query=probe, key=hidden_state, value=hidden_state, key_padding_mask=~attention_mask
+        )[0]
+
+        residual = hidden_state
+        hidden_state = self.layernorm(hidden_state)
+        hidden_state = residual + self.mlp(hidden_state)
+
+        return hidden_state[:, 0]
+
+
+@add_start_docstrings(
+    """The vision model from SigLIP without any head or projection on top.""",
+    SIGLIP_START_DOCSTRING,
+)
+class SiglipVisionModel(SiglipPreTrainedModel):
+    config_class = SiglipVisionConfig
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: SiglipVisionConfig):
+        super().__init__(config)
+
+        self.vision_model = SiglipVisionTransformer(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.vision_model.embeddings.patch_embedding
+
+    @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipVisionConfig)
+    def forward(
+        self,
+        pixel_values,
+        patch_attention_mask: Optional[torch.BoolTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+        Examples:
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, SiglipVisionModel
+        >>> model = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224")
+        >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+        >>> inputs = processor(images=image, return_tensors="pt")
+        >>> outputs = model(**inputs)
+        >>> last_hidden_state = outputs.last_hidden_state
+        >>> pooled_output = outputs.pooler_output  # pooled features
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        return self.vision_model(
+            pixel_values=pixel_values,
+            patch_attention_mask=patch_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+
+@add_start_docstrings(SIGLIP_START_DOCSTRING)
+class SiglipModel(SiglipPreTrainedModel):
+    config_class = SiglipConfig
+
+    def __init__(self, config: SiglipConfig):
+        super().__init__(config)
+
+        if not isinstance(config.text_config, SiglipTextConfig):
+            raise ValueError(
+                "config.text_config is expected to be of type SiglipTextConfig but is of type"
+                f" {type(config.text_config)}."
+            )
+
+        if not isinstance(config.vision_config, SiglipVisionConfig):
+            raise ValueError(
+                "config.vision_config is expected to be of type SiglipVisionConfig but is of type"
+                f" {type(config.vision_config)}."
+            )
+
+        text_config = config.text_config
+        vision_config = config.vision_config
+
+        self.text_model = SiglipTextTransformer(text_config)
+        self.vision_model = SiglipVisionTransformer(vision_config)
+
+        self.logit_scale = nn.Parameter(torch.randn(1))
+        self.logit_bias = nn.Parameter(torch.randn(1))
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)
+    def get_text_features(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.FloatTensor:
+        r"""
+        Returns:
+            text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
+            applying the projection layer to the pooled output of [`SiglipTextModel`].
+        Examples:
+        ```python
+        >>> from transformers import AutoTokenizer, AutoModel
+        >>> import torch
+        >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224")
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224")
+        >>> # important: make sure to set padding="max_length" as that's how the model was trained
+        >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt")
+        >>> with torch.no_grad():
+        ...     text_features = model.get_text_features(**inputs)
+        ```"""
+        # Use SigLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = text_outputs[1]
+
+        return pooled_output
+
+    @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
+    def get_image_features(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> torch.FloatTensor:
+        r"""
+        Returns:
+            image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by
+            applying the projection layer to the pooled output of [`SiglipVisionModel`].
+        Examples:
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, AutoModel
+        >>> import torch
+        >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224")
+        >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+        >>> inputs = processor(images=image, return_tensors="pt")
+        >>> with torch.no_grad():
+        ...     image_features = model.get_image_features(**inputs)
+        ```"""
+        # Use SiglipModel's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = vision_outputs[1]
+
+        return pooled_output
+
+    @add_start_docstrings_to_model_forward(SIGLIP_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=SiglipOutput, config_class=SiglipConfig)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        return_loss: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SiglipOutput]:
+        r"""
+        Returns:
+        Examples:
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, AutoModel
+        >>> import torch
+        >>> model = AutoModel.from_pretrained("google/siglip-base-patch16-224")
+        >>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+        >>> texts = ["a photo of 2 cats", "a photo of 2 dogs"]
+        >>> # important: we pass `padding=max_length` since the model was trained with this
+        >>> inputs = processor(text=texts, images=image, padding="max_length", return_tensors="pt")
+        >>> with torch.no_grad():
+        ...     outputs = model(**inputs)
+        >>> logits_per_image = outputs.logits_per_image
+        >>> probs = torch.sigmoid(logits_per_image) # these are the probabilities
+        >>> print(f"{probs[0][0]:.1%} that image 0 is '{texts[0]}'")
+        31.9% that image 0 is 'a photo of 2 cats'
+        ```"""
+        # Use SigLIP model's config for some fields (if specified) instead of those of vision & text components.
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        image_embeds = vision_outputs[1]
+        text_embeds = text_outputs[1]
+
+        # normalized features
+        image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+
+        # cosine similarity as logits
+        logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * self.logit_scale.exp() + self.logit_bias
+        logits_per_image = logits_per_text.t()
+
+        loss = None
+        if return_loss:
+            raise NotImplementedError("SigLIP loss to be implemented")
+
+        if not return_dict:
+            output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
+            return ((loss,) + output) if loss is not None else output
+
+        return SiglipOutput(
+            loss=loss,
+            logits_per_image=logits_per_image,
+            logits_per_text=logits_per_text,
+            text_embeds=text_embeds,
+            image_embeds=image_embeds,
+            text_model_output=text_outputs,
+            vision_model_output=vision_outputs,
+        )
+
+
+def get_siglip_vision_model(_flash_attn_2_enabled=True, **kwargs):
+    siglip_vision_config = {
+        "hidden_size": 1152,
+        "image_size": 448,
+        "intermediate_size": 4304,
+        "model_type": "siglip_vision_model",
+        "num_attention_heads": 16,
+        "num_hidden_layers": 27,
+        "patch_size": 14,
+    }
+
+    model_config = SiglipVisionConfig(**siglip_vision_config, _flash_attn_2_enabled=_flash_attn_2_enabled, **kwargs)
+
+    vision_model = SiglipVisionModel(model_config).vision_model
+
+    return vision_model