Upload folder using huggingface_hub

README.md

@@ -0,0 +1,149 @@
+---
+datasets:
+- liuhaotian/LLaVA-Pretrain
+- liuhaotian/LLaVA-Instruct-150K
+language:
+- en
+tags:
+- llava
+- phi
+license: mit
+library_name: transformers
+widget:
+- text: "What animal is it?"
+  src: "https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg"
+- text: "Where is it?"
+  src: "https://huggingface.co/datasets/mishig/sample_images/resolve/main/palace.jpg"
+---
+
+# Multi-crop LLaVA-3b
+
+<a target="_blank" href="https://colab.research.google.com/drive/1W7JQrFXwFunAY1XvS31mwC7mrXBgGD_M">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
+
+## Model details
+
+The core idea behind multi-crop LLaVA (MC-LLaVA) is that instead of N visual token embeddings per image, I generate one token embedding per N parts of the image. 
+Having high-quality embeddings for smaller parts of the image helps to extract more details and understand the scene better.
+
+For every crop of the image, I generate an embedding from the full SigLIP encoder (size [1, 1152]) and then push all N embeddings through the LLaVA adapter, which 
+gives the token embedding of size [N, 2560]. Right now, the tokens do not contain explicit information about their position in the original image. I plan to add it later.
+
+MC-LLaVA-3b was fine-tuned from [Dolphin 2.6 Phi](https://huggingface.co/cognitivecomputations/dolphin-2_6-phi-2) using vision tower from
+[SigLIP 400M](https://huggingface.co/timm/ViT-SO400M-14-SigLIP-384). 
+
+The context length during training was 1200 tokens, as the L4 GPUs I used didn't allow me to get more.
+
+As Dolphin 2.6 Phi, LLaVA-3b uses ChatML prompt format:
+
+```
+<|im_start|>system
+You are Dolphin, a helpful AI assistant.<|im_end|>
+<|im_start|>user
+{prompt}<|im_end|>
+<|im_start|>assistant
+```
+
+## How to use
+
+**Install dependencies**
+
+```bash
+!pip install -q open_clip_torch timm einops
+```
+
+**Download modeling files**
+
+```python
+from huggingface_hub import hf_hub_download
+
+hf_hub_download(repo_id="visheratin/LLaVA-3b", filename="configuration_llava.py", local_dir="./", force_download=True)
+hf_hub_download(repo_id="visheratin/LLaVA-3b", filename="configuration_phi.py", local_dir="./", force_download=True)
+hf_hub_download(repo_id="visheratin/LLaVA-3b", filename="modeling_llava.py", local_dir="./", force_download=True)
+hf_hub_download(repo_id="visheratin/LLaVA-3b", filename="modeling_phi.py", local_dir="./", force_download=True)
+hf_hub_download(repo_id="visheratin/LLaVA-3b", filename="processing_llava.py", local_dir="./", force_download=True)
+```
+
+**Create a model**
+
+```python
+from modeling_llava import LlavaForConditionalGeneration
+import torch
+
+model = LlavaForConditionalGeneration.from_pretrained("visheratin/LLaVA-3b", torch_dtype=torch.float16)
+model = model.to("cuda")
+```
+
+**Create processors**
+
+```python
+from transformers import AutoTokenizer
+from processing_llava import LlavaProcessor, OpenCLIPImageProcessor
+
+tokenizer = AutoTokenizer.from_pretrained("visheratin/LLaVA-3b")
+image_processor = OpenCLIPImageProcessor(model.config.preprocess_config)
+processor = LlavaProcessor(image_processor, tokenizer)
+```
+
+**Set image and text**
+
+```python
+from PIL import Image
+import requests
+
+image_file = "https://images.unsplash.com/photo-1439246854758-f686a415d9da"
+raw_image = Image.open(requests.get(image_file, stream=True).raw)
+
+prompt = """<|im_start|>system
+A chat between a curious human and an artificial intelligence assistant.
+The assistant gives helpful, detailed, and polite answers to the human's questions.
+The assistant does not hallucinate and pays very close attention to the details.<|im_end|>
+<|im_start|>user
+<image>
+Describe the image.<|im_end|>
+<|im_start|>assistant
+"""
+```
+
+**Process inputs**
+
+```python
+inputs = processor(prompt, raw_image, model, return_tensors='pt')
+
+inputs['input_ids'] = inputs['input_ids'].to(model.device)
+inputs['attention_mask'] = inputs['attention_mask'].to(model.device)
+```
+
+**Generate the data**
+
+```python
+import torch
+
+with torch.inference_mode():
+  output = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.4, pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id)
+```
+
+## Benchmarks
+
+- TextVQA - 38.59%
+- GQA - 49.6%
+- VQAv2 - 64.24%
+- VizWiz - 24.88%
+- POPE - 80.59%
+- V*-bench - 52.25% (OCR - 46.66%, GPT4V-hard - 41.17%, direct attributes - 43.48%, relative position - 65.79%)
+
+## Examples
+
+<a target="_blank" href="https://colab.research.google.com/drive/1sXDvVl5s9fTcE0N2bQGOlXhnNlKEdeun">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
+
+## License
+
+The model is licensed under MIT license, but since the data used for model training is largely synthetic, you should also follow OpenAI and Google Gemini terms of service. 
+Which means don't create competitor models for them.
+
+## Acknowledgments
+
+Thanks to [ML Collective](https://mlcollective.org/) for providing credits for computing resources.

added_tokens.json

@@ -0,0 +1,44 @@
+{
+  "\t\t": 50294,
+  "\t\t\t": 50293,
+  "\t\t\t\t": 50292,
+  "\t\t\t\t\t": 50291,
+  "\t\t\t\t\t\t": 50290,
+  "\t\t\t\t\t\t\t": 50289,
+  "\t\t\t\t\t\t\t\t": 50288,
+  "\t\t\t\t\t\t\t\t\t": 50287,
+  "  ": 50286,
+  "   ": 50285,
+  "    ": 50284,
+  "     ": 50283,
+  "      ": 50282,
+  "       ": 50281,
+  "        ": 50280,
+  "         ": 50279,
+  "          ": 50278,
+  "           ": 50277,
+  "            ": 50276,
+  "             ": 50275,
+  "              ": 50274,
+  "               ": 50273,
+  "                ": 50272,
+  "                 ": 50271,
+  "                  ": 50270,
+  "                   ": 50269,
+  "                    ": 50268,
+  "                     ": 50267,
+  "                      ": 50266,
+  "                       ": 50265,
+  "                        ": 50264,
+  "                         ": 50263,
+  "                          ": 50262,
+  "                           ": 50261,
+  "                            ": 50260,
+  "                             ": 50259,
+  "                              ": 50258,
+  "                               ": 50257,
+  "<image>": 50297,
+  "<pad>": 50298,
+  "<|im_end|>": 50295,
+  "<|im_start|>": 50296
+}

config.json

@@ -0,0 +1,119 @@
+{
+  "architectures": [
+    "LlavaForConditionalGeneration"
+  ],
+  "ignore_index": -100,
+  "image_token_index": 50297,
+  "max_image_tokens": 100,
+  "model_type": "llava",
+  "projector_hidden_act": "gelu",
+  "projector_tokens_num": 5,
+  "text_config": {
+    "_name_or_path": "cognitivecomputations/dolphin-2_6-phi-2",
+    "activation_function": "gelu_new",
+    "add_cross_attention": false,
+    "architectures": [
+      "PhiForCausalLM"
+    ],
+    "attn_pdrop": 0.0,
+    "auto_map": {
+      "AutoConfig": "cognitivecomputations/dolphin-2_6-phi-2--configuration_phi.PhiConfig",
+      "AutoModelForCausalLM": "cognitivecomputations/dolphin-2_6-phi-2--modeling_phi.PhiForCausalLM"
+    },
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "embd_pdrop": 0.0,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "flash_attn": false,
+    "flash_rotary": false,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "fused_dense": false,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "img_processor": null,
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_epsilon": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "phi-msft",
+    "n_embd": 2560,
+    "n_head": 32,
+    "n_head_kv": null,
+    "n_inner": null,
+    "n_layer": 32,
+    "n_positions": 2048,
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "resid_pdrop": 0.1,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rotary_dim": 32,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": false,
+    "vocab_size": 51200
+  },
+  "preprocess_config": {
+    "mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "interpolation": "bicubic",
+    "resize_mode": "squash",
+    "size": 384
+  },
+  "torch_dtype": "float16",
+  "transformers_version": "4.36.2",
+  "vision_embed_dim": 1152,
+  "vision_tower_name": "ViT-SO400M-14-SigLIP-384",
+  "vocab_size": 51200
+}

configuration_llava.py

@@ -0,0 +1,41 @@
+# coding=utf-8
+
+from transformers.configuration_utils import PretrainedConfig
+from open_clip import get_model_config
+from configuration_phi import PhiConfig
+
+
+class LlavaConfig(PretrainedConfig):
+    model_type = "llava"
+    is_composition = False
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_tower_name="ViT-SO400M-14-SigLIP-384",
+        ignore_index=-100,
+        image_token_index=50297,
+        projector_hidden_act="gelu",
+        projector_tokens_num=1,
+        vocab_size=51200,
+        **kwargs,
+    ):
+        self.ignore_index = ignore_index
+        self.image_token_index = image_token_index
+        self.projector_hidden_act = projector_hidden_act
+        self.projector_tokens_num = projector_tokens_num
+        self.vocab_size = vocab_size
+
+        self.vision_tower_name = vision_tower_name
+        vision_config = get_model_config(vision_tower_name)
+        self.vision_embed_dim = vision_config["embed_dim"]
+
+        self.vocab_size = self.vocab_size
+
+        self.text_config = text_config
+        if isinstance(self.text_config, dict):
+            text_config["model_type"] = text_config["model_type"] if "model_type" in text_config else "llama"
+            self.text_config = PhiConfig(**text_config)
+            self.vocab_size = self.text_config.vocab_size
+
+        super().__init__(**kwargs)

configuration_phi.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import math
+from typing import Optional
+
+from transformers import PretrainedConfig
+
+
+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 = 51200,
+        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)

convert_model.py

@@ -0,0 +1,102 @@
+# Copyright 2023 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.
+import argparse
+
+import torch
+
+from transformers import (
+    AddedToken,
+    AutoConfig,
+    AutoTokenizer,
+)
+from configuration_llava import LlavaConfig
+from modeling_llava import LlavaForConditionalGeneration
+
+
+KEYS_TO_MODIFY_MAPPING = {
+    "transformer.vision_tower.vision_tower": "vision_model",
+    "transformer.mm_projector": "multi_modal_projector",
+    "transformer": "language_model.transformer",
+    "lm_head": "language_model.lm_head",
+    "model.model": "language_model.transformer",
+    "multi_modal_projector.0": "multi_modal_projector.linear_1",
+    "multi_modal_projector.2": "multi_modal_projector.linear_2",
+}
+
+
+def convert_state_dict_to_hf(state_dict):
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
+            if key_to_modify in key:
+                key = key.replace(key_to_modify, new_key)
+
+        new_state_dict[key] = value
+    return new_state_dict
+
+
+def convert_llava_llama_to_hf(text_model_id, vision_model_id, projector_tokens_num, output_path, old_state_dict_path):
+    torch.set_default_dtype(torch.float16)
+    text_config = AutoConfig.from_pretrained(text_model_id, trust_remote_code=True)
+
+    tokenizer = AutoTokenizer.from_pretrained(text_model_id)
+    tokenizer.add_tokens(AddedToken("<image>", special=True, normalized=False), special_tokens=True)
+    tokenizer.add_special_tokens({"pad_token": "<pad>"})
+
+    config = LlavaConfig(text_config=text_config, vocab_size=51200, vision_tower_name=vision_model_id, projector_tokens_num=projector_tokens_num)
+    config.text_config.vocab_size = config.vocab_size
+
+    with torch.device("cuda"):
+        model = LlavaForConditionalGeneration(config)
+    
+    state_dict = torch.load(old_state_dict_path, map_location="cpu")
+    state_dict = convert_state_dict_to_hf(state_dict)
+    model.load_state_dict(state_dict, strict=True, assign=True)
+
+    model.config.vocab_size = model.config.vocab_size
+    model.config.text_config.vocab_size = model.config.text_config.vocab_size
+
+    model.save_pretrained(output_path)
+    tokenizer.save_pretrained(output_path)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--text_model_id",
+        help="Hub location of the text model",
+    )
+    parser.add_argument(
+        "--vision_model_id",
+        help="Hub location of the vision model",
+    )
+    parser.add_argument(
+        "--output_path",
+        help="Location of the converted model",
+    )
+    parser.add_argument(
+        "--old_state_dict_path",
+        help="Location on the hub of the raw state dict of the original model. The filename needs to be `model_state_dict.bin`",
+    )
+    parser.add_argument(
+        "--tokens_num",
+        type=int,
+        default=1
+    )
+    args = parser.parse_args()
+    convert_llava_llama_to_hf(args.text_model_id, args.vision_model_id, args.tokens_num, args.output_path, args.old_state_dict_path)
+
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.36.2",
+  "use_cache": false
+}

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,678 @@
+{
+  "metadata": {
+    "total_size": 6773041280
+  },
+  "weight_map": {
+    "language_model.lm_head.linear.bias": "model-00002-of-00002.safetensors",
+    "language_model.lm_head.linear.weight": "model-00002-of-00002.safetensors",
+    "language_model.lm_head.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.lm_head.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.embd.wte.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.0.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.1.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.10.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.11.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.12.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.13.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.14.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.15.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.16.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.17.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.18.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.19.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.2.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.20.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.21.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.22.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.22.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.22.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.22.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.22.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.22.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.22.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.22.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.22.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.22.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.23.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.24.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.25.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.26.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.27.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.28.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.29.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.3.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.3.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.30.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.30.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.ln.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.ln.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mixer.Wqkv.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mixer.Wqkv.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mixer.out_proj.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mixer.out_proj.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mlp.fc1.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mlp.fc1.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mlp.fc2.bias": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.31.mlp.fc2.weight": "model-00002-of-00002.safetensors",
+    "language_model.transformer.h.4.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.4.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.5.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.6.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.7.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.8.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.ln.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.ln.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mixer.Wqkv.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mixer.Wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mixer.out_proj.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mixer.out_proj.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "language_model.transformer.h.9.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "multi_modal_projector.linear_1.bias": "model-00001-of-00002.safetensors",
+    "multi_modal_projector.linear_1.weight": "model-00001-of-00002.safetensors",
+    "multi_modal_projector.linear_2.bias": "model-00001-of-00002.safetensors",
+    "multi_modal_projector.linear_2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.kv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.kv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.latent": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.norm.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.norm.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.q.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.attn_pool.q.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.0.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.1.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.10.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.11.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.12.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.13.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.14.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.15.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.16.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.17.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.18.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.19.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.2.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.20.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.21.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.22.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.23.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.24.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.25.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.26.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.3.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.4.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.5.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.6.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.7.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.8.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.blocks.9.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.norm.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.norm.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.patch_embed.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.patch_embed.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.trunk.pos_embed": "model-00001-of-00002.safetensors"
+  }
+}

modeling_llava.py

@@ -0,0 +1,307 @@
+# coding=utf-8
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from modeling_phi import PhiForCausalLM
+from configuration_llava import LlavaConfig
+from open_clip import create_model
+
+
+@dataclass
+class LlavaCausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[List[torch.FloatTensor]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    image_features: Optional[torch.FloatTensor] = None
+
+
+class LlavaMultiModalProjector(nn.Module):
+    def __init__(self, config: LlavaConfig):
+        super().__init__()
+
+        self.linear_1 = nn.Linear(
+            config.vision_embed_dim,
+            config.text_config.n_embd * config.projector_tokens_num,
+            bias=True,
+        )
+        self.act = nn.GELU()
+        self.linear_2 = nn.Linear(
+            config.text_config.n_embd * 5,
+            config.text_config.n_embd,
+            bias=True,
+        )
+        self.projector_tokens_num = config.projector_tokens_num
+
+    def forward(self, image_features):
+        hidden_states = self.linear_1(image_features)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class LlavaPreTrainedModel(PreTrainedModel):
+    config_class = LlavaConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LlavaVisionAttention"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def _init_weights(self, module):
+        return
+
+    @property
+    def _supports_sdpa(self):
+        """
+        Retrieve language_model's attribute to check whether the model supports
+        SDPA or not.
+        """
+        return self.language_model._supports_sdpa
+
+
+class LlavaForConditionalGeneration(LlavaPreTrainedModel):
+    def __init__(self, config: LlavaConfig):
+        super().__init__(config)
+        clip_model = create_model(config.vision_tower_name)
+        self.vision_model = clip_model.visual
+
+        self.multi_modal_projector = LlavaMultiModalProjector(config)
+        self.vocab_size = config.vocab_size
+        self.language_model = PhiForCausalLM(config.text_config)
+        self.pad_token_id = (
+            self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        )
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.language_model.transformer = decoder
+
+    def get_decoder(self):
+        return self.language_model.transformer
+
+    def tie_weights(self):
+        return self.language_model.tie_weights()
+
+    def resize_token_embeddings(
+        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of=None
+    ) -> nn.Embedding:
+        model_embeds = self.language_model.resize_token_embeddings(
+            new_num_tokens, pad_to_multiple_of
+        )
+        # update vocab size
+        self.config.text_config.vocab_size = model_embeds.num_embeddings
+        self.config.vocab_size = model_embeds.num_embeddings
+        self.vocab_size = model_embeds.num_embeddings
+        return model_embeds
+
+    def _merge_input_ids_with_image_features(
+        self, image_features, inputs_embeds, input_ids, attention_mask, position_ids
+    ):
+        num_images, num_image_patches, embed_dim = image_features.shape
+        batch_size, sequence_length = input_ids.shape
+        left_padding = not torch.sum(
+            input_ids[:, -1] == torch.tensor(self.pad_token_id)
+        )
+        # 1. Create a mask to know where special image tokens are
+        special_image_token_mask = input_ids == self.config.image_token_index
+        num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1)
+        # Compute the maximum embed dimension
+        max_embed_dim = (
+            num_special_image_tokens.max() * (num_image_patches - 1)
+        ) + sequence_length
+        batch_indices, non_image_indices = torch.where(
+            input_ids != self.config.image_token_index
+        )
+
+        # 2. Compute the positions where text should be written
+        # Calculate new positions for text tokens in merged image-text sequence.
+        # `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens.
+        # `torch.cumsum` computes how each image token shifts subsequent text token positions.
+        # - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one.
+        new_token_positions = (
+            torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1)
+            - 1
+        )
+        nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1]
+        if left_padding:
+            new_token_positions += nb_image_pad[:, None]  # offset for left padding
+        text_to_overwrite = new_token_positions[batch_indices, non_image_indices]
+
+        # 3. Create the full embedding, already padded to the maximum position
+        final_embedding = torch.zeros(
+            batch_size,
+            max_embed_dim,
+            embed_dim,
+            dtype=inputs_embeds.dtype,
+            device=inputs_embeds.device,
+        )
+        final_attention_mask = torch.zeros(
+            batch_size,
+            max_embed_dim,
+            dtype=attention_mask.dtype,
+            device=inputs_embeds.device,
+        )
+        # In case the Vision model or the Language model has been offloaded to CPU, we need to manually
+        # set the corresponding tensors into their correct target device.
+        target_device = inputs_embeds.device
+        batch_indices, non_image_indices, text_to_overwrite = (
+            batch_indices.to(target_device),
+            non_image_indices.to(target_device),
+            text_to_overwrite.to(target_device),
+        )
+        attention_mask = attention_mask.to(target_device)
+
+        # 4. Fill the embeddings based on the mask. If we have ["hey" "<image>", "how", "are"]
+        # we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features
+        final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[
+            batch_indices, non_image_indices
+        ]
+        final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[
+            batch_indices, non_image_indices
+        ]
+
+        # 5. Fill the embeddings corresponding to the images. Anything that is still zeros needs filling
+        image_to_overwrite = torch.all(final_embedding == 0, dim=-1)
+        image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[
+            :, None
+        ].to(target_device)
+
+        if image_to_overwrite.sum() != image_features.shape[:-1].numel():
+            raise ValueError(
+                f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while"
+                f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation."
+            )
+
+        final_embedding[image_to_overwrite] = (
+            image_features.contiguous().reshape(-1, embed_dim).to(target_device)
+        )
+        final_attention_mask |= image_to_overwrite
+        position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_(
+            (final_attention_mask == 0), 1
+        )
+        return final_embedding, final_attention_mask, position_ids
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        image_features: torch.FloatTensor = 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,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, LlavaCausalLMOutputWithPast]:
+        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 inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+            if image_features is not None and input_ids.shape[1] != 1:
+                (
+                    inputs_embeds,
+                    attention_mask,
+                    position_ids,
+                ) = self._merge_input_ids_with_image_features(
+                    image_features,
+                    inputs_embeds,
+                    input_ids,
+                    attention_mask,
+                    position_ids,
+                )
+
+        outputs = self.language_model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        logits = outputs[0]
+
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return output
+
+        return LlavaCausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_features=image_features,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        attention_mask=None,
+        image_features=None,
+        **kwargs,
+    ):
+        res = self.language_model.prepare_inputs_for_generation(input_ids, past_key_values, attention_mask, **kwargs)
+        input_ids = res["input_ids"]
+        past_key_values = res["past_key_values"]
+        attention_mask = res["attention_mask"]
+
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "image_features": image_features,
+            }
+        )
+        return model_inputs
+
+    def _reorder_cache(self, *args, **kwargs):
+        return self.language_model._reorder_cache(*args, **kwargs)

modeling_phi.py

@@ -0,0 +1,988 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+#
+# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
+# Licensed under the BSD 3-Clause License.
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from configuration_phi import PhiConfig
+
+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
+    print("Using Flash Attention!")
+except Exception as exc:
+    print(exc)
+    pad_input, unpad_input = None, None
+    FlashRotaryEmbedding = None
+    FlashSelfAttention, FlashCrossAttention = None, None
+    FusedDense = None
+    print("Not using Flash Attention!")
+
+
+@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 = True,
+    ) -> 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, None, cu_seqlens, max_seqlen, use_reentrant=False
+                )
+            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, None, key_padding_mask, use_reentrant=False)
+
+        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,
+                    cu_seqlens_q,
+                    max_seqlen_q,
+                    cu_seqlens_k,
+                    max_seqlen_k,
+                    use_reentrant=False,
+                )
+            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,
+                causal,
+                key_padding_mask, 
+                use_reentrant=False,
+            )
+
+        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 CausalLMLoss(nn.Module):
+    """Causal Language Modeling loss.
+
+    Reference:
+        Improving Language Understanding by Generative Pre-Training.
+        https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.
+
+    """
+
+    def __init__(self, shift_labels: bool = True) -> None:
+        super().__init__()
+
+        self.shift_labels = shift_labels
+        self.loss_fct = nn.CrossEntropyLoss()
+
+    def forward(self, logits: torch.FloatTensor, labels: torch.LongTensor) -> torch.FloatTensor:
+        if self.shift_labels:
+            logits = logits[..., :-1, :].contiguous()
+            labels = labels[..., 1:].contiguous()
+
+        loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        return loss
+
+
+class PhiPreTrainedModel(PreTrainedModel):
+    """Phi pre-trained model."""
+
+    config_class = PhiConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["ParallelBlock"]
+
+    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`
+        #     past_key_values.seqlen_offset = input_ids.shape[1] - 1
+        #     input_ids = input_ids[:, -1].unsqueeze(-1)
+        #     attention_mask = attention_mask[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+        }
+
+
+class PhiModel(PhiPreTrainedModel):
+    """Phi model."""
+
+    _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: PhiConfig) -> None:
+        config.flash_attn = True
+        config.flash_rotary = True
+        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 = True
+        self.post_init()
+
+    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,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+    ) -> torch.FloatTensor:
+        if input_ids is not None:
+            hidden_states = self.embd(input_ids)
+        elif inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        for layer in self.h:
+            if self.gradient_checkpointing:
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    layer.__call__,
+                    hidden_states,
+                    past_key_values,
+                    attention_mask,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states = layer(
+                    hidden_states,
+                    past_key_values=past_key_values,
+                    attention_mask=attention_mask,
+                )
+
+        return hidden_states
+
+
+class PhiForCausalLM(PhiPreTrainedModel):
+    """Phi 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)"]
+    
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["ParallelBlock"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def __init__(self, config: PhiConfig) -> None:
+        super().__init__(config)
+
+        self.transformer = PhiModel(config)
+        self.lm_head = CausalLMHead(config)
+        self.loss = CausalLMLoss()
+
+        self.post_init()
+
+    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 forward(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states = self.transformer(input_ids, inputs_embeds=inputs_embeds, past_key_values=past_key_values, attention_mask=attention_mask)
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss(lm_logits, labels)
+
+        return CausalLMOutputWithPast(loss=loss, logits=lm_logits, past_key_values=past_key_values)

processing_llava.py

@@ -0,0 +1,152 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# 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 Llava.
+"""
+
+
+from typing import List, Optional, Union
+
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.tokenization_utils_base import (
+    PaddingStrategy,
+    PreTokenizedInput,
+    TextInput,
+    TruncationStrategy,
+)
+from transformers.utils import TensorType
+import torch
+from open_clip.transform import PreprocessCfg, image_transform_v2
+from modeling_llava import LlavaForConditionalGeneration
+from PIL import Image
+import math
+
+
+class OpenCLIPImageProcessor:
+    def __init__(self, config, crop_size=384, max_tokens=100):
+        cfg = PreprocessCfg(**config)
+        transform = image_transform_v2(cfg=cfg, is_train=False)
+        self.transform = transform
+        self.crop_size = crop_size
+        self.max_tokens = max_tokens
+
+    def __call__(self, image: Image.Image):
+        output = self.transform_func(image)
+        return {
+            "pixel_values": output,
+        }
+
+    def transform_func(self, image: Image.Image):
+        outputs = []
+        outputs.append(self.transform(image))
+        width, height = image.size
+        crop_size = self.crop_size
+        if width <= crop_size and height <= crop_size:
+            outputs = torch.stack(outputs, dim=0)
+            return outputs
+        total_tokens = math.inf
+        while total_tokens > self.max_tokens:
+            total_tokens = math.floor(
+                (2 * width - crop_size)
+                / crop_size
+                * (2 * height - crop_size)
+                / crop_size
+            )
+            if total_tokens > self.max_tokens:
+                crop_size += 10
+        stride = crop_size // 2
+        x_steps = int(round((2 * width - crop_size) / crop_size))
+        if x_steps < 1:
+            x_steps = 1
+        y_steps = int(round((2 * height - crop_size) / crop_size))
+        if y_steps < 1:
+            y_steps = 1
+        x_coords = []
+        y_coords = []
+        for i in range(x_steps):
+            x_coords.append([i * stride, i * stride + crop_size])
+        if x_coords[-1][1] != width:
+            x_coords[-1][1] = width
+        for i in range(y_steps):
+            y_coords.append([i * stride, i * stride + crop_size])
+        if y_coords[-1][1] != height:
+            y_coords[-1][1] = height
+        image_parts = []
+        for i in range(len(x_coords)):
+            for j in range(len(y_coords)):
+                image_parts.append(
+                    image.crop(
+                        (x_coords[i][0], y_coords[j][0], x_coords[i][1], y_coords[j][1])
+                    )
+                )
+        for image_part in image_parts:
+            outputs.append(self.transform(image_part))
+        outputs = torch.stack(outputs, dim=0)
+        return outputs
+
+    @property
+    def model_input_names(self):
+        return ["pixel_values"]
+
+
+class LlavaProcessor:
+    def __init__(self, image_processor: OpenCLIPImageProcessor, tokenizer):
+        self.image_processor = image_processor
+        self.tokenizer = tokenizer
+
+    def __call__(
+        self,
+        text: Union[
+            TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]
+        ] = None,
+        images: ImageInput = None,
+        model: LlavaForConditionalGeneration = None,
+        padding: Union[bool, str, PaddingStrategy] = False,
+        truncation: Union[bool, str, TruncationStrategy] = None,
+        max_length=None,
+        return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
+    ) -> BatchFeature:
+        if images is not None:
+            pixel_values = self.image_processor(images)[
+                "pixel_values"
+            ]
+            pixel_values = pixel_values.to(model.device).to(model.dtype)
+            image_outputs = model.vision_model(pixel_values)
+            image_features = model.multi_modal_projector(image_outputs)
+            image_features = image_features.unsqueeze(0)
+        else:
+            image_features = None
+        text_inputs = self.tokenizer(
+            text,
+            return_tensors=return_tensors,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+        )
+
+        return BatchFeature(data={**text_inputs, "image_features": image_features})
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,357 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50257": {
+      "content": "                               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "                         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50280": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50281": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50282": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50283": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50284": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50285": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50286": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50287": {
+      "content": "\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50288": {
+      "content": "\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50289": {
+      "content": "\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50290": {
+      "content": "\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50291": {
+      "content": "\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50292": {
+      "content": "\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50293": {
+      "content": "\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50294": {
+      "content": "\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50295": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50296": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50297": {
+      "content": "<image>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50298": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|im_end|>",
+  "model_max_length": 1200,
+  "pad_token": "<pad>",
+  "tokenizer_class": "CodeGenTokenizer",
+  "unk_token": "<|endoftext|>"
+}