model_files

README.md

@@ -1,62 +1,3 @@
 ---
-license: creativeml-openrail-m
-language:
-- en
-metrics:
-- bleu
-tags:
-  - endpoints
-  - text-generation-inference
-inference: true
+license: apache-2.0
 ---
-
-<h3 align='center' style='font-size: 24px;'>Blazzing Fast Tiny Vision Language Model</h3>
-
-
-<p align='center', style='font-size: 16px;' >A Custom 3B parameter Model. Built by <a href="https://www.linkedin.com/in/manishkumarthota/">@Manish</a> The model is released for research purposes only, commercial use is not allowed. </p>
-
-## How to use
-
-
-**Install dependencies**
-```bash
-pip install transformers # latest version is ok, but we recommend v4.31.0
-pip install -q pillow accelerate einops
-```
-
-You can use the following code for model inference. The format of text instruction is similar to [LLaVA](https://github.com/haotian-liu/LLaVA).
-
-```Python
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from PIL import Image
-
-torch.set_default_device("cuda")
-
-#Create model
-model = AutoModelForCausalLM.from_pretrained(
-    "ManishThota/CustomModel", 
-    torch_dtype=torch.float16, 
-    device_map="auto",
-    trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("ManishThota/CustomModel", trust_remote_code=True)
-
-#function to generate the answer
-def predict(question, image_path):
-    #Set inputs
-    text = f"USER: <image>\n{question}? ASSISTANT:"
-    image = Image.open(image_path)
-    
-    input_ids = tokenizer(text, return_tensors='pt').input_ids.to('cuda')
-    image_tensor = model.image_preprocess(image)
-    
-    #Generate the answer
-    output_ids = model.generate(
-        input_ids,
-        max_new_tokens=25,
-        images=image_tensor,
-        use_cache=True)[0]
-    
-    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
-
-```

README_old.md

@@ -1,62 +0,0 @@
----
-license: creativeml-openrail-m
-language:
-- en
-metrics:
-- bleu
----
-<h1 align='center' style='font-size: 36px; font-weight: bold;'>Sparrow</h1>
-<h3 align='center' style='font-size: 24px;'>Blazzing Fast Tiny Vision Language Model</h3>
-
-
-<p align="center">
-  <img src="https://cdn-uploads.huggingface.co/production/uploads/650c7fbb8ffe1f53bdbe1aec/DTjDSq2yG-5Cqnk6giPFq.jpeg" width="50%" height="auto"/>
-</p>
-
-<p align='center', style='font-size: 16px;' >A Custom 3B parameter Model Enhanced for Educational Contexts: This specialized model integrates slide-text pairs from machine learning classes, leveraging a unique training approach. It connects a frozen pre-trained vision encoder (SigLip) with a frozen language model (Phi-2) through an innovative projector. The model employs attention mechanisms and language modeling loss to deeply understand and generate educational content, specifically tailored to the context of machine learning education. Built by <a href="https://www.linkedin.com/in/manishkumarthota/">@Manish</a> The model is released for research purposes only, commercial use is not allowed. </p>
-
-## How to use
-
-
-**Install dependencies**
-```bash
-pip install transformers # latest version is ok, but we recommend v4.31.0
-pip install -q pillow accelerate einops
-```
-
-You can use the following code for model inference. The format of text instruction is similar to [LLaVA](https://github.com/haotian-liu/LLaVA).
-
-```Python
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from PIL import Image
-
-torch.set_default_device("cuda")
-
-#Create model
-model = AutoModelForCausalLM.from_pretrained(
-    "ManishThota/Sparrow", 
-    torch_dtype=torch.float16, 
-    device_map="auto",
-    trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("ManishThota/SparrowVQE", trust_remote_code=True)
-
-#function to generate the answer
-def predict(question, image_path):
-    #Set inputs
-    text = f"USER: <image>\n{question}? ASSISTANT:"
-    image = Image.open(image_path)
-    
-    input_ids = tokenizer(text, return_tensors='pt').input_ids.to('cuda')
-    image_tensor = model.image_preprocess(image)
-    
-    #Generate the answer
-    output_ids = model.generate(
-        input_ids,
-        max_new_tokens=25,
-        images=image_tensor,
-        use_cache=True)[0]
-    
-    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
-
-```

config.json

@@ -1,13 +1,13 @@
 {
-  "_name_or_path": "ManishThota/Sparrow",
+  "_name_or_path": "MILVLG/imp-v1-3b",
   "activation_function": "gelu_new",
   "architectures": [
     "ImpForCausalLM"
   ],
   "attn_pdrop": 0.0,
   "auto_map": {
-    "AutoConfig": "configuration_imp.ImpConfig",
-    "AutoModelForCausalLM": "modeling_imp.ImpForCausalLM"
+    "AutoConfig": "MILVLG/imp-v1-3b--configuration_imp.ImpConfig",
+    "AutoModelForCausalLM": "MILVLG/imp-v1-3b--modeling_imp.ImpForCausalLM"
   },
   "embd_pdrop": 0.0,
   "eos_token_id": 50295,
@@ -29,7 +29,7 @@
   "mm_vision_select_feature": "patch",
   "mm_vision_select_layer": -2,
   "mm_vision_tower": "google/siglip-so400m-patch14-384",
-  "model_type": "Sparrow",
+  "model_type": "imp",
   "n_embd": 2560,
   "n_head": 32,
   "n_head_kv": null,

configuration_imp.py

@@ -1,175 +0,0 @@
-
-# ------------------------------- Phi-2 ---------------------------------------------
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT license.
-# https://huggingface.co/google/siglip-so400m-patch14-384
-# 
-# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
-# Licensed under the BSD 3-Clause License.
-# ------------------------------- SigLIP --------------------------------------------
-# 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.
-# ------------------------------- Llava ---------------------------------------------
-# Copyright 2023 Haotian Liu
-#
-# 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.
-# -----------------------------------------------------------------------------------
-
-
-import os
-import math
-from typing import Optional, Union
-
-from transformers import PretrainedConfig
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-class PhiConfig(PretrainedConfig):
-    """Phi configuration."""
-
-    model_type = "phi-msft"
-    attribute_map = {
-        "max_position_embeddings": "n_positions",
-        "hidden_size": "n_embd",
-        "num_attention_heads": "n_head",
-        "num_hidden_layers": "n_layer",
-    }
-
-    def __init__(
-        self,
-        vocab_size: int = 50304,
-        n_positions: int = 2048,
-        n_embd: int = 1024,
-        n_layer: int = 20,
-        n_inner: Optional[int] = None,
-        n_head: int = 16,
-        n_head_kv: Optional[int] = None,
-        rotary_dim: Optional[int] = 32,
-        activation_function: Optional[str] = "gelu_new",
-        flash_attn: bool = False,
-        flash_rotary: bool = False,
-        fused_dense: bool = False,
-        attn_pdrop: float = 0.0,
-        embd_pdrop: float = 0.0,
-        resid_pdrop: float = 0.0,
-        layer_norm_epsilon: float = 1e-5,
-        initializer_range: float = 0.02,
-        tie_word_embeddings: bool = False,
-        pad_vocab_size_multiple: int = 64,
-        **kwargs
-    ) -> None:
-        self.vocab_size = int(math.ceil(vocab_size / pad_vocab_size_multiple) * pad_vocab_size_multiple)
-        self.n_positions = n_positions
-        self.n_embd = n_embd
-        self.n_layer = n_layer
-        self.n_inner = n_inner
-        self.n_head = n_head
-        self.n_head_kv = n_head_kv
-        self.rotary_dim = min(rotary_dim, n_embd // n_head)
-        self.activation_function = activation_function
-        self.flash_attn = flash_attn
-        self.flash_rotary = flash_rotary
-        self.fused_dense = fused_dense
-        self.attn_pdrop = attn_pdrop
-        self.embd_pdrop = embd_pdrop
-        self.resid_pdrop = resid_pdrop
-        self.layer_norm_epsilon = layer_norm_epsilon
-        self.initializer_range = initializer_range
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-
-
-class SiglipVisionConfig(PretrainedConfig):
-
-    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,
-        **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
-
-    @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 ImpConfig(PhiConfig):
-    model_type = "imp"
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.image_token_index = getattr(self, "image_token_index", 50296)
-        self.image_token = getattr(self, "image_token", "<image>")
-
-        if not hasattr(self, "vision_tower_config") and hasattr(self, "mm_vision_tower"):
-            vision_tower_config = SiglipVisionConfig.from_pretrained(self.mm_vision_tower)
-            self.vision_tower_config = vision_tower_config.to_diff_dict()
-    
-    @property
-    def vision_tower_cfg(self):
-        cfg = SiglipVisionConfig.from_dict(self.vision_tower_config)
-        # imp-v1 only supports `patch` feature for now w/o cls token
-        # cfg.mm_vision_select_feature = self.mm_vision_select_feature
-        cfg.mm_vision_select_layer = self.mm_vision_select_layer
-        cfg.mm_vision_tower = self.mm_vision_tower
-        return cfg

model.safetensors

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

modeling_imp.py

@@ -1,1262 +0,0 @@
-# Copyright (c) MILVLG team.
-# Licensed under the Apache 2.0 license.
-#
-# Some code here is copied from the project Phi-2 (https://huggingface.co/microsoft/phi-2),
-# SigLIP@transformers==4.37.0.dev0 (https://huggingface.co/google/siglip-so400m-patch14-384),
-# and Llava (https://github.com/haotian-liu/LLaVA), and modified by 
-# Zhenwei Shao (shaozw@hdu.edu.cn) @ MILVLG. We thank them for their great works.
-# And their original licenses and copyright should be inherited (see the statements
-# in `configuration_imp.py` for more details).
-
-
-# Be careful: The way how `past_key_values.seqlen_offset` is updated is modified from
-# the implementation of original Phi-2. See the comments below for details.
-
-from __future__ import annotations
-import os
-import math
-import re
-from dataclasses import dataclass, field
-from typing import Any, Dict, Optional, Tuple, Union, List
-from abc import ABC, abstractmethod
-
-import torch
-import torch.nn as nn
-from einops import rearrange, repeat
-from transformers import (
-    PretrainedConfig, 
-    PreTrainedModel,
-    AutoConfig,
-    AutoModelForCausalLM
-)    
-from transformers.activations import ACT2FN
-from transformers.modeling_outputs import CausalLMOutputWithPast
-import sys
-from .configuration_imp import PhiConfig, ImpConfig
-from .vision_encoder import VisionTower
-
-try:
-    from flash_attn.bert_padding import pad_input, unpad_input
-    from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
-    from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention
-    from flash_attn.ops.fused_dense import FusedDense
-except:
-    pad_input, unpad_input = None, None
-    FlashRotaryEmbedding = None
-    FlashSelfAttention, FlashCrossAttention = None, None
-    FusedDense = None
-
-
-@dataclass
-class InferenceParams:
-    """Inference parameters passed to model to efficiently calculate
-    and store context during inference.
-
-    Reference:
-        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.
-
-    Args:
-        max_seqlen: Maximum sequence length.
-        max_batch_size: Maximum batch size.
-        seqlen_offset: Sequence length offset.
-        batch_size_offset: Batch size offset.
-        key_value_memory_dict: Key value memory dictionary.
-        lengths_per_sample: Lengths per sample.
-
-    """
-
-    max_seqlen: int = field(metadata={"help": "Maximum sequence length."})
-
-    max_batch_size: int = field(metadata={"help": "Maximum batch size."})
-
-    seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."})
-
-    batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})
-
-    key_value_memory_dict: Dict[str, Any] = field(
-        default_factory=dict, metadata={"help": "Key value memory dictionary."}
-    )
-
-    lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})
-
-
-class Embedding(nn.Module):
-    """Token embedding with dropout."""
-
-    def __init__(self, config: PretrainedConfig) -> None:
-        super().__init__()
-
-        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
-        self.drop = nn.Dropout(config.embd_pdrop)
-
-    def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
-        input_shape = input_ids.size()
-        input_ids = input_ids.view(-1, input_shape[-1])
-
-        hidden_states = self.wte(input_ids)
-        hidden_states = self.drop(hidden_states)
-
-        return hidden_states
-
-
-
-def _apply_rotary_emb(
-    x: torch.FloatTensor,
-    cos: torch.FloatTensor,
-    sin: torch.FloatTensor,
-) -> torch.FloatTensor:
-    _, seqlen, _, _ = x.shape
-    _, rotary_dim = cos.shape
-    rotary_dim *= 2
-
-    x_rot = x[:, :, :, :rotary_dim]
-    x_pass = x[:, :, :, rotary_dim:]
-
-    x1, x2 = x_rot.chunk(2, dim=-1)
-    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
-    x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]]
-
-    x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype)
-
-    return torch.cat([x_rot, x_pass], axis=-1)
-
-
-def _apply_rotary_emb_kv(
-    kv: torch.FloatTensor,
-    cos: torch.FloatTensor,
-    sin: torch.FloatTensor,
-    cos_k: Optional[torch.FloatTensor] = None,
-    sin_k: Optional[torch.FloatTensor] = None,
-) -> torch.FloatTensor:
-    _, seqlen, _, _, _ = kv.shape
-    _, rotary_dim = cos.shape
-    rotary_dim *= 2
-
-    k_rot = kv[:, :, 0, :, :rotary_dim]
-    k_pass = kv[:, :, 0, :, rotary_dim:]
-
-    k1, k2 = k_rot.chunk(2, dim=-1)
-    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
-    k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]]
-
-    k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype)
-
-    return torch.cat(
-        [
-            torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
-            kv[:, :, 1:2, :, :],
-        ],
-        axis=2,
-    )
-
-
-def _apply_rotary_emb_qkv(
-    qkv: torch.FloatTensor,
-    cos: torch.FloatTensor,
-    sin: torch.FloatTensor,
-    cos_k: Optional[torch.FloatTensor] = None,
-    sin_k: Optional[torch.FloatTensor] = None,
-) -> torch.FloatTensor:
-    _, seqlen, _, _, _ = qkv.shape
-    _, rotary_dim = cos.shape
-    rotary_dim *= 2
-
-    q_rot = qkv[:, :, 0, :, :rotary_dim]
-    q_pass = qkv[:, :, 0, :, rotary_dim:]
-
-    k_rot = qkv[:, :, 1, :, :rotary_dim]
-    k_pass = qkv[:, :, 1, :, rotary_dim:]
-
-    q1, q2 = q_rot.chunk(2, dim=-1)
-    k1, k2 = k_rot.chunk(2, dim=-1)
-    c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")
-    q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]
-
-    q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)
-    k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)
-
-    return torch.cat(
-        [
-            torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),
-            torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),
-            qkv[:, :, 2:3, :, :],
-        ],
-        axis=2,
-    )
-
-
-class RotaryEmbedding(nn.Module):
-    """Rotary positional embedding (RoPE).
-
-    Reference:
-        RoFormer: Enhanced Transformer with Rotary Position Embedding.
-        https://arxiv.org/pdf/2104.09864.pdf.
-
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        base: int = 10000,
-        scale_base: Optional[float] = None,
-        pos_idx_in_fp32: bool = True,
-        max_position_embeddings: int = 2048,
-        device: Optional[str] = None,
-        **kwargs,
-    ) -> None:
-        super().__init__()
-
-        if scale_base is not None:
-            raise NotImplementedError
-
-        self.dim = dim
-        self.base = float(base)
-        self.scale_base = scale_base
-        self.pos_idx_in_fp32 = pos_idx_in_fp32
-        self.max_position_embeddings = max_position_embeddings
-        self.device = device
-
-        # Generate and save the inverse frequency buffer (non-trainable)
-        inv_freq = self._compute_inv_freq(device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-
-        # Generate and save the scale buffer (non-trainable)
-        scale = (
-            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
-            if scale_base is not None
-            else None
-        )
-        self.register_buffer("scale", scale, persistent=False)
-
-        # Initialize cached attributes since ONNX can't rely on dynamic initialization
-        self._update_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.float32)
-
-    def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:
-        return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
-
-    def _update_cos_sin_cache(
-        self,
-        seqlen: int,
-        device: Optional[str] = None,
-        dtype: Optional[torch.dtype] = None,
-    ) -> None:
-        self._seq_len_cached = seqlen
-
-        # fp32 is preferred since the output of `torch.arange` can be quite large
-        # and bf16 would lose a lot of precision
-        if self.pos_idx_in_fp32:
-            t = torch.arange(seqlen, device=device, dtype=torch.float32)
-            if self.inv_freq.dtype != torch.float32:
-                inv_freq = self._compute_inv_freq(device=device)
-            else:
-                inv_freq = self.inv_freq
-        else:
-            t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
-            inv_freq = self.inv_freq
-
-        # `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP
-        freqs = torch.outer(t, inv_freq)
-        if self.scale is None:
-            self._cos_cached = torch.cos(freqs).to(dtype)
-            self._sin_cached = torch.sin(freqs).to(dtype)
-        else:
-            power = (
-                torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
-            ) / self.scale_base
-            scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
-
-            # Force the scale multiplication to happen in fp32
-            self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
-            self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
-            self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
-            self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
-
-    def forward(
-        self,
-        qkv: torch.Tensor,
-        kv: Optional[torch.Tensor] = None,
-        seqlen_offset: int = 0,
-        **kwargs,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if (
-            self._seq_len_cached < qkv.shape[1] + seqlen_offset
-            or self._cos_cached.device != qkv.device
-            or self._cos_cached.dtype != qkv.dtype
-            or (self.training and self._cos_cached.is_inference())
-        ):
-            self._update_cos_sin_cache(qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
-
-        if kv is None:
-            return _apply_rotary_emb_qkv(
-                qkv,
-                self._cos_cached[seqlen_offset:],
-                self._sin_cached[seqlen_offset:],
-            )
-        else:
-            q = _apply_rotary_emb(
-                qkv,
-                self._cos_cached[seqlen_offset:],
-                self._sin_cached[seqlen_offset:],
-            )
-            kv = _apply_rotary_emb_kv(
-                kv,
-                self._cos_cached[seqlen_offset:],
-                self._sin_cached[seqlen_offset:],
-            )
-
-            return q, kv
-
-
-class MLP(nn.Module):
-    """Multi-Layer Perceptron.
-
-    Reference:
-        Attention Is All You Need.
-        https://arxiv.org/pdf/1706.03762.pdf.
-
-    """
-
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        n_inner: Optional[int] = None,
-        act_fn: Optional[str] = None,
-    ) -> None:
-        super().__init__()
-
-        act_fn = config.activation_function if act_fn is None else act_fn
-
-        n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner
-        n_inner = n_inner if n_inner is not None else 4 * config.n_embd
-
-        self.fc1 = nn.Linear(config.n_embd, n_inner)
-        self.fc2 = nn.Linear(n_inner, config.n_embd)
-        self.act = ACT2FN[act_fn]
-
-    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
-        hidden_states = self.fc1(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.fc2(hidden_states)
-
-        return hidden_states
-
-
-class SelfAttention(nn.Module):
-    """Self-attention layer (compatible with PyTorch).
-
-    Reference:
-        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
-
-    """
-
-    def __init__(
-        self,
-        causal: bool = True,
-        softmax_scale: Optional[float] = None,
-        attention_dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-        self.drop = nn.Dropout(attention_dropout)
-
-    @torch.autocast("cpu", enabled=False)
-    @torch.autocast("cuda", enabled=False)
-    def forward(
-        self,
-        qkv: torch.FloatTensor,
-        causal: bool = None,
-        key_padding_mask: Optional[torch.BoolTensor] = None,
-        **kwargs,
-    ) -> torch.FloatTensor:
-        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
-        q, k, v = qkv.unbind(dim=2)
-
-        q = q.to(torch.float32)
-        k = k.to(torch.float32)
-
-        causal = self.causal if causal is None else causal
-        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
-
-        # Autocast is manually disabled to avoid `torch.einsum` performing the operation
-        # using float16, which might lead to overflow
-        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
-
-        if key_padding_mask is not None:
-            padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)
-            padding_mask.masked_fill_(key_padding_mask, 0.0)
-
-            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
-
-        if causal:
-            causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)
-            scores = scores + causal_mask.to(dtype=scores.dtype)
-
-        attention = torch.softmax(scores, dim=-1).to(v.dtype)
-        attention = self.drop(attention)
-
-        output = torch.einsum("bhts,bshd->bthd", attention, v)
-
-        return output
-
-
-class CrossAttention(nn.Module):
-    """Cross-attention layer (compatible with PyTorch).
-
-    Reference:
-        https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.
-
-    """
-
-    def __init__(
-        self,
-        causal: bool = True,
-        softmax_scale: Optional[float] = None,
-        attention_dropout: float = 0.0,
-    ) -> None:
-        super().__init__()
-
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-        self.drop = nn.Dropout(attention_dropout)
-
-    @torch.autocast("cpu", enabled=False)
-    @torch.autocast("cuda", enabled=False)
-    def forward(
-        self,
-        q: torch.FloatTensor,
-        kv: torch.FloatTensor,
-        causal: bool = None,
-        key_padding_mask: Optional[torch.BoolTensor] = None,
-        **kwargs,
-    ) -> torch.FloatTensor:
-        batch_size, seqlen_q = q.shape[0], q.shape[1]
-        seqlen_k = kv.shape[1]
-
-        if kv.shape[3] != q.shape[2]:
-            kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])
-        k, v = kv.unbind(dim=2)
-
-        q = q.to(torch.float32)
-        k = k.to(torch.float32)
-
-        causal = self.causal if causal is None else causal
-        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])
-
-        # Autocast is manually disabled to avoid `torch.einsum` performing the operation
-        # using float16, which might lead to overflow
-        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
-
-        if key_padding_mask is not None:
-            padding_mask = torch.full(
-                (batch_size, seqlen_k),
-                -10000.0,
-                dtype=scores.dtype,
-                device=scores.device,
-            )
-            padding_mask.masked_fill_(key_padding_mask, 0.0)
-
-            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")
-
-        if causal:
-            rows = rearrange(torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1")
-            cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)
-            causal_mask = cols > rows + seqlen_k - seqlen_q
-
-            scores = scores.masked_fill(causal_mask, -10000.0)
-
-        attention = torch.softmax(scores, dim=-1).to(v.dtype)
-        attention = self.drop(attention)
-
-        output = torch.einsum("bhts,bshd->bthd", attention, v)
-
-        return output
-
-
-def _find_mha_dims(
-    config: PretrainedConfig,
-    n_head: Optional[int] = None,
-    n_head_kv: Optional[int] = None,
-    head_dim: Optional[int] = None,
-) -> Tuple[int, int]:
-    if n_head is None and head_dim is None:
-        head_dim = config.n_embd // config.n_head
-        n_head = config.n_head
-    elif n_head is None or head_dim is None:
-        raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")
-
-    if n_head_kv is None:
-        n_head_kv = getattr(config, "n_head_kv", None) or n_head
-
-    return n_head, n_head_kv, head_dim
-
-
-def _update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:
-    num_heads, head_dim = kv.shape[-2:]
-
-    if layer_idx not in inference_params.key_value_memory_dict:
-        inference_params.key_value_memory_dict[layer_idx] = torch.empty(
-            inference_params.max_batch_size,
-            inference_params.max_seqlen,
-            2,
-            num_heads,
-            head_dim,
-            dtype=kv.dtype,
-            device=kv.device,
-        )
-
-    batch_start = inference_params.batch_size_offset
-    batch_end = batch_start + kv.shape[0]
-
-    sequence_start = inference_params.seqlen_offset
-    sequence_end = sequence_start + kv.shape[1]
-
-    # When the current sequence length is equal to or larger than the maximum sequence length,
-    # we need to concatenate the current `kv` with the cached `kv` to expand its length
-    if sequence_end >= inference_params.max_seqlen:
-        inference_params.key_value_memory_dict[layer_idx] = torch.concatenate((inference_params.key_value_memory_dict[layer_idx], kv), dim=1)
-
-    inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...] = kv
-    kv = inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, :sequence_end, ...]
-        
-    return kv
-
-
-class MHA(nn.Module):
-    """Multi-head attention layer."""
-
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[str] = None,
-        rotary_dim: Optional[int] = None,
-        rotary_base: float = 10000.0,
-        rotary_scale_base: Optional[float] = None,
-        n_head: Optional[int] = None,
-        n_head_kv: Optional[int] = None,
-        head_dim: Optional[int] = None,
-        bias: bool = True,
-        causal: bool = True,
-        softmax_scale: Optional[float] = None,
-        layer_idx: Optional[int] = None,
-        return_residual: bool = False,
-        checkpointing: bool = False,
-    ) -> None:
-        super().__init__()
-
-        # Rotary embedding
-        self.rotary_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)
-        if self.rotary_dim > 0:
-            rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding
-            if rotary_cls is None:
-                rotary_cls = RotaryEmbedding
-
-            rotary_kwargs = {}
-            if rotary_cls is RotaryEmbedding:
-                rotary_kwargs["max_position_embeddings"] = config.n_positions
-
-            self.rotary_emb = rotary_cls(
-                self.rotary_dim,
-                base=rotary_base,
-                scale_base=rotary_scale_base,
-                device=device,
-                **rotary_kwargs,
-            )
-
-        # MLP
-        self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(
-            config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim
-        )
-        op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)
-        hidden_size = config.n_embd
-
-        linear_cls = FusedDense if config.fused_dense else nn.Linear
-        if linear_cls is None:
-            linear_cls = nn.Linear
-
-        self.Wqkv = linear_cls(hidden_size, op_size, bias=bias, device=device, dtype=dtype)
-        self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)
-
-        # Attention
-        attn_cls = FlashSelfAttention if config.flash_attn else SelfAttention
-        if attn_cls is None:
-            attn_cls = SelfAttention
-
-        cross_attn_cls = FlashCrossAttention if config.flash_attn else CrossAttention
-        if cross_attn_cls is None:
-            cross_attn_cls = CrossAttention
-
-        self.inner_attn = attn_cls(
-            causal=causal,
-            softmax_scale=softmax_scale,
-            attention_dropout=config.attn_pdrop,
-        )
-        self.inner_cross_attn = cross_attn_cls(
-            causal=causal,
-            softmax_scale=softmax_scale,
-            attention_dropout=config.attn_pdrop,
-        )
-
-        self.flash_attn = config.flash_attn and attn_cls is FlashSelfAttention
-        self.layer_idx = layer_idx
-        self.return_residual = return_residual
-        self.checkpointing = checkpointing
-
-    def _forward_self_attn(
-        self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]
-    ) -> torch.FloatTensor:
-        qkv = self.Wqkv(x)
-        qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)
-
-        if self.rotary_dim > 0:
-            qkv = self.rotary_emb(qkv)
-
-        if self.flash_attn:
-            batch_size, seqlen = qkv.shape[0], qkv.shape[1]
-
-            cu_seqlens, max_seqlen = None, None
-            if key_padding_mask is not None:
-                # If `key_padding_mask` is supplied, we need to unpad the input and retrieve
-                # the `cu_seqlens` and `max_seqlen` to be used by `flash-attn`
-                qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, key_padding_mask)
-
-            if self.checkpointing:
-                attn_output = torch.utils.checkpoint.checkpoint(
-                    self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen
-                )
-            else:
-                attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)
-
-            # If `key_padding_mask` is supplied, we need to pad the output back to the original shape
-            return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output
-
-        if self.checkpointing:
-            return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask)
-
-        return self.inner_attn(qkv, key_padding_mask=key_padding_mask)
-
-    def _forward_cross_attn(
-        self,
-        x: torch.FloatTensor,
-        past_key_values: Optional[InferenceParams],
-        key_padding_mask: Optional[torch.BoolTensor],
-    ) -> torch.FloatTensor:
-        batch_size = x.shape[0]
-
-        qkv = self.Wqkv(x)
-
-        q = qkv[..., : self.n_head * self.head_dim]
-        q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
-
-        kv = qkv[..., self.n_head * self.head_dim :]
-        kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
-
-        seqlen_offset = past_key_values.seqlen_offset if past_key_values is not None else 0
-        causal = None if seqlen_offset == 0 else False
-        if self.rotary_dim > 0:
-            q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)
-
-        if past_key_values is not None:
-            kv = _update_kv_cache(kv, past_key_values, self.layer_idx)
-
-        if self.flash_attn:
-            batch_size, seqlen_q = q.shape[0], q.shape[1]
-            seqlen_k = kv.shape[1]
-
-            cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = (
-                None,
-                None,
-                None,
-                None,
-            )
-            if key_padding_mask is not None:
-                kv, _, cu_seqlens_k, max_seqlen_k = unpad_input(kv, key_padding_mask)
-
-                if seqlen_q == 1:
-                    key_padding_mask = torch.ones(batch_size, 1, device=q.device)
-                elif seqlen_q != seqlen_k:
-                    key_padding_mask = key_padding_mask[:, -seqlen_q:]
-
-                q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, key_padding_mask)
-
-            if self.checkpointing:
-                attn_output = torch.utils.checkpoint.checkpoint(
-                    self.inner_cross_attn,
-                    q,
-                    kv,
-                    causal=causal,
-                    cu_seqlens=cu_seqlens_q,
-                    max_seqlen=max_seqlen_q,
-                    cu_seqlens_k=cu_seqlens_k,
-                    max_seqlen_k=max_seqlen_k,
-                )
-            else:
-                attn_output = self.inner_cross_attn(
-                    q,
-                    kv,
-                    causal=causal,
-                    cu_seqlens=cu_seqlens_q,
-                    max_seqlen=max_seqlen_q,
-                    cu_seqlens_k=cu_seqlens_k,
-                    max_seqlen_k=max_seqlen_k,
-                )
-
-            return (
-                pad_input(attn_output, indices_q, batch_size, max_seqlen_q)
-                if key_padding_mask is not None
-                else attn_output
-            )
-
-        if self.checkpointing:
-            return torch.utils.checkpoint.checkpoint(
-                self.inner_cross_attn,
-                q,
-                kv,
-                key_padding_mask=key_padding_mask,
-                causal=causal,
-            )
-
-        return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)
-
-    def forward(
-        self,
-        x: torch.FloatTensor,
-        past_key_values: Optional[InferenceParams] = None,
-        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
-        if attention_mask is not None:
-            attention_mask = attention_mask.bool()
-        else:
-            attention_mask = None
-
-        # MHA
-        if self.n_head == self.n_head_kv:
-            if past_key_values is None:
-                # If `past_key_values` are not supplied, we run self-attention
-                attn_output = self._forward_self_attn(x, attention_mask)
-            else:
-                # If `past_key_values` are supplied, it means that we might have cached values and
-                # could take advantage of cross-attention
-                attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
-        # MQA / GQA
-        else:
-            # Regardless of `past_key_values` being supplied or not, it always use cross-attention
-            # because `q` and `kv` lengths might be different
-            attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)
-
-        output = rearrange(attn_output, "... h d -> ... (h d)")
-        output = self.out_proj(output)
-
-        return output if not self.return_residual else (output, x)
-
-
-class ParallelBlock(nn.Module):
-    """Parallel block.
-
-    This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).
-
-    """
-
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        block_idx: Optional[int] = None,
-    ) -> None:
-        super().__init__()
-
-        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-        self.resid_dropout = nn.Dropout(config.resid_pdrop)
-        self.block_idx = block_idx
-
-        self.mixer = MHA(config, layer_idx=block_idx)
-        self.mlp = MLP(config)
-
-    def forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
-        attention_mask: Optional[torch.BoolTensor] = None,
-        **kwargs,
-    ) -> torch.FloatTensor:
-        residual = hidden_states
-        hidden_states = self.ln(hidden_states)
-
-        attn_outputs = self.mixer(
-            hidden_states,
-            past_key_values=past_key_values,
-            attention_mask=attention_mask,
-        )
-        if isinstance(attn_outputs, tuple):
-            attn_outputs = attn_outputs[0]
-
-        attn_outputs = self.resid_dropout(attn_outputs)
-        feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
-
-        hidden_states = attn_outputs + feed_forward_hidden_states + residual
-
-        return hidden_states
-
-
-class CausalLMHead(nn.Module):
-    """Causal Language Modeling head.
-
-    Reference:
-        Improving Language Understanding by Generative Pre-Training.
-        https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
-
-    """
-
-    def __init__(self, config: PretrainedConfig) -> None:
-        super().__init__()
-
-        self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
-        self.linear = nn.Linear(config.n_embd, config.vocab_size)
-
-    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
-        hidden_states = self.ln(hidden_states)
-        logits = self.linear(hidden_states).to(torch.float32)
-
-        return logits
-
-
-class PhiPreTrainedModel(PreTrainedModel):
-    """Phi pre-trained model."""
-
-    config_class = PhiConfig
-    base_model_prefix = "transformer"
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["ParallelBlock", "CLIPEncoderLayer", "Block"]
-
-    def __init__(self, *inputs, **kwargs) -> None:
-        super().__init__(*inputs, **kwargs)
-
-    def _init_weights(self, module: nn.Module) -> None:
-        if isinstance(module, (nn.Linear,)):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, nn.LayerNorm):
-            if module.bias is not None:
-                module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor,
-        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
-        attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,
-        **kwargs,
-    ) -> Dict[str, Any]:
-        if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):
-            past_key_values = InferenceParams(
-                max_seqlen=self.config.n_positions,
-                max_batch_size=input_ids.shape[0],
-                seqlen_offset=0,
-                batch_size_offset=0,
-                key_value_memory_dict={},
-                lengths_per_sample=None,
-            )
-        else:
-            # ======================================================================
-            # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`
-            # inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...]
-            # past_key_values.seqlen_offset = input_ids.shape[1] - 1
-            # ======================================================================
-            # I change the way of updating `past_key_values.seqlen_offset` to make the inference of imp work.
-            # [Edited by zhenwei - 2024-01-20 21:15]
-            input_ids = input_ids[:, -1].unsqueeze(-1)
-
-        return {
-            "input_ids": input_ids,
-            "past_key_values": past_key_values,
-            "attention_mask": attention_mask,
-        }
-
-
-class LlavaMetaModel(ABC):
-    """
-    Define the APIs for building components that are related to image perceiving.
-    This implementation is based on the implementation from the Llave project.
-    """
-
-    def get_vision_tower(self):
-        vision_tower = getattr(self, 'vision_tower', None)
-        if type(vision_tower) is list:
-            vision_tower = vision_tower[0]
-        return vision_tower
-    
-    def build_vision_tower(self, config):
-        self.vision_tower = VisionTower(config.vision_tower_cfg)
-
-    def build_vision_projector(self, config):
-        projector_type = getattr(config, 'mm_projector_type', 'linear')
-
-        if projector_type == 'linear':
-            self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)
-            return
-
-        mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
-        if mlp_gelu_match:
-            mlp_depth = int(mlp_gelu_match.group(1))
-            modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
-            for _ in range(1, mlp_depth):
-                modules.append(nn.GELU())
-                modules.append(nn.Linear(config.hidden_size, config.hidden_size))
-            self.mm_projector = nn.Sequential(*modules)
-            return
-
-        if projector_type == 'identity':
-            self.mm_projector = nn.Identity()
-            return
-
-        raise ValueError(f'Unknown projector type: {projector_type}')
-
-
-class ImpModel(PhiPreTrainedModel, LlavaMetaModel):
-    """Imp model. This implementation is modified from the implementation of Phi-2"""
-
-    config_class = ImpConfig
-    # _keys_to_ignore_on_load_missing = [""]
-    # _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
-
-    def __init__(self, config: ImpConfig) -> None:
-        super().__init__(config)
-
-        self.embd = Embedding(config)
-        self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])
-        self.gradient_checkpointing = False
-
-        if hasattr(config, "mm_vision_tower"):
-            self.build_vision_tower(config)
-            self.build_vision_projector(config)
-
-        self.post_init()
-
-    def embed_tokens(self, input_ids: torch.LongTensor) -> torch.FloatTensor:
-        return self.embd(input_ids)[0]
-    
-    def get_input_embeddings(self) -> nn.Embedding:
-        return self.embd.wte
-    
-    def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
-        self.embd.wte = new_embeddings
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
-        attention_mask: Optional[torch.BoolTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None
-    ) -> torch.FloatTensor:
-
-        if inputs_embeds is None:
-            hidden_states = self.embd(input_ids)
-        else:
-            hidden_states = inputs_embeds
-
-        for layer in self.h:
-            if self.gradient_checkpointing and self.training:
-
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs)
-
-                    return custom_forward
-
-                hidden_states = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(layer),
-                    hidden_states,
-                    None,
-                    attention_mask,
-                )
-            else:
-                hidden_states = layer(
-                    hidden_states,
-                    past_key_values=past_key_values,
-                    attention_mask=attention_mask,
-                )
-
-        # I change the way of updating `past_key_values.seqlen_offset` to make the inference of imp work.
-        # [Edited by zhenwei - 2024-01-20 21:15]
-        if past_key_values is not None: # FIXME: when multi-batch inference, it is a bug
-            past_key_values.seqlen_offset += hidden_states.shape[1]
-
-        return hidden_states
-
-
-class LlavaMetaForCausalLM(ABC):
-    """This implementation is based on the implementation from the Llave project."""
-
-    def init_constants(self, config):
-        self.IGNORE_INDEX = getattr(config, 'ignore_index', -100)
-        self.IMAGE_TOKEN_INDEX = getattr(config, 'image_token_index', 50296)
-        self.DEFAULT_IMAGE_TOKEN = getattr(config, 'image_token', "<image>")
-
-    @abstractmethod
-    def get_model(self):
-        pass
-
-    def get_vision_tower(self):
-        return self.get_model().get_vision_tower()
-
-    def encode_images(self, images):
-        image_features = self.get_model().get_vision_tower()(images)
-        image_features = self.get_model().mm_projector(image_features)
-        return image_features
-
-    def prepare_inputs_labels_for_multimodal(
-        self, input_ids, position_ids, attention_mask, past_key_values, labels, images
-    ):
-        vision_tower = self.get_vision_tower()
-        # if vision_tower is None or images is None or past_key_values.seqlen_offset != 0:
-        if vision_tower is None or images is None or input_ids.shape[1] == 1:
-            if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1:
-                target_shape = past_key_values.seqlen_offset + 1
-                # inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...]
-                attention_mask = torch.cat((attention_mask, torch.ones(
-                    (attention_mask.shape[0], target_shape - attention_mask.shape[1]),
-                    dtype=attention_mask.dtype,
-                    device=attention_mask.device
-                )), dim=1)
-                position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
-            return input_ids, position_ids, attention_mask, past_key_values, None, labels
-
-        if type(images) is list or images.ndim == 5:
-            concat_images = torch.cat([image for image in images], dim=0)
-            concat_images = concat_images.to(device=self.device, dtype=vision_tower.dtype)
-            image_features = self.encode_images(concat_images)
-            split_sizes = [image.shape[0] for image in images]
-            image_features = torch.split(image_features, split_sizes, dim=0)
-            image_features = [x.flatten(0, 1).to(self.device) for x in image_features]
-        else:
-            images = images.to(device=self.device, dtype=vision_tower.dtype)
-            image_features = self.encode_images(images).to(self.device)
-
-        # TODO: image start / end is not implemented here to support pretraining.
-        if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
-            raise NotImplementedError
-
-        # Let's just add dummy tensors if they do not exist,
-        # it is a headache to deal with None all the time.
-        # But it is not ideal, and if you have a better idea,
-        # please open an issue / submit a PR, thanks.
-        _labels = labels
-        _position_ids = position_ids
-        _attention_mask = attention_mask
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
-        else:
-            attention_mask = attention_mask.bool()
-        if position_ids is None:
-            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
-        if labels is None:
-            labels = torch.full_like(input_ids, self.IGNORE_INDEX)
-
-        # remove the padding using attention_mask -- TODO: double check
-        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]
-        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]
-
-        new_input_embeds = []
-        new_labels = []
-        cur_image_idx = 0
-        for batch_idx, cur_input_ids in enumerate(input_ids):
-            num_images = (cur_input_ids == self.IMAGE_TOKEN_INDEX).sum()
-            if num_images == 0:
-                cur_image_features = image_features[cur_image_idx]
-                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
-                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)
-                new_input_embeds.append(cur_input_embeds)
-                new_labels.append(labels[batch_idx])
-                cur_image_idx += 1
-                continue
-
-            image_token_indices = [-1] + torch.where(cur_input_ids == self.IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]
-            cur_input_ids_noim = []
-            cur_labels = labels[batch_idx]
-            cur_labels_noim = []
-            for i in range(len(image_token_indices) - 1):
-                cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])
-                cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])
-            split_sizes = [x.shape[0] for x in cur_labels_noim]
-            cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
-            # print(cur_input_embeds.shape)
-            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)
-            cur_new_input_embeds = []
-            cur_new_labels = []
-
-            for i in range(num_images + 1):
-                cur_new_input_embeds.append(cur_input_embeds_no_im[i])
-                cur_new_labels.append(cur_labels_noim[i])
-                if i < num_images:
-                    cur_image_features = image_features[cur_image_idx]
-                    cur_image_idx += 1
-                    cur_new_input_embeds.append(cur_image_features)
-                    cur_new_labels.append(torch.full((cur_image_features.shape[0],), self.IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))
-
-            cur_new_input_embeds = torch.cat(cur_new_input_embeds)
-            cur_new_labels = torch.cat(cur_new_labels)
-
-            new_input_embeds.append(cur_new_input_embeds)
-            new_labels.append(cur_new_labels)
-
-        # Truncate sequences to max length as image embeddings can make the sequence longer
-        tokenizer_model_max_length = getattr(self.config, 'tokenizer_model_max_length', None)
-        if tokenizer_model_max_length is not None:
-            new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]
-            new_labels = [x[:tokenizer_model_max_length] for x in new_labels]
-
-        # Combine them
-        max_len = max(x.shape[0] for x in new_input_embeds)
-        batch_size = len(new_input_embeds)
-
-        new_input_embeds_padded = []
-        new_labels_padded = torch.full((batch_size, max_len), self.IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)
-        attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)
-        position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)
-
-        for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):
-            cur_len = cur_new_embed.shape[0]
-            if getattr(self.config, 'tokenizer_padding_side', 'right') == "left":
-                new_input_embeds_padded.append(torch.cat((
-                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device),
-                    cur_new_embed
-                ), dim=0))
-                if cur_len > 0:
-                    new_labels_padded[i, -cur_len:] = cur_new_labels
-                    attention_mask[i, -cur_len:] = True
-                    position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
-            else:
-                new_input_embeds_padded.append(torch.cat((
-                    cur_new_embed,
-                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)
-                ), dim=0))
-                if cur_len > 0:
-                    new_labels_padded[i, :cur_len] = cur_new_labels
-                    attention_mask[i, :cur_len] = True
-                    position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
-
-        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)
-
-        if _labels is None:
-            new_labels = None
-        else:
-            new_labels = new_labels_padded
-
-        if _attention_mask is None:
-            attention_mask = None
-        else:
-            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)
-
-        if _position_ids is None:
-            position_ids = None
-
-        return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels
-
-
-class ImpForCausalLM(PhiPreTrainedModel, LlavaMetaForCausalLM):
-    """Imp for Causal Language Modeling."""
-
-    # _keys_to_ignore_on_load_missing = [""]
-    # _keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
-    config_class = ImpConfig
-
-    def __init__(self, config: ImpConfig) -> None:
-        super().__init__(config)
-
-        self.transformer = ImpModel(config)
-        self.lm_head = CausalLMHead(config)
-        
-        self.post_init()
-        self.init_constants(config)
-
-    def get_output_embeddings(self) -> nn.Linear:
-        return self.lm_head.linear
-
-    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
-        self.lm_head.linear = new_embeddings
-
-    def get_model(self):
-        return self.transformer
-    
-    def image_preprocess(self, images):
-        return self.get_vision_tower().image_processor(images)['pixel_values']
-    
-    def backbone_forward(
-        self,
-        input_ids: torch.LongTensor,
-        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
-        attention_mask: Optional[torch.BoolTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        **kwargs,
-    ) -> CausalLMOutputWithPast:
-        hidden_states = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds)
-        lm_logits = self.lm_head(hidden_states)
-
-        return CausalLMOutputWithPast(loss=None, logits=lm_logits, past_key_values=past_key_values)
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        images: Optional[torch.FloatTensor] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, CausalLMOutputWithPast]:
-
-        if inputs_embeds is None:
-            (
-                input_ids,
-                position_ids,
-                attention_mask,
-                past_key_values,
-                inputs_embeds,
-                labels
-            ) = self.prepare_inputs_labels_for_multimodal(
-                input_ids,
-                position_ids,
-                attention_mask,
-                past_key_values,
-                labels,
-                images
-            )
-
-        return self.backbone_forward(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            labels=labels,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict
-        )
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
-        images = kwargs.pop("images", None)
-        _inputs = super().prepare_inputs_for_generation(
-            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
-        )
-        if images is not None:
-            _inputs['images'] = images
-        return _inputs
-
-
-AutoConfig.register("imp", ImpConfig)
-AutoModelForCausalLM.register(ImpConfig, ImpForCausalLM)

pytorch_model.bin

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

vision_encoder.py

@@ -1,593 +0,0 @@
-# Copyright (c) MILVLG team.
-# Licensed under the Apache 2.0 license.
-#
-# Some code here is copied from the project Phi-2 (https://huggingface.co/microsoft/phi-2),
-# SigLIP@transformers==4.37.0.dev0 (https://huggingface.co/google/siglip-so400m-patch14-384),
-# and Llava (https://github.com/haotian-liu/LLaVA), and modified by 
-# Zhenwei Shao (shaozw@hdu.edu.cn) @ MILVLG. We thank them for their great works.
-# And their original licenses and copyright should be inherited (see the statements
-# in `configuration_imp.py` for more details).
-
-
-from typing import Any, Optional, Tuple, Union, List, Dict
-from dataclasses import dataclass
-import math
-import warnings
-from functools import partial, reduce
-
-
-import numpy as np
-from PIL import Image
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from transformers.image_processing_utils import BatchFeature
-from transformers.image_transforms import (
-    convert_to_rgb,
-    normalize,
-    rescale,
-    resize,
-    to_channel_dimension_format,
-)
-from transformers.image_utils import (
-    ChannelDimension,
-    PILImageResampling,
-    to_numpy_array,
-)
-from transformers.activations import ACT2FN
-from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import ModelOutput
-
-from .configuration_imp import SiglipVisionConfig
-
-
-# ============================================================================
-# A simple image preprocessor for SigLIP models. 
-# ============================================================================
-
-def simple_image_processor(
-        images, 
-        image_mean=(0.5, 0.5, 0.5), 
-        image_std=(0.5, 0.5, 0.5), 
-        size=(384, 384), 
-        resample=PILImageResampling.BICUBIC, 
-        rescale_factor=1 / 255, 
-        data_format=ChannelDimension.FIRST,
-        return_tensors="pt"
-    ):
-
-    if isinstance(images, Image.Image):
-        images = [images]
-    else:
-        assert isinstance(images, list)
-    
-    transforms = [
-        convert_to_rgb,
-        to_numpy_array,
-        partial(resize, size=size, resample=resample, data_format=data_format),
-        partial(rescale, scale=rescale_factor, data_format=data_format),
-        partial(normalize, mean=image_mean, std=image_std, data_format=data_format),
-        partial(to_channel_dimension_format, channel_dim=data_format, input_channel_dim=data_format),
-    ]
-
-    images = reduce(lambda x, f: [*map(f, x)], transforms, images)
-    data = {"pixel_values": images}
-    
-    return BatchFeature(data=data, tensor_type=return_tensors)
-
-# ============================================================================
-# Definitions for SigLIP models. 
-# ============================================================================
-
-@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
-
-
-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 = (self.image_size // self.patch_size) ** 2
-        self.num_positions = self.num_patches
-        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
-        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
-
-    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
-        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
-        embeddings = patch_embeds.flatten(2).transpose(1, 2)
-
-        embeddings = embeddings + self.position_embedding(self.position_ids)
-        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
-
-
-# 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: SiglipVisionConfig):
-        super().__init__()
-        self.embed_dim = config.hidden_size
-        self.self_attn = SiglipAttention(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)
-
-    # Ignore copy
-    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 = SiglipVisionConfig
-    base_model_prefix = "siglip"
-    supports_gradient_checkpointing = True
-
-    def _init_weights(self, module):
-        """Initialize the weights"""
-        pass
-
-# 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: SiglipVisionConfig
-    """
-
-    def __init__(self, config: SiglipVisionConfig):
-        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 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)
-
-    def forward(
-        self,
-        pixel_values,
-        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
-
-        hidden_states = self.embeddings(pixel_values)
-
-        encoder_outputs = self.encoder(
-            inputs_embeds=hidden_states,
-            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(last_hidden_state)
-
-        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):
-        batch_size = hidden_state.shape[0]
-        probe = self.probe.repeat(batch_size, 1, 1)
-
-        hidden_state = self.attention(probe, hidden_state, hidden_state)[0]
-
-        residual = hidden_state
-        hidden_state = self.layernorm(hidden_state)
-        hidden_state = residual + self.mlp(hidden_state)
-
-        return hidden_state[:, 0]
-
-
-class SiglipVisionModel(SiglipPreTrainedModel):
-    config_class = SiglipVisionConfig
-    main_input_name = "pixel_values"
-    _no_split_modules = ["SiglipEncoderLayer"]
-
-    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
-
-    def forward(
-        self,
-        pixel_values,
-        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,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-
-# ============================================================================
-# VisionTower module for Imp
-# ============================================================================
-
-class VisionTower(nn.Module):
-    def __init__(self, vision_tower_cfg, delay_load=False):
-        super().__init__()
-
-        self.is_loaded = False
-
-        self.config = vision_tower_cfg
-        self.vision_tower_name = vision_tower_cfg.mm_vision_tower
-        self.select_layer = vision_tower_cfg.mm_vision_select_layer
-        # self.select_feature = getattr(vision_tower_cfg, 'mm_vision_select_feature', 'patch')
-
-        self.image_processor = simple_image_processor
-
-        if not delay_load:
-            self.load_model()
-        else:
-            raise NotImplementedError("delay load is not implemented yet.")
-
-    def load_model(self):
-        if self.is_loaded:
-            return
-
-        # "google/siglip-so400m-patch14-384"
-        # self.vision_tower = SiglipVisionModel.from_pretrained(self.vision_tower_name)
-        self.vision_tower = SiglipVisionModel(self.config)
-        del self.vision_tower.vision_model.encoder.layers[(self.select_layer + 1):]
-        self.vision_tower.vision_model.head = nn.Identity()
-        self.vision_tower.requires_grad_(False)
-        self.vision_tower.eval()
-
-        self.is_loaded = True
-
-    @torch.no_grad()
-    def forward(self, images):
-        if type(images) is list:
-            image_features = []
-            for image in images:
-                image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)
-                image_feature = image_forward_out.hidden_states[-1].to(image.dtype)
-                assert image_features.shape[-2] == 729
-                image_features.append(image_feature)
-        else:
-            image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)
-            image_features = image_forward_outs.hidden_states[-1].to(images.dtype)
-            assert image_features.shape[-2] == 729
-
-        return image_features
-
-    @property
-    def dummy_feature(self):
-        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
-
-    @property
-    def dtype(self):
-        for p in self.vision_tower.parameters():
-            return p.dtype
-
-    @property
-    def device(self):
-        for p in self.vision_tower.parameters():
-            return p.device
-
-    @property
-    def hidden_size(self):
-        return self.config.hidden_size
-
-    @property
-    def num_patches(self):
-        return (self.config.image_size // self.config.patch_size) ** 2