Upload turing-motors/Heron-NVILA-Lite-33B

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+llm/tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,150 @@
+---
+license: apache-2.0
+license_link: https://huggingface.co/Qwen/Qwen2.5-32B-Instruct/blob/main/LICENSE
+language:
+- ja
+- en
+tags:
+- vila
+- nvila
+- conversational
+- multimodal
+base_model:
+- Qwen/Qwen2.5-32B-Instruct
+- google/siglip2-so400m-patch16-512
+---
+# Heron-NVILA-Lite-33B
+
+Heron-NVILA-Lite-33B is a vision language model trained for Japanese, based on the [NVILA](https://arxiv.org/abs/2412.04468)-Lite architecture.
+
+## Model Overview
+
+* **Developer**: [Turing Inc.](https://www.turing-motors.com/)
+* **Vision Encoder**: [siglip2-so400m-patch16-512](https://huggingface.co/google/siglip2-so400m-patch16-512)
+* **Projector**: mlp_downsample_2x2_fix
+* **LLM**: [Qwen2.5-32B-Instruct](https://huggingface.co/Qwen/Qwen2.5-32B-Instruct)
+* **Supported Languages**: Japanese, English
+
+## Setup
+
+```bash
+# I have confirmed that 4.46.0 and 4.49.0 also work. Other versions of Transformer may also work, but I have not tested them.
+pip install transformers==4.45.0 accelerate opencv-python torchvision einops pillow
+pip install git+https://github.com/bfshi/scaling_on_scales.git
+```
+
+## Usage
+
+```python
+from transformers import AutoConfig, AutoModel
+
+model_path = "turing-motors/Heron-NVILA-Lite-33B"
+
+# you can use config
+config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModel.from_config(config, trust_remote_code=True, device_map="auto")
+
+# or directly from_pretrained
+model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
+
+# show chat_template
+print(model.tokenizer.chat_template)
+
+# examples generate with raw text
+response = model.generate_content(["こんにちは"])
+print(response)
+print("---" * 40)
+
+# examples generate with text + image
+from PIL import Image
+import requests
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+response = model.generate_content([image, "画像を説明してください。"])
+print(response)
+print("---" * 40)
+
+# examples generate using generation_config
+from PIL import Image
+import requests
+from transformers import GenerationConfig
+generation_config = {
+    "max_new_tokens": 512,
+    "temperature": 0.5,
+    "do_sample": True,
+}
+generation_config = GenerationConfig(**generation_config)
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+response = model.generate_content(
+    [image, "画像を説明してください。"],
+    generation_config=generation_config
+)
+print(response)
+print("---" * 40)
+
+# examples generate with text + image + text + image + text
+from PIL import Image
+import requests
+url_list = [
+    "https://images.unsplash.com/photo-1694831404826-3400c48c188d?q=80&w=2070&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D",
+    "https://images.unsplash.com/photo-1693240876439-473af88b4ed7?q=80&w=1974&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
+]
+images = [
+   Image.open(requests.get(url, stream=True).raw).convert("RGB") for url in url_list
+]
+response = model.generate_content([
+    images[0],
+    "これは日本の画像です",
+    images[1],
+    "これはオーストリアの画像です",
+    "各画像の違いを説明して"])
+print(response)
+print("---" * 40)
+```
+
+## Training Summary
+
+| Stage  | Training                      | Data Sources                  | Samples     |
+|--------|-------------------------------|-------------------------------|-------------|
+| Stage1 | Projector                     | [Japanese image text pairs](https://gitlab.llm-jp.nii.ac.jp/datasets/llm-jp-japanese-image-text-pairs), [LLaVA-Pretrain](https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain)                          | 1.1M      |
+| Stage2 | Projector, LLM                | Filtered [MOMIJI](https://huggingface.co/datasets/turing-motors/MOMIJI) (CC-MAIN-2024-42)  | 3M     |
+|        |                               | [Japanese image text pairs (subset)](https://gitlab.llm-jp.nii.ac.jp/datasets/llm-jp-japanese-image-text-pairs), [Japanese interleaved data (subset)](https://gitlab.llm-jp.nii.ac.jp/datasets/llm-jp-japanese-interleaved-data), [mmc4-core (subset)](https://github.com/allenai/mmc4), [coyo-700m (subset)](https://huggingface.co/datasets/kakaobrain/coyo-700m), [wikipedia_ja](https://huggingface.co/datasets/turing-motors/Wikipedia-Vision-JA), [llava_pretrain_ja](https://huggingface.co/datasets/turing-motors/LLaVA-Pretrain-JA), [stair_captions](http://captions.stair.center/)  | 20M     |
+| Stage3 | Vision Encoder, Projector, LLM | [llava-instruct-v1_5-en-subset-358k](https://huggingface.co/datasets/llm-jp/llava-instruct-v1_5-en-subset-358k), [llava-instruct-ja](https://huggingface.co/datasets/llm-jp/llava-instruct-ja), [japanese-photos-conv](https://huggingface.co/datasets/llm-jp/japanese-photos-conversation), [ja-vg-vqa](https://huggingface.co/datasets/llm-jp/ja-vg-vqa-conversation), [synthdog-ja (subset)](https://huggingface.co/datasets/naver-clova-ix/synthdog-ja), [ai2d](https://huggingface.co/datasets/lmms-lab/ai2d), [synthdog-en](https://huggingface.co/datasets/naver-clova-ix/synthdog-en), [sherlock](https://github.com/allenai/sherlock)    | 1.1M      |
+
+## Evaluation
+
+I used [llm-jp-eval-mm](https://github.com/llm-jp/llm-jp-eval-mm) for this evaluation. Scores for models other than Heron-NVILA-Lite and Sarashina2-Vision-14B were taken from [llm-jp-eval-mm leaderboard](https://llm-jp.github.io/llm-jp-eval-mm/) as of March 2025 and the [Asagi website](https://uehara-mech.github.io/asagi-vlm?v=1). Heron-NVILA-Lite and Sarashina2-Vision-14B were evaluated using llm-as-a-judge with "gpt-4o-2024-05-13". Sarashina2-Vision-14B was evaluated on the [official blog](https://www.sbintuitions.co.jp/blog/entry/2025/03/17/111703) using "gpt-4o-2024-08-06"; please note that due to differing evaluation conditions, the results for Sarashina2-Vision-14B should be treated as reference only.
+
+| Model                          | LLM Size | Heron-Bench overall LLM (%) | JA-VLM-Bench-In-the-Wild LLM (/5.0) | JA-VG-VQA-500 LLM (/5.0) |
+|--------------------------------|----------|------------------------------|-------------------------------------|--------------------------|
+| **[Heron-NVILA-Lite-1B](https://huggingface.co/turing-motors/Heron-NVILA-Lite-1B)**        | 0.5B     | 45.9                         | 2.92                                | 3.16                     |
+| **[Heron-NVILA-Lite-2B](https://huggingface.co/turing-motors/Heron-NVILA-Lite-2B)**        | 1.5B     | 52.8                         | 3.52                                | 3.50                     |
+| **[Heron-NVILA-Lite-15B](https://huggingface.co/turing-motors/Heron-NVILA-Lite-15B)**       | 14B      | 59.6                         | 4.2                                 | 3.82                     |
+| **Heron-NVILA-Lite-33B**       | 32B      | 61.1                         | 4.0                                 | 3.85                     |
+| [LLaVA-CALM2-SigLIP](https://huggingface.co/cyberagent/llava-calm2-siglip)             | 7B      | 43.3                        | 3.15                                | 3.21                     |
+| [Llama-3-EvoVLM-JP-v2](https://huggingface.co/SakanaAI/Llama-3-EvoVLM-JP-v2)           | 8B      | 39.3                        | 2.92                                | 2.96                     |
+| [VILA-jp](https://huggingface.co/llm-jp/llm-jp-3-vila-14b)                        | 13B     | 57.2                        | 3.69                                | 3.62                     |
+| [Asagi-14B](https://huggingface.co/MIL-UT/Asagi-14B)                      | 13B     | 55.8                        | 3.44                                | 3.84                     |
+| [Sarashina2-Vision-14B](https://huggingface.co/sbintuitions/sarashina2-vision-14b)                      | 13B     | 50.9                        | 4.1                                | 3.43                     |
+| [Qwen2-VL 7B Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct)                         | 7B       | 55.5                        | 3.61                                | 3.6                     |
+| GPT-4o                         | -       | 87.6                        | 3.85                                | 3.58                     |
+
+
+## Risks and Limitations
+
+This model is experimental and has not been thoroughly calibrated for ethical compliance or legal standards. Caution is advised for sensitive applications.
+
+## License
+
+- Model weights are licensed under [Apache License 2.0](https://huggingface.co/Qwen/Qwen2.5-32B-Instruct/blob/main/LICENSE).
+- Users must comply with  [OpenAI terms of use](https://openai.com/policies/terms-of-use) due to the inclusion of GPT-4-generated synthetic data.
+
+## Acknowledgements
+
+This model is based on results obtained from a project, JPNP20017, subsidized by the New Energy and Industrial Technology Development Organization (NEDO).
+
+I would like to acknowledge the use of the following open-source repositories:
+
+- [VILA](https://github.com/NVlabs/VILA)
+- [llm-jp-eval-mm](https://github.com/llm-jp/llm-jp-eval-mm)

auto_processor.py

@@ -0,0 +1,330 @@
+import copy
+import os
+import os.path as osp
+import warnings
+from collections import defaultdict
+from typing import List, Union
+
+import torch
+from transformers import AutoConfig, AutoImageProcessor, AutoModel, AutoProcessor, AutoTokenizer
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput, VideoInput
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.utils import logging
+
+from .constants import DEFAULT_IMAGE_TOKEN, MEDIA_TOKENS
+from .media import Image, Video, extract_media
+from .mm_utils import process_image, process_images
+from .tokenizer_utils import tokenize_conversation
+
+def fetch_image_url_or_fpath(url_or_fpath):
+    if url_or_fpath.startswith("http") or url_or_fpath.startswith("https"):
+        import tempfile
+        import requests
+        
+        # Download the image to a temporary file
+        temp_dir = tempfile.mkdtemp()
+        temp_file = os.path.join(temp_dir, os.path.basename(url_or_fpath))
+        
+        response = requests.get(url_or_fpath, stream=True)
+        response.raise_for_status()
+        
+        with open(temp_file, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+                
+        return temp_file
+    elif url_or_fpath.startswith("file://"):
+        fpath = url_or_fpath.replace("file://", "")
+        assert osp.exists(fpath), f"File {fpath} does not exist"
+        return fpath
+    elif osp.exists(url_or_fpath):
+        assert osp.isfile(url_or_fpath), f"File {url_or_fpath} is not a file"
+        return url_or_fpath
+    else:
+        raise ValueError(f"Unsupported image path: {url_or_fpath}")
+    
+
+def __pad_fn(input_ids_list, padding_value=0, target_len=None, padding_side="left"):
+    # tensor shape is (batch_size, seq_len)
+    max_len = max([ids.shape[1] for ids in input_ids_list])
+    if target_len is not None:
+        assert target_len >= max_len, "target_len must be greater than or equal to max_len"
+        max_len = target_len
+
+    new_input_ids_list = []
+    for i, input_ids in enumerate(input_ids_list):
+        pad_tensor = torch.ones_like(input_ids) * padding_value
+        curr_len = input_ids.shape[1]
+        pad_tensor = pad_tensor[:, : max_len - curr_len]
+        if padding_side == "right":
+            input_ids = torch.cat((input_ids, pad_tensor), dim=1)
+        else:
+            input_ids = torch.cat((pad_tensor, input_ids), dim=1)
+        new_input_ids_list.append(input_ids)
+    return torch.cat(new_input_ids_list, dim=0)
+
+
+class VILAProcessorKwargs(ProcessingKwargs, total=False):
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+    }
+
+
+
+
+class VILAProcessor(ProcessorMixin):
+    # attributes = ["image_processor", "tokenizer"]
+    attributes = []
+    # valid_kwargs = ["chat_template"]
+    valid_kwargs = []
+    # image_processor_class = "VILAImageProcessor"
+    # tokenizer_class = ("VILATokenizer", "VILATokenizerFast")
+
+    def __init__(self, image_processor=None, tokenizer=None, chat_template=None, config=None, **kwargs):
+        # self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
+        # self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
+        self.image_token = MEDIA_TOKENS["image"]
+        self.video_token = MEDIA_TOKENS["video"]
+        self.config = config
+        self.image_processor = image_processor
+        self.tokenizer = tokenizer
+
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        if os.path.isdir(pretrained_model_name_or_path):
+            pretrained_model_name_or_path = pretrained_model_name_or_path
+        else:
+            print(f"pretrained_model_name_or_path {pretrained_model_name_or_path} is not a directory, downloading")
+            from huggingface_hub import snapshot_download
+
+            pretrained_model_name_or_path = snapshot_download(pretrained_model_name_or_path)
+
+        image_processor = AutoImageProcessor.from_pretrained(
+            osp.join(pretrained_model_name_or_path, "vision_tower"), trust_remote_code=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            osp.join(pretrained_model_name_or_path, "llm"), trust_remote_code=True
+        )
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        return cls(image_processor=image_processor, tokenizer=tokenizer, config=config)
+
+    def __repr__(self):
+        return (
+            f"VILAProcessor(image_processor={self.image_processor}, tokenizer={self.tokenizer}, config={self.config})"
+        )
+
+    def __call__(
+        self,
+        conversation,
+        images: ImageInput = None,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+        videos: VideoInput = None,
+        **kwargs: Unpack[VILAProcessorKwargs],
+    ) -> BatchFeature:
+        if images is not None:
+            warnings.warn("images is not supported in __call__")
+
+        input_ids = []
+        media = defaultdict(list)
+        media_config = defaultdict(dict)
+        for conv in conversation:
+            feat = self.__single_call__(conv, images, text, videos, **kwargs)
+            input_ids.append(feat.input_ids)
+            for name in feat.media:
+                media[name] += feat.media[name]
+            for name in feat.media_config:
+                media_config[name].update(feat.media_config[name])
+
+        return BatchFeature(
+            data={
+                # "input_ids": torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=self.pad_token_id),
+                "input_ids": __pad_fn(
+                    input_ids,
+                    padding_value=self.tokenizer.pad_token_id,
+                    padding_side="left",
+                ),
+                "media": media,
+                "media_config": media_config,
+            }
+        )
+
+    def __single_call__(
+        self,
+        conversation,
+        images: ImageInput = None,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+        videos: VideoInput = None,
+        **kwargs: Unpack[VILAProcessorKwargs],
+    ) -> BatchFeature:
+        # TODO: should be merged with llava_arch.py/generate_content()
+        # TODO (extract and preprocess should be done together, as the preprocess of image and video can be different, i.e. when dynamic res is used)
+        conversation = copy.deepcopy(conversation)
+        media = extract_media(conversation, self.config)
+        # Process media
+        media_config = defaultdict(dict)
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in ["dynamic", "dynamic_s2"]:
+                    self.config.image_processor = self.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(media["image"][0], self.config, None, enable_dynamic_res=True).half()
+                        conversation[0]["value"] = conversation[0]["value"].replace(
+                            DEFAULT_IMAGE_TOKEN, f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0]
+                        )
+                    else:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(map(int, self.config.s2_scales.split(",")))
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(media["image"], self.vision_tower.image_processor, self.config).half()
+                media[name] = [image for image in images]
+            elif name == "video":
+                media[name] = [
+                    process_images(images, self.vision_tower.image_processor, self.config).half()
+                    for images in media[name]
+                ]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+        input_ids = tokenize_conversation(conversation, self.tokenizer, add_generation_prompt=True).cuda().unsqueeze(0)
+        # Set up the generation config
+        return BatchFeature(data={"input_ids": input_ids, "media": media, "media_config": media_config})
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def post_process_image_text_to_text(self, generated_outputs):
+        """
+        Post-process the output of the model to decode the text.
+
+        Args:
+            generated_outputs (`torch.Tensor` or `np.ndarray`):
+                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+                or `(sequence_length,)`.
+
+        Returns:
+            `List[str]`: The decoded text.
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+
+    @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))
+
+    #     inputs = processor(conversation=llavaconv, padding=True, return_tensors="pt")
+    def apply_chat_template(self, conversation, add_generation_prompt=True, **kwargs):
+        vila_conv = []
+        for chat in conversation:
+            vila_chat = {"from": "", "value": []}
+            if chat["role"] == "user":
+                # user allows to input image and text
+                vila_chat["from"] = "human"
+                for content in chat["content"]:
+                    if content["type"] == "image":
+                        if "path" in content:
+                            # VILA style
+                            vila_chat["value"].append(Image(fetch_image_url_or_fpath(content["path"])))
+                        elif "image" in content:
+                            # Qwen style
+                            vila_chat["value"].append(Image(fetch_image_url_or_fpath(content["image"])))
+                        else:
+                            raise ValueError(f"Unsupported content type `image`: {content}, `image` and `path` are required")
+                    elif content["type"] == "text":
+                        vila_chat["value"].append(content["text"])
+                    # NOTE(ligeng): video supports are needed here
+                    else:
+                        raise ValueError(f"Unsupported content type: {content['type']}")
+            elif chat["role"] == "assistant":
+                vila_chat["from"] = "gpt"
+                for content in chat["content"]:
+                    assert content["type"] == "text", f"Unsupported content type: {content['type']}"
+                    vila_chat["value"].append(content["text"])
+            vila_conv.append(vila_chat)
+
+        return vila_conv
+
+
+if __name__ == "__main__":
+    # gpt style: user, assistant
+    # vila style: human, gpt
+    gpt_conv = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "path": "demo_images/demo_img_1.png"},
+                {"type": "text", "text": "Describe this image."},
+            ],
+        }
+    ]
+
+    llavaconv = [
+        {
+            "from": "human",
+            "value": [
+                PIL.Image.open("demo_images/demo_img_1.png"),
+                "Describe this image.",
+            ],
+        }
+    ]
+
+    processor = AutoProcessor.from_pretrained(output_dir, trust_remote_code=True)
+    inputs = processor.apply_chat_template(conversation=gpt_conv, padding=True, return_tensors="pt")
+    # model = llava.load("Efficient-Large-Model/qwen25_2B_3x3-sft").cuda()
+    # print(model)
+    model_path = "NVILA-Lite-2B-hf-preview"
+    model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
+    # res = model.generate_content(["how are you today?"])
+    # print(model.config)
+    # print(model.tokenizer)
+    # print(res)
+    # exit(0)
+
+    processor = VILAProcessor(
+        config=model.config,
+        image_processor=model.vision_tower.image_processor,
+        tokenizer=model.tokenizer,
+    )
+
+    # TODO: add padding, return_tensors,
+    inputs = processor(conversation=llavaconv, padding=True, return_tensors="pt")
+    print(inputs.keys(), inputs.input_ids.shape, [_.shape for _ in inputs.image])
+    print("vila conv pass")
+
+    inputs = processor.apply_chat_template(conversation=gpt_conv, padding=True, return_tensors="pt")
+    print(inputs.keys(), inputs.input_ids.shape, [_.shape for _ in inputs.image])
+    print("gpt conv pass")
+
+    output_ids = model.generate(
+        input_ids=inputs.input_ids,
+        media={
+            "image": inputs.image,
+        },
+        media_config={"image": {}},
+        generation_config=model.generation_config,
+        max_new_tokens=100,
+    )
+    print(output_ids)

base_projector.py

@@ -0,0 +1,228 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+
+
+class IdentityMap(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    @property
+    def config(self):
+        return {"mm_projector_type": "identity"}
+
+
+class SimpleResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(channels)
+
+        self.proj = nn.Sequential(nn.Linear(channels, channels), nn.GELU(), nn.Linear(channels, channels))
+
+    def forward(self, x):
+        x = self.pre_norm(x)
+        return x + self.proj(x)
+
+
+class DownSampleBlock(nn.Module):
+    def forward(self, x):
+        vit_embeds = x
+        h = w = int(vit_embeds.shape[1] ** 0.5)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        vit_embeds = self.flat_square(vit_embeds)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
+        return vit_embeds
+
+    def flat_square(self, x):
+        n, w, h, c = x.size()
+        if w % 2 == 1:
+            x = torch.concat([x, torch.zeros((n, 1, h, c), dtype=x.dtype).to(x.device)], dim=1).contiguous()
+            n, w, h, c = x.size()
+        if h % 2 == 1:
+            x = torch.concat([x, torch.zeros((n, w, 1, c), dtype=x.dtype).to(x.device)], dim=2).contiguous()
+            n, w, h, c = x.size()
+        x = x.contiguous()
+        x = x.view(n, w, int(h / 2), int(c * 2))
+        x = x.permute(0, 2, 1, 3).contiguous()
+        x = x.view(n, int(h / 2), int(w / 2), int(c * 4))
+        x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+
+
+class DownSample2x2BlockFix(nn.Module):
+    def forward(self, x):
+        vit_embeds = x
+        h = w = int(vit_embeds.shape[1] ** 0.5)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        vit_embeds = flat_square_2x2(vit_embeds)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
+        return vit_embeds
+
+
+def flat_square_2x2(x):
+    n, w, h, c = x.size()
+    if w % 2 == 1:
+        x = torch.concat([x, torch.zeros((n, 1, h, c), dtype=x.dtype).to(x.device)], dim=1).contiguous()
+        n, w, h, c = x.size()
+    x = x.contiguous()
+    if h % 2 == 1:
+        x = torch.concat([x, torch.zeros((n, w, 1, c), dtype=x.dtype).to(x.device)], dim=2).contiguous()
+        n, w, h, c = x.size()
+    x = x.view(n, w, int(h / 2), int(c * 2))
+    x = x.permute(0, 2, 1, 3).contiguous()
+    x = x.view(n, int(h / 2), int(w / 2), int(c * 4))
+    x = x.permute(0, 2, 1, 3).contiguous()
+    return x
+
+
+class DownSample3x3BlockFix(nn.Module):
+    def forward(self, x):
+        vit_embeds = x
+        h = w = int(vit_embeds.shape[1] ** 0.5)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        vit_embeds = flat_square_3x3(vit_embeds)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
+        return vit_embeds
+
+
+def flat_square_3x3(x):
+    n, w, h, c = x.size()
+    if w % 3 != 0:
+        x = torch.concat([x, torch.zeros((n, 3 - (w % 3), h, c), dtype=x.dtype).to(x.device)], dim=1).contiguous()
+        n, w, h, c = x.size()
+    x = x.contiguous()
+    if h % 3 != 0:
+        x = torch.concat([x, torch.zeros((n, w, 3 - (h % 3), c), dtype=x.dtype).to(x.device)], dim=2).contiguous()
+        n, w, h, c = x.size()
+    x = x.view(n, w, int(h / 3), int(c * 3))
+    x = x.permute(0, 2, 1, 3).contiguous()
+    x = x.view(n, int(h / 3), int(w / 3), int(c * 9))
+    x = x.permute(0, 2, 1, 3).contiguous()
+    return x
+
+
+class MultimodalProjectorConfig(PretrainedConfig):
+    model_type = "v2l_projector"
+
+    def __init__(self, mm_projector_type: str = None, **kwargs):
+        super().__init__()
+        self.mm_projector_type = mm_projector_type
+
+
+class MultimodalProjector(PreTrainedModel):
+    config_class = MultimodalProjectorConfig
+
+    def __init__(self, mm_projector_cfg: MultimodalProjectorConfig, config: PretrainedConfig):
+        super().__init__(mm_projector_cfg)
+        mm_projector_type = mm_projector_cfg.mm_projector_type
+        self.downsample_rate = 1
+        if mm_projector_type == "identity":
+            self.layers = IdentityMap()
+        elif mm_projector_type == "linear":
+            self.layers = nn.Linear(config.mm_hidden_size, config.hidden_size)
+        elif mm_projector_type == "mlp_downsample":
+            self.layers = nn.Sequential(
+                DownSampleBlock(),
+                nn.LayerNorm(config.mm_hidden_size * 4),
+                nn.Linear(config.mm_hidden_size * 4, config.hidden_size),
+                nn.GELU(),
+                nn.Linear(config.hidden_size, config.hidden_size),
+            )
+            self.downsample_rate = 2
+        elif mm_projector_type == "mlp_downsample_2x2_fix":
+            self.layers = nn.Sequential(
+                DownSample2x2BlockFix(),
+                nn.LayerNorm(config.mm_hidden_size * 4),
+                nn.Linear(config.mm_hidden_size * 4, config.hidden_size),
+                nn.GELU(),
+                nn.Linear(config.hidden_size, config.hidden_size),
+            )
+            self.downsample_rate = 2
+        elif mm_projector_type == "mlp_downsample_3x3_fix":
+            self.layers = nn.Sequential(
+                DownSample3x3BlockFix(),
+                nn.LayerNorm(config.mm_hidden_size * 9),
+                nn.Linear(config.mm_hidden_size * 9, config.mm_hidden_size * 3),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size * 3),
+                nn.Linear(config.mm_hidden_size * 3, config.hidden_size),
+                nn.GELU(),
+                nn.Linear(config.hidden_size, config.hidden_size),
+            )
+            self.downsample_rate = 3
+        elif mm_projector_type == "mlp_downsample_3x3_s2":
+            self.layers = nn.Sequential(
+                DownSample3x3BlockFix(),
+                nn.LayerNorm(config.mm_hidden_size * 9),
+                nn.Linear(config.mm_hidden_size * 9, config.mm_hidden_size * 3),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size * 3),
+                nn.Linear(config.mm_hidden_size * 3, config.mm_hidden_size),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size),
+                nn.Linear(config.mm_hidden_size, config.mm_hidden_size // 3),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size // 3),
+                nn.Linear(config.mm_hidden_size // 3, config.hidden_size),
+                nn.GELU(),
+                nn.Linear(config.hidden_size, config.hidden_size),
+            )
+        elif mm_projector_type == "mlp_downsample_3x3_s2_new":
+            self.layers = nn.Sequential(
+                DownSample3x3BlockFix(),
+                nn.LayerNorm(config.mm_hidden_size * 9),
+                nn.Linear(config.mm_hidden_size * 9, config.mm_hidden_size * 4),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size * 4),
+                nn.Linear(config.mm_hidden_size * 4, config.mm_hidden_size * 2),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size * 2),
+                nn.Linear(config.mm_hidden_size * 2, config.mm_hidden_size),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size),
+                nn.Linear(config.mm_hidden_size, config.mm_hidden_size // 3),
+                nn.GELU(),
+                nn.LayerNorm(config.mm_hidden_size // 3),
+                nn.Linear(config.mm_hidden_size // 3, config.hidden_size),
+                nn.GELU(),
+                nn.Linear(config.hidden_size, config.hidden_size),
+            )
+        else:
+            mlp_gelu_match = re.match(r"^mlp(\d+)x_gelu$", mm_projector_type)
+            if mlp_gelu_match:
+                mlp_depth = int(mlp_gelu_match.group(1))
+                modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
+                for _ in range(1, mlp_depth):
+                    modules.append(nn.GELU())
+                    modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+                self.layers = nn.Sequential(*modules)
+            else:
+                raise ValueError(f"Unknown projector type: {mm_projector_type}")
+
+    def forward(self, x, *args, **kwargs):
+        return self.layers(x)
+
+
+# AutoConfig.register("v2l_projector", MultimodalProjectorConfig)
+# AutoModel.register(MultimodalProjectorConfig, MultimodalProjector)

builder.py

@@ -0,0 +1,245 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import os
+import os.path as osp
+import warnings
+from dataclasses import asdict
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import torch
+import transformers
+from huggingface_hub import file_exists, repo_exists
+from huggingface_hub.utils import HFValidationError
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+# from .conversation import *
+from .conversation import SeparatorStyle, default_conversation
+
+SENTINEL_TOKEN = "<vila/sentinel>"
+MEDIA_TOKENS = {
+    "image": "<image>",
+    "video": "<vila/video>",
+}
+
+# from llava.model.utils import packing
+# from llava.utils.logging import logger
+# from llava.utils.tokenizer import infer_stop_tokens
+
+DUMMY_CONVERSATION = [
+    {"from": "human", "value": "question"},
+    {"from": "gpt", "value": "answer"},
+] * 10
+
+
+def tokenizer_image_token(prompt, tokenizer, return_tensors=None):
+    return tokenizer(prompt, return_tensors=return_tensors).input_ids[0]
+
+
+def has_tokenizer(repo_id_or_path: str) -> bool:
+    # Check if the tokenizer is in a local directory
+    if osp.exists(osp.join(repo_id_or_path, "tokenizer_config.json")):
+        return True
+
+    # Check if the tokenizer is in a Hugging Face Hub repo
+    try:
+        return repo_exists(repo_id_or_path) and file_exists(repo_id_or_path, "tokenizer_config.json")
+    except HFValidationError:
+        return False
+
+
+def _maybe_add_sentinel_token(tokenizer: transformers.PreTrainedTokenizer) -> None:
+    if not hasattr(tokenizer, "sentinel_token"):
+        tokenizer.add_tokens([SENTINEL_TOKEN], special_tokens=True)
+        tokenizer.sentinel_token = SENTINEL_TOKEN
+        tokenizer.sentinel_token_id = tokenizer.convert_tokens_to_ids(SENTINEL_TOKEN)
+
+
+def tokenize_conversation_legacy(
+    messages: Sequence[Dict[str, str]],
+    tokenizer: transformers.PreTrainedTokenizer,
+    add_generation_prompt: bool = False,
+    overrides: Optional[Dict[str, str]] = None,
+    no_system_prompt: bool = False,
+) -> torch.Tensor:
+    conv = default_conversation.copy()
+    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
+
+    if no_system_prompt:
+        conv.system = ""
+
+    # Skip the first message if it is not from human
+    if messages[0]["from"] != "human":
+        messages = messages[1:]
+
+    # Add a generation prompt if needed
+    if add_generation_prompt:
+        messages.append({"from": "gpt", "value": None})
+
+    conv.messages = []
+    for turn, message in enumerate(messages):
+        role = roles[message["from"]]
+        assert role == conv.roles[turn % 2]
+        if overrides is not None and message["from"] in overrides:
+            conv.append_message(role, overrides[message["from"]])
+        else:
+            conv.append_message(role, message["value"])
+
+    return tokenizer_image_token(conv.get_prompt(), tokenizer, return_tensors="pt")
+
+
+def tokenize_conversation(
+    messages: Sequence[Dict[str, str]],
+    tokenizer: transformers.PreTrainedTokenizer,
+    add_generation_prompt: bool = False,
+    overrides: Optional[Dict[str, str]] = None,
+    no_system_prompt: bool = False,
+) -> torch.Tensor:
+    # Normalize the conversation before tokenization
+    for message in messages:
+        message["value"] = message["value"].strip()
+
+    if default_conversation.sep_style != SeparatorStyle.AUTO:
+        return tokenize_conversation_legacy(
+            messages,
+            tokenizer,
+            add_generation_prompt=add_generation_prompt,
+            overrides=overrides,
+            no_system_prompt=no_system_prompt,
+        )
+
+    conversation = []
+    for m in messages:
+        message = {}
+        if m["from"] == "human":
+            message["role"] = "user"
+        elif m["from"] == "gpt":
+            message["role"] = "assistant"
+        else:
+            raise ValueError(f"Unexpected sender '{m['from']}' in conversation entry.")
+
+        message["content"] = m["value"]
+        if overrides is not None and m["from"] in overrides:
+            message["content"] = overrides[m["from"]]
+        conversation.append(message)
+
+    if no_system_prompt:
+        conversation = [{"role": "system", "content": ""}] + conversation
+
+    text = tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=add_generation_prompt,
+        tokenize=False,
+    )
+    return tokenizer_image_token(text, tokenizer, return_tensors="pt")
+
+
+def infer_stop_tokens(tokenizer: transformers.PreTrainedTokenizer) -> List[str]:
+    _maybe_add_sentinel_token(tokenizer)
+    template = tokenize_conversation(DUMMY_CONVERSATION, tokenizer, overrides={"gpt": SENTINEL_TOKEN})
+
+    stop_tokens = {tokenizer.eos_token}
+    for k in range(template.size(0) - 1):
+        if template[k] == tokenizer.sentinel_token_id:
+            stop_token = tokenizer.decode(template[k + 1])
+            stop_tokens.add(stop_token)
+    return list(stop_tokens)
+
+
+def context_length_extension(config):
+    orig_ctx_len = getattr(config, "max_position_embeddings", None)
+    model_max_length = getattr(config, "model_max_length", None)
+    if orig_ctx_len and model_max_length > orig_ctx_len:
+        print(f"Scaling RoPE from {orig_ctx_len} to {model_max_length}")
+        scaling_factor = float(math.ceil(model_max_length / orig_ctx_len))
+        config.rope_scaling = {"type": "linear", "factor": scaling_factor}
+    return config
+
+
+def build_llm_and_tokenizer(
+    model_name_or_path: str,
+    config: PretrainedConfig,
+    attn_implementation=None,
+    model_max_length=None,
+    *args,
+    **kwargs,
+) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+    # print(model_name_or_path)
+    llm_cfg = AutoConfig.from_pretrained(model_name_or_path)
+    llm_cfg._attn_implementation = attn_implementation
+    llm_cfg.model_max_length = model_max_length
+    if model_max_length is not None:
+        context_length_extension(llm_cfg)
+
+    # Quantization related
+    quantization_restore_from_checkpoint = False
+
+    if quantization_restore_from_checkpoint:
+        fp8_model_name_or_path = kwargs.pop("fp8_llm_cfg", None)
+
+        llm = AutoModelForCausalLM.from_pretrained(
+            fp8_model_name_or_path, config=llm_cfg, torch_dtype=eval(config.model_dtype), *args, **kwargs
+        )
+    else:
+        llm = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path, config=llm_cfg, torch_dtype=eval(config.model_dtype), *args, **kwargs
+        )
+    # NOTE(ligeng): not sure whether it affects the training
+    # packing.patch(llm)
+
+    # Locate the tokenizer.
+    llm_path = model_name_or_path
+    if not has_tokenizer(llm_path):
+        llm_path = osp.join(llm_path, "llm")
+    if not has_tokenizer(llm_path):
+        raise ValueError(f"Cannot find tokenizer in {llm_path}.")
+
+    tokenizer = AutoTokenizer.from_pretrained(llm_path, padding_side="right", use_fast=True, legacy=False)
+    if model_max_length is not None:
+        tokenizer.model_max_length = model_max_length
+
+    # Load chat template if specified.
+    if getattr(config, "chat_template", None) is not None:
+        print(f"Using chat template: {config.chat_template}")
+        fpath = os.path.join(os.path.dirname(__file__), "chat_templates", f"{config.chat_template}.jinja")
+        if not os.path.exists(fpath):
+            fpath = os.path.join(os.path.dirname(model_name_or_path), f"{config.chat_template}.jinja")
+        with open(fpath) as fd:
+            chat_template = fd.read()
+        tokenizer.chat_template = chat_template.replace("    ", "").replace("\n", "")
+
+    # Set stop tokens for the tokenizer
+    tokenizer.stop_tokens = infer_stop_tokens(tokenizer)
+    tokenizer.stop_token_ids = tokenizer.convert_tokens_to_ids(tokenizer.stop_tokens)
+
+    # Add media tokens to the tokenizer
+    tokenizer.media_tokens = MEDIA_TOKENS
+    tokenizer.media_token_ids = {}
+    for name, token in MEDIA_TOKENS.items():
+        tokenizer.add_tokens([token], special_tokens=True)
+        tokenizer.media_token_ids[name] = tokenizer.convert_tokens_to_ids(token)
+
+    # TODO(ligeng): is this necessary for llava?
+    config.hidden_size = llm.config.hidden_size
+    return llm, tokenizer

config.json

@@ -0,0 +1,318 @@
+{
+  "Ubit": 100,
+  "_name_or_path": "runs/train/NVILA-Lite_32b_qwen25vl/sft_32b_siglip2/model",
+  "architectures": [
+    "VILAForCasualLM"
+  ],
+  "babit": "E5M2",
+  "bobit": "E5M2",
+  "bwbit": "E5M2",
+  "chat_template": null,
+  "col_blocksize": -1,
+  "col_blocksize_optimizer": 128,
+  "draw_distribution_backward": false,
+  "draw_distribution_forward": false,
+  "drop_path_rate": 0.0,
+  "dynamic_s2": false,
+  "epsilon": 1e-10,
+  "epsilon_optimizer": 1e-15,
+  "fabit": "E4M3",
+  "first_order_bit": null,
+  "first_order_quant_type": null,
+  "fobit": "E4M3",
+  "fps": 0.0,
+  "fwbit": "E4M3",
+  "group_size": -1,
+  "hidden_size": 5120,
+  "image_aspect_ratio": "dynamic",
+  "image_encoder": {
+    "_target_": "llava.model.encoders.BasicImageEncoder"
+  },
+  "interpolate_mode": "linear",
+  "llm_cfg": {
+    "_name_or_path": "runs/train/NVILA-Lite_32b_qwen25vl/sft_32b_siglip2/model/llm",
+    "add_cross_attention": false,
+    "architectures": [
+      "Qwen2ForCausalLM"
+    ],
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 151643,
+    "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,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 151645,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "silu",
+    "hidden_size": 5120,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 27648,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 32768,
+    "max_window_layers": 70,
+    "min_length": 0,
+    "model_max_length": 4096,
+    "model_type": "qwen2",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 40,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 64,
+    "num_key_value_heads": 8,
+    "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,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000.0,
+    "sep_token_id": null,
+    "sliding_window": 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": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 151670
+  },
+  "max_tiles": 12,
+  "min_blockunit_col": 4,
+  "min_blockunit_row": 4,
+  "min_tiles": 1,
+  "mlp_path": null,
+  "mm_hidden_size": 1152,
+  "mm_projector": "mlp_downsample_2x2_fix",
+  "mm_projector_cfg": {
+    "_name_or_path": "runs/train/NVILA-Lite_32b_qwen25vl/sft_32b_siglip2/model/mm_projector",
+    "add_cross_attention": false,
+    "architectures": [
+      "MultimodalProjector"
+    ],
+    "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,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "mm_projector_type": "mlp_downsample_2x2_fix",
+    "model_type": "v2l_projector",
+    "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,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "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": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "mm_projector_lr": null,
+  "mm_use_im_patch_token": false,
+  "mm_use_im_start_end": false,
+  "mm_vision_select_feature": "cls_patch",
+  "mm_vision_select_layer": -2,
+  "model_dtype": "torch.bfloat16",
+  "model_name_or_path": "runs/train/NVILA-Lite_32b_qwen25vl/pretrain_32b_siglip2/model",
+  "model_type": "vila",
+  "num_time_tokens": 0,
+  "num_video_frames": 8,
+  "pad_block": false,
+  "pad_to_multiple_of": 0,
+  "qchoice": "none",
+  "quantize_model": false,
+  "refine_attn_blocksize": false,
+  "refine_col_blocksize": 4,
+  "refine_ln_blocksize": false,
+  "refine_ln_blocksize_but_only_backward": false,
+  "refine_ln_blocksize_but_only_forward": false,
+  "refine_ln_pertoken": false,
+  "refine_mlp_blocksize": false,
+  "refine_residual_fp": false,
+  "refine_row_blocksize": 4,
+  "resume_path": "runs/train/NVILA-Lite_32b_qwen25vl/sft_32b_siglip2/model",
+  "row_blocksize": -1,
+  "row_blocksize_optimizer": 1,
+  "s2": false,
+  "s2_max_split_size": 336,
+  "s2_resize_output_to_scale_idx": 0,
+  "s2_scales": "336,672,1008",
+  "second_order_bit": null,
+  "second_order_quant_type": null,
+  "soft_ce_std": 1.0,
+  "symm": true,
+  "time_token_format": "<t{t}>",
+  "time_token_ids": [],
+  "transformers_version": "4.45.0",
+  "tune_language_model": true,
+  "tune_mm_projector": true,
+  "tune_vision_tower": true,
+  "use_quantize_optimizer": false,
+  "version": "2.0",
+  "video_encoder": {
+    "_target_": "llava.model.encoders.BasicVideoEncoder"
+  },
+  "vision_resolution": -1,
+  "vision_tower": "/data/models/google/siglip2-so400m-patch16-512",
+  "vision_tower_cfg": {
+    "_name_or_path": "runs/train/NVILA-Lite_32b_qwen25vl/sft_32b_siglip2/model/vision_tower",
+    "add_cross_attention": false,
+    "architectures": [
+      "SiglipVisionModel"
+    ],
+    "attention_dropout": 0.0,
+    "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,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 512,
+    "intermediate_size": 4304,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "siglip_vision_model",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 27,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 16,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "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": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "vision_tower_lr": null,
+  "weight_memory_efficient": true,
+  "auto_map": {
+    "AutoProcessor": "auto_processor.VILAProcessor",
+    "AutoConfig": "modeling_vila.VILAConfig",
+    "AutoModel": "modeling_vila.VILAForCasualLM",
+    "AutoModelForCausalLM": "modeling_vila.VILAForCasualLM"
+  }
+}

configuration_vila.py

@@ -0,0 +1,93 @@
+import json
+import math
+import os
+import os.path as osp
+from copy import deepcopy
+from threading import Thread
+from typing import List, Optional
+
+import torch
+import torchvision
+from PIL import Image
+from transformers import (
+    AutoProcessor,
+    PretrainedConfig,
+    PreTrainedModel,
+    Qwen2Config,
+    Qwen2ForCausalLM,
+    Qwen2PreTrainedModel,
+    TextIteratorStreamer,
+)
+
+
+class VILAConfig(PretrainedConfig):
+    model_type = "vila"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        llm_cfg=None,
+        vision_tower_cfg=None,
+        mm_projector_cfg=None,
+        architectures=None,
+        resume_path=None,
+        hidden_size=None,
+        mm_hidden_size=None,
+        image_aspect_ratio=None,
+        num_video_frames=None,
+        fps=None,
+        mm_vision_select_layer=None,
+        mm_vision_select_feature=None,
+        mm_use_im_start_end=False,
+        mm_use_im_patch_token=False,
+        mm_projector_lr=None,
+        vision_tower_lr=None,
+        vision_resolution=None,
+        interpolate_mode=None,
+        s2=None,
+        dynamic_s2=None,
+        s2_scales=None,
+        s2_max_split_size=None,
+        s2_resize_output_to_scale_idx=0,
+        min_tiles: Optional[int] = 1,
+        max_tiles: Optional[int] = 12,
+        num_time_tokens=None,
+        time_token_format=None,
+        image_encoder: str = '{"_target_": "llava.model.encoders.BasicImageEncoder"}',
+        video_encoder: str = '{"_target_": "llava.model.encoders.BasicVideoEncoder"}',
+        **kwargs,
+    ):
+        super().__init__()
+        self.architectures = architectures
+        self.llm_cfg = llm_cfg
+        self.vision_tower_cfg = vision_tower_cfg
+        self.mm_projector_cfg = mm_projector_cfg
+        self.resume_path = resume_path
+
+        self.hidden_size = hidden_size
+        self.mm_hidden_size = mm_hidden_size
+        self.image_aspect_ratio = image_aspect_ratio
+        self.num_video_frames = num_video_frames
+        self.fps = fps
+        self.mm_vision_select_layer = mm_vision_select_layer
+        self.mm_vision_select_feature = mm_vision_select_feature
+        self.mm_use_im_start_end = mm_use_im_start_end
+        self.mm_use_im_patch_token = mm_use_im_patch_token
+        self.mm_projector_lr = mm_projector_lr
+        self.vision_tower_lr = vision_tower_lr
+        self.vision_resolution = vision_resolution
+        self.interpolate_mode = interpolate_mode
+        self.s2 = s2
+        self.dynamic_s2 = dynamic_s2
+        self.s2_scales = s2_scales
+        self.s2_max_split_size = s2_max_split_size
+        self.s2_resize_output_to_scale_idx = s2_resize_output_to_scale_idx
+        self.min_tiles = min_tiles
+        self.max_tiles = max_tiles
+        self.num_time_tokens = num_time_tokens
+        self.time_token_format = time_token_format
+
+        self.image_encoder = image_encoder
+        self.video_encoder = video_encoder
+
+        super().__init__(**kwargs)

constants.py

@@ -0,0 +1,43 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+DEFAULT_IMAGE_TOKEN = "<image>"
+
+SENTINEL_TOKEN = "<vila/sentinel>"
+MEDIA_TOKENS = {
+    "image": "<image>",
+    "video": "<vila/video>",
+}
+# <image> <vila/video> <vila/sentinel>
+# TODO(ligeng): need to discuss with Zhijian for the following tokens for different models.
+"""
+151643: AddedToken("<|endoftext|>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151644: AddedToken("<|im_start|>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151645: AddedToken("<|im_end|>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151646: AddedToken("[BOS]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151647: AddedToken("[PAD]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151648: AddedToken("<vila/sentinel>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151649: AddedToken("<image>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+151650: AddedToken("<vila/video>", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
+"""
+NUM_EXTRA_TOKENS = 8

conversation.py

@@ -0,0 +1,191 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+# This file is modified from https://github.com/haotian-liu/LLaVA/
+
+import dataclasses
+from enum import Enum, auto
+from typing import List
+
+# from llava.utils.logging import logger
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+
+    AUTO = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_3 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    sep_style: SeparatorStyle = SeparatorStyle.AUTO
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.LLAMA_3:
+            ret = self.system + self.sep
+            for rid, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message = message[0]
+                    sep = self.sep if rid < len(messages) - 1 else self.sep2
+                    ret += role + message + sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version,
+        )
+
+
+conv_auto = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(),
+    sep_style=SeparatorStyle.AUTO,
+    sep="\n",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(),
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+hermes_2 = Conversation(
+    system="<|im_start|>system\nAnswer the questions.",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+    messages=(),
+    version="hermes-2",
+)
+
+# Template added by Yukang. Note (kentang-mit@): sep is <|eot_id|> for official template.
+llama_3_chat = Conversation(
+    system="<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful language and vision assistant. "
+    "You are able to understand the visual content that the user provides, "
+    "and assist the user with a variety of tasks using natural language.",
+    roles=("<|start_header_id|>user<|end_header_id|>\n\n", "<|start_header_id|>assistant<|end_header_id|>\n\n"),
+    version="llama_v3",
+    messages=(),
+    sep_style=SeparatorStyle.LLAMA_3,
+    sep="<|eot_id|>",
+    sep2="<|end_of_text|>",
+)
+
+
+default_conversation = conv_auto
+conv_templates = {
+    "auto": conv_auto,
+    "hermes-2": hermes_2,
+    "llama_3": llama_3_chat,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "plain": conv_llava_plain,
+}
+
+
+CONVERSATION_MODE_MAPPING = {
+    "vila1.5-3b": "vicuna_v1",
+    "vila1.5-8b": "llama_3",
+    "vila1.5-13b": "vicuna_v1",
+    "vila1.5-40b": "hermes-2",
+    "llama-3": "llama_3",
+    "llama3": "llama_3",
+}
+
+
+def auto_set_conversation_mode(model_name_or_path: str) -> str:
+    global default_conversation
+    for k, v in CONVERSATION_MODE_MAPPING.items():
+        if k in model_name_or_path.lower():
+            print(f"Setting conversation mode to `{v}` based on model name/path `{model_name_or_path}`.")
+            default_conversation = conv_templates[v]
+            return

distributed.py

@@ -0,0 +1,73 @@
+import os
+import warnings
+from typing import Any, List, Optional
+
+from torch import distributed as dist
+
+__all__ = [
+    "init",
+    "is_initialized",
+    "size",
+    "rank",
+    "local_size",
+    "local_rank",
+    "is_main",
+    "barrier",
+    "gather",
+    "all_gather",
+]
+
+
+def init() -> None:
+    if "RANK" not in os.environ:
+        warnings.warn("Environment variable `RANK` is not set. Skipping distributed initialization.")
+        return
+    dist.init_process_group(backend="nccl", init_method="env://")
+
+
+def is_initialized() -> bool:
+    return dist.is_initialized()
+
+
+def size() -> int:
+    return int(os.environ.get("WORLD_SIZE", 1))
+
+
+def rank() -> int:
+    return int(os.environ.get("RANK", 0))
+
+
+def local_size() -> int:
+    return int(os.environ.get("LOCAL_WORLD_SIZE", 1))
+
+
+def local_rank() -> int:
+    return int(os.environ.get("LOCAL_RANK", 0))
+
+
+def is_main() -> bool:
+    return rank() == 0
+
+
+def barrier() -> None:
+    dist.barrier()
+
+
+def gather(obj: Any, dst: int = 0) -> Optional[List[Any]]:
+    if not is_initialized():
+        return [obj]
+    if is_main():
+        objs = [None for _ in range(size())]
+        dist.gather_object(obj, objs, dst=dst)
+        return objs
+    else:
+        dist.gather_object(obj, dst=dst)
+        return None
+
+
+def all_gather(obj: Any) -> List[Any]:
+    if not is_initialized():
+        return [obj]
+    objs = [None for _ in range(size())]
+    dist.all_gather_object(objs, obj)
+    return objs

evaluation.jsonl

@@ -0,0 +1 @@
+{"heron-bench": {"overall_score": 74.79865133099491, "details": {"conv": 8.0, "conv_rel": 80.76923076923076, "detail": 6.380952380952381, "detail_rel": 65.3658536585366, "complex": 7.65, "complex_rel": 78.26086956521739, "parse_error_count": 0, "overall": 7.533980582524272, "overall_rel": 74.79865133099491}}}

llm/added_tokens.json

@@ -0,0 +1,29 @@
+{
+  "</tool_call>": 151658,
+  "<image>": 151666,
+  "<tool_call>": 151657,
+  "<vila/sentinel>": 151665,
+  "<vila/video>": 151667,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652,
+  "[BOS]": 151668,
+  "[PAD]": 151669
+}

llm/config.json

@@ -0,0 +1,30 @@
+{
+  "_name_or_path": "runs/train/NVILA-Lite_32b_qwen25vl/sft_32b_siglip2/model/llm",
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 27648,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 70,
+  "model_max_length": 4096,
+  "model_type": "qwen2",
+  "num_attention_heads": 40,
+  "num_hidden_layers": 64,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.45.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151670
+}

llm/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.05,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "4.45.0"
+}

llm/model-00001-of-00014.safetensors

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

llm/model-00002-of-00014.safetensors

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

llm/model-00003-of-00014.safetensors

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

llm/model-00004-of-00014.safetensors

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

llm/model-00005-of-00014.safetensors

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

llm/model-00006-of-00014.safetensors

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

llm/model-00007-of-00014.safetensors

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

llm/model-00008-of-00014.safetensors

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

llm/model-00009-of-00014.safetensors

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

llm/model-00010-of-00014.safetensors

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

llm/model-00011-of-00014.safetensors

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

llm/model-00012-of-00014.safetensors

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

llm/model-00013-of-00014.safetensors

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

llm/model-00014-of-00014.safetensors

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

llm/model.safetensors.index.json

@@ -0,0 +1,778 @@
+{
+  "metadata": {
+    "total_size": 65519683584
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00014-of-00014.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.self_attn.k_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.self_attn.q_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.0.self_attn.v_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.self_attn.k_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.self_attn.q_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.1.self_attn.v_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.self_attn.k_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.self_attn.q_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.10.self_attn.v_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.self_attn.k_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.self_attn.q_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.11.self_attn.v_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.self_attn.k_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.self_attn.q_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.12.self_attn.v_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.13.self_attn.k_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.13.self_attn.q_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.13.self_attn.v_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.self_attn.k_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.self_attn.q_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.14.self_attn.v_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.self_attn.k_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.self_attn.q_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.15.self_attn.v_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.self_attn.k_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.self_attn.q_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.16.self_attn.v_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.self_attn.k_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.self_attn.q_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.17.self_attn.v_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.18.self_attn.k_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.18.self_attn.q_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.18.self_attn.v_proj.bias": "model-00004-of-00014.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00004-of-00014.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.self_attn.k_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.self_attn.q_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.19.self_attn.v_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.self_attn.k_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.self_attn.q_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.2.self_attn.v_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.self_attn.k_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.self_attn.q_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.20.self_attn.v_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.self_attn.k_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.self_attn.q_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.21.self_attn.v_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.self_attn.k_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.self_attn.q_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.22.self_attn.v_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.23.self_attn.k_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.23.self_attn.q_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.23.self_attn.v_proj.bias": "model-00005-of-00014.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00005-of-00014.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.self_attn.k_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.self_attn.q_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.24.self_attn.v_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.self_attn.k_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.self_attn.q_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.25.self_attn.v_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.self_attn.k_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.self_attn.q_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.26.self_attn.v_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.self_attn.k_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.self_attn.q_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.27.self_attn.v_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.28.self_attn.k_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.28.self_attn.q_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.28.self_attn.v_proj.bias": "model-00006-of-00014.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00006-of-00014.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.self_attn.k_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.self_attn.q_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.29.self_attn.v_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.3.self_attn.k_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.3.self_attn.q_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.3.self_attn.v_proj.bias": "model-00001-of-00014.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00014.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.self_attn.k_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.self_attn.q_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.30.self_attn.v_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.self_attn.k_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.self_attn.q_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.31.self_attn.v_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.mlp.down_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.mlp.up_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.self_attn.k_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.32.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.self_attn.q_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.32.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.32.self_attn.v_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.32.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.33.mlp.down_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.33.mlp.gate_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.33.mlp.up_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.33.self_attn.k_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.33.self_attn.k_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.33.self_attn.o_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.33.self_attn.q_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.33.self_attn.q_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.33.self_attn.v_proj.bias": "model-00007-of-00014.safetensors",
+    "model.layers.33.self_attn.v_proj.weight": "model-00007-of-00014.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.mlp.down_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.mlp.up_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.self_attn.k_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.34.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.self_attn.q_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.34.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.34.self_attn.v_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.34.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.input_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.mlp.down_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.mlp.up_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.self_attn.k_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.35.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.self_attn.q_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.35.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.35.self_attn.v_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.35.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.input_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.mlp.down_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.mlp.up_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.self_attn.k_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.36.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.self_attn.q_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.36.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.36.self_attn.v_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.36.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.input_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.mlp.down_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.mlp.up_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.post_attention_layernorm.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.self_attn.k_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.37.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.self_attn.q_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.37.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.37.self_attn.v_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.37.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.38.input_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.38.mlp.down_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.38.mlp.gate_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.38.mlp.up_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.38.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.38.self_attn.k_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.38.self_attn.k_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.38.self_attn.o_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.38.self_attn.q_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.38.self_attn.q_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.38.self_attn.v_proj.bias": "model-00008-of-00014.safetensors",
+    "model.layers.38.self_attn.v_proj.weight": "model-00008-of-00014.safetensors",
+    "model.layers.39.input_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.mlp.down_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.mlp.up_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.self_attn.k_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.39.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.self_attn.q_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.39.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.39.self_attn.v_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.39.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.self_attn.k_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.self_attn.q_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.4.self_attn.v_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.40.input_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.mlp.down_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.mlp.up_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.self_attn.k_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.40.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.self_attn.q_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.40.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.40.self_attn.v_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.40.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.input_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.mlp.down_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.mlp.up_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.self_attn.k_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.41.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.self_attn.q_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.41.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.41.self_attn.v_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.41.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.input_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.mlp.down_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.mlp.up_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.post_attention_layernorm.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.self_attn.k_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.42.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.self_attn.q_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.42.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.42.self_attn.v_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.42.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.43.input_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.43.mlp.down_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.43.mlp.gate_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.43.mlp.up_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.43.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.43.self_attn.k_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.43.self_attn.k_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.43.self_attn.o_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.43.self_attn.q_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.43.self_attn.q_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.43.self_attn.v_proj.bias": "model-00009-of-00014.safetensors",
+    "model.layers.43.self_attn.v_proj.weight": "model-00009-of-00014.safetensors",
+    "model.layers.44.input_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.mlp.down_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.mlp.up_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.self_attn.k_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.44.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.self_attn.q_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.44.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.44.self_attn.v_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.44.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.input_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.mlp.down_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.mlp.up_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.self_attn.k_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.45.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.self_attn.q_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.45.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.45.self_attn.v_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.45.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.input_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.mlp.down_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.mlp.up_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.self_attn.k_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.46.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.self_attn.q_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.46.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.46.self_attn.v_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.46.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.input_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.mlp.down_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.mlp.up_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.post_attention_layernorm.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.self_attn.k_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.47.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.self_attn.q_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.47.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.47.self_attn.v_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.47.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.48.input_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.48.mlp.down_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.48.mlp.gate_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.48.mlp.up_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.48.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.48.self_attn.k_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.48.self_attn.k_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.48.self_attn.o_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.48.self_attn.q_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.48.self_attn.q_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.48.self_attn.v_proj.bias": "model-00010-of-00014.safetensors",
+    "model.layers.48.self_attn.v_proj.weight": "model-00010-of-00014.safetensors",
+    "model.layers.49.input_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.mlp.down_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.mlp.up_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.self_attn.k_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.49.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.self_attn.q_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.49.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.49.self_attn.v_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.49.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.self_attn.k_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.self_attn.q_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.5.self_attn.v_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.50.input_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.mlp.down_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.mlp.up_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.self_attn.k_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.50.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.self_attn.q_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.50.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.50.self_attn.v_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.50.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.input_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.mlp.down_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.mlp.up_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.self_attn.k_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.51.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.self_attn.q_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.51.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.51.self_attn.v_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.51.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.input_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.mlp.down_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.mlp.up_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.post_attention_layernorm.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.self_attn.k_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.52.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.self_attn.q_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.52.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.52.self_attn.v_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.52.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.53.input_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.53.mlp.down_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.53.mlp.gate_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.53.mlp.up_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.53.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.53.self_attn.k_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.53.self_attn.k_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.53.self_attn.o_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.53.self_attn.q_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.53.self_attn.q_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.53.self_attn.v_proj.bias": "model-00011-of-00014.safetensors",
+    "model.layers.53.self_attn.v_proj.weight": "model-00011-of-00014.safetensors",
+    "model.layers.54.input_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.mlp.down_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.mlp.up_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.self_attn.k_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.54.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.self_attn.q_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.54.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.54.self_attn.v_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.54.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.input_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.mlp.down_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.mlp.up_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.self_attn.k_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.55.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.self_attn.q_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.55.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.55.self_attn.v_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.55.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.input_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.mlp.down_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.mlp.up_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.self_attn.k_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.56.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.self_attn.q_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.56.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.56.self_attn.v_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.56.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.input_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.mlp.down_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.mlp.up_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.post_attention_layernorm.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.self_attn.k_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.57.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.self_attn.q_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.57.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.57.self_attn.v_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.57.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.58.input_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.58.mlp.down_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.58.mlp.gate_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.58.mlp.up_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.58.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.58.self_attn.k_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.58.self_attn.k_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.58.self_attn.o_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.58.self_attn.q_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.58.self_attn.q_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.58.self_attn.v_proj.bias": "model-00012-of-00014.safetensors",
+    "model.layers.58.self_attn.v_proj.weight": "model-00012-of-00014.safetensors",
+    "model.layers.59.input_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.mlp.down_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.mlp.up_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.self_attn.k_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.59.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.self_attn.q_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.59.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.59.self_attn.v_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.59.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.self_attn.k_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.self_attn.q_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.6.self_attn.v_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.60.input_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.mlp.down_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.mlp.up_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.self_attn.k_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.60.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.self_attn.q_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.60.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.60.self_attn.v_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.60.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.input_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.mlp.down_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.mlp.up_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.self_attn.k_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.61.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.self_attn.q_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.61.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.61.self_attn.v_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.61.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.input_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.mlp.down_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.mlp.up_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.post_attention_layernorm.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.self_attn.k_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.62.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.self_attn.q_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.62.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.62.self_attn.v_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.62.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.63.input_layernorm.weight": "model-00014-of-00014.safetensors",
+    "model.layers.63.mlp.down_proj.weight": "model-00014-of-00014.safetensors",
+    "model.layers.63.mlp.gate_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.63.mlp.up_proj.weight": "model-00014-of-00014.safetensors",
+    "model.layers.63.post_attention_layernorm.weight": "model-00014-of-00014.safetensors",
+    "model.layers.63.self_attn.k_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.63.self_attn.k_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.63.self_attn.o_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.63.self_attn.q_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.63.self_attn.q_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.63.self_attn.v_proj.bias": "model-00013-of-00014.safetensors",
+    "model.layers.63.self_attn.v_proj.weight": "model-00013-of-00014.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.self_attn.k_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.self_attn.q_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.7.self_attn.v_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.8.self_attn.k_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.8.self_attn.q_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.8.self_attn.v_proj.bias": "model-00002-of-00014.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00002-of-00014.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.self_attn.k_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.self_attn.q_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00003-of-00014.safetensors",
+    "model.layers.9.self_attn.v_proj.bias": "model-00003-of-00014.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00003-of-00014.safetensors",
+    "model.norm.weight": "model-00014-of-00014.safetensors"
+  }
+}

llm/special_tokens_map.json

@@ -0,0 +1,41 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<vila/sentinel>",
+    "<image>",
+    "<vila/video>"
+  ],
+  "bos_token": {
+    "content": "[BOS]",
+    "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
+  }
+}

llm/tokenizer.json

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

llm/tokenizer_config.json

@@ -0,0 +1,252 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<vila/sentinel>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151666": {
+      "content": "<image>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151667": {
+      "content": "<vila/video>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151668": {
+      "content": "[BOS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151669": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<vila/sentinel>",
+    "<image>",
+    "<vila/video>"
+  ],
+  "bos_token": "[BOS]",
+  "chat_template": "{% if messages[0]['role'] != 'system' %}{{ '<|im_start|>system\\n以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。<|im_end|>\\n' }}{% endif %}{% for message in messages if message['content'] is not none %}{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "legacy": false,
+  "model_max_length": 4096,
+  "pad_token": "[PAD]",
+  "padding_side": "right",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

loss.py

@@ -0,0 +1,48 @@
+from typing import List, Union
+
+import torch
+from torch.nn.functional import cross_entropy
+
+from .constants import IGNORE_INDEX
+
+__all__ = ["soft_cross_entropy"]
+
+
+def soft_cross_entropy(
+    outputs: torch.Tensor,
+    targets: torch.Tensor,
+    soft_tokens: Union[torch.Tensor, List[int]],
+    std: float = 1,
+    ignore_index: int = IGNORE_INDEX,
+) -> torch.Tensor:
+    # Remove last token from outputs and first token from targets
+    outputs = outputs[..., :-1, :].contiguous()
+    targets = targets[..., 1:].contiguous()
+
+    # Flatten outputs and targets
+    targets = targets.view(-1)
+    outputs = outputs.view(targets.size(0), -1)
+
+    # Remove outputs and targets with ignore_index
+    indices = targets != ignore_index
+    outputs = outputs[indices]
+    targets = targets[indices]
+
+    # Convert soft token IDs to tensor
+    if isinstance(soft_tokens, list):
+        soft_tokens = torch.tensor(soft_tokens).to(targets)
+
+    # Calculate loss for non-soft tokens
+    indices = torch.isin(targets, soft_tokens, invert=True)
+    loss = cross_entropy(outputs[indices], targets[indices], reduction="sum")
+
+    # Calculate loss for soft tokens
+    indices = torch.isin(targets, soft_tokens)
+    targets_indices = torch.zeros_like(outputs[indices])
+    for k, target in enumerate(targets[indices]):
+        dist = torch.exp(-((target - soft_tokens) ** 2) / (2 * std**2))
+        targets_indices[k][soft_tokens] = dist / dist.sum()
+    loss += cross_entropy(outputs[indices], targets_indices, reduction="sum")
+
+    # Return average loss
+    return loss / targets.size(0)

media.py

@@ -0,0 +1,129 @@
+import glob
+import os
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, Union
+
+import cv2
+import numpy as np
+import PIL
+import PIL.Image
+import requests
+from transformers import PretrainedConfig
+
+# from llava.constants import MEDIA_TOKENS
+# from llava.media import Image, Video
+# from llava.utils import make_list
+# from llava.utils.logging import logger
+
+MEDIA_TOKENS = {
+    "image": "<image>",
+    "video": "<vila/video>",
+}
+
+
+class Media:
+    pass
+
+
+class File(Media):
+    def __init__(self, path: str) -> None:
+        self.path = path
+
+
+class Image(File):
+    pass
+
+
+class Video(File):
+    pass
+
+
+def make_list(obj: Any) -> List:
+    return obj if isinstance(obj, list) else [obj]
+
+
+def _extract_image(image: Union[Image, PIL.Image.Image]) -> PIL.Image.Image:
+    if isinstance(image, Image):
+        if image.path.startswith("http://") or image.path.startswith("https://"):
+            image = PIL.Image.open(requests.get(image.path, stream=True).raw)
+        else:
+            image = PIL.Image.open(image.path)
+    return image
+
+
+def _load_video(video_path: str, *, num_frames: int) -> List[PIL.Image.Image]:
+    # Load video frames from a directory
+    if os.path.isdir(video_path):
+        frame_paths = sorted(glob.glob(os.path.join(video_path, "*")))
+        indices = np.round(np.linspace(0, len(frame_paths) - 1, num_frames)).astype(int)
+        return [PIL.Image.open(frame_paths[index]) for index in indices]
+
+    # Load video frames from a video file
+    vidcap = cv2.VideoCapture(video_path)
+
+    # Find the last frame as frame count might not be accurate
+    frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    while frame_count > 0:
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_count - 1)
+        if vidcap.grab():
+            break
+        frame_count -= 1
+    else:
+        raise ValueError(f"Video '{video_path}' has no frames.")
+
+    # Extract frames uniformly
+    indices = np.round(np.linspace(0, frame_count - 1, num_frames)).astype(int)
+    frames = {}
+    for index in indices:
+        if index in frames:
+            continue
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, index)
+        success, frame = vidcap.read()
+        if not success:
+            print(f"Failed to read frame {index} from video '{video_path}'. Skipped.")
+            continue
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames[index] = PIL.Image.fromarray(frame)
+    return [frames[index] for index in indices if index in frames]
+
+
+def _extract_video(video: Video, config: PretrainedConfig) -> List[PIL.Image.Image]:
+    num_frames = config.num_video_frames
+    if getattr(config, "fps") != 0:
+        print("Extracting frames from video with specified FPS is not supported yet. Ignored.")
+
+    frames = _load_video(video.path, num_frames=num_frames)
+    return frames
+
+
+def extract_media(
+    messages: List[Dict[str, Any]],
+    config: Optional[PretrainedConfig] = None,
+    draft: bool = False,
+) -> Dict[str, List[Any]]:
+    media = defaultdict(list)
+    for message in messages:
+        text = ""
+        for part in make_list(message["value"]):
+            if isinstance(part, str):
+                for token in MEDIA_TOKENS.values():
+                    if token in part:
+                        print(f"Media token '{token}' found in text: '{part}'. Removed.")
+                        part = part.replace(token, "").strip()
+                text += part
+            elif isinstance(part, (Image, PIL.Image.Image)):
+                if draft:
+                    media["image"].append(part)
+                else:
+                    media["image"].append(_extract_image(part))
+                text += MEDIA_TOKENS["image"]
+            elif isinstance(part, Video):
+                if draft:
+                    media["video"].append(part)
+                else:
+                    media["video"].append(_extract_video(part, config))
+                text += MEDIA_TOKENS["video"]
+            else:
+                raise ValueError(f"Unsupported prompt part type: {type(part)}")
+        message["value"] = text
+    return media

media_encoder.py

@@ -0,0 +1,101 @@
+from functools import partial
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch import nn
+
+
+class BaseEncoder(nn.Module):
+    def __init__(self, parent: nn.Module) -> None:
+        super().__init__()
+        self._parent = [parent]
+
+    @property
+    def parent(self) -> nn.Module:
+        return self._parent[0]
+
+
+class BasicImageEncoder(BaseEncoder):
+    def __init__(
+        self,
+        parent: torch.nn.Module,
+        start_tokens: Optional[str] = None,
+        end_tokens: Optional[str] = "\n",
+    ) -> None:
+        super().__init__(parent)
+        self.start_tokens = start_tokens
+        self.end_tokens = end_tokens
+
+    def embed_tokens(self, tokens: Optional[str]) -> Optional[torch.Tensor]:
+        if tokens is None:
+            return None
+        token_ids = self.parent.tokenizer(tokens).input_ids
+        token_ids = torch.tensor(token_ids, device=self.parent.device)
+        return self.parent.llm.model.embed_tokens(token_ids)
+
+    def _process_features(
+        self,
+        features: torch.Tensor,
+        start_token_embeds: Optional[torch.Tensor],
+        end_token_embeds: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if start_token_embeds is not None:
+            features = torch.cat([start_token_embeds, features], dim=0)
+        if end_token_embeds is not None:
+            features = torch.cat([features, end_token_embeds], dim=0)
+        return features
+
+    def forward(self, images: List[torch.Tensor], config: Dict[str, Any]) -> List[torch.Tensor]:
+        images = torch.stack(images, dim=0)
+        features = self.parent.encode_images(images, block_sizes=config.get("block_sizes"))
+        process_features = partial(
+            self._process_features,
+            start_token_embeds=self.embed_tokens(self.start_tokens),
+            end_token_embeds=self.embed_tokens(self.end_tokens),
+        )
+        return [process_features(f) for f in features]
+
+
+class BasicVideoEncoder(BaseEncoder):
+    def __init__(
+        self,
+        parent: torch.nn.Module,
+        start_tokens: Optional[str] = None,
+        end_tokens: Optional[str] = "\n",
+    ) -> None:
+        super().__init__(parent)
+        self.start_tokens = start_tokens
+        self.end_tokens = end_tokens
+
+    def embed_tokens(self, tokens: Optional[str]) -> Optional[torch.Tensor]:
+        if tokens is None:
+            return None
+        token_ids = self.parent.tokenizer(tokens).input_ids
+        token_ids = torch.tensor(token_ids, device=self.parent.device)
+        return self.parent.llm.model.embed_tokens(token_ids)
+
+    def _process_features(
+        self,
+        features: torch.Tensor,
+        start_token_embeds: Optional[torch.Tensor],
+        end_token_embeds: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        if start_token_embeds is not None:
+            start_embeds = torch.stack([start_token_embeds] * features.shape[0], dim=0)
+            features = torch.cat([start_embeds, features], dim=1)
+        if end_token_embeds is not None:
+            end_embeds = torch.stack([end_token_embeds] * features.shape[0], dim=0)
+            features = torch.cat([features, end_embeds], dim=1)
+        return features.flatten(0, 1)
+
+    def forward(self, videos: List[torch.Tensor], config: Dict[str, Any]) -> List[torch.Tensor]:
+        num_frames = [video.shape[0] for video in videos]
+        images = torch.cat(videos, dim=0)
+        features = self.parent.encode_images(images)
+        features = torch.split(features, num_frames)
+        process_features = partial(
+            self._process_features,
+            start_token_embeds=self.embed_tokens(self.start_tokens),
+            end_token_embeds=self.embed_tokens(self.end_tokens),
+        )
+        return [process_features(f) for f in features]

mm_projector/config.json

@@ -0,0 +1,10 @@
+{
+  "_name_or_path": "runs/train/NVILA-Lite_32b_qwen25vl/sft_32b_siglip2/model/mm_projector",
+  "architectures": [
+    "MultimodalProjector"
+  ],
+  "mm_projector_type": "mlp_downsample_2x2_fix",
+  "model_type": "v2l_projector",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.45.0"
+}

mm_projector/model.safetensors

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

mm_utils.py

@@ -0,0 +1,572 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# dynamic_preprocess and find_closest_aspect_ratio are referenced from https://github.com/OpenGVLab/InternVL
+
+import base64
+import os
+import tempfile
+from io import BytesIO
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import StoppingCriteria
+
+from .constants import DEFAULT_IMAGE_TOKEN
+
+
+def get_frame_from_vcap(vidcap, num_frames=10, max_fps=0.0, fps=None, frame_count=None, video_file_name=None):
+    import cv2
+
+    if fps == None or frame_count == None:
+        # if one of fps or frame_count is None, still recompute
+        fps = vidcap.get(cv2.CAP_PROP_FPS)
+        frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if fps == 0 or frame_count == 0:
+        print(f"Video file not found. return empty images. {video_file_name}")
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * num_frames, 0
+
+    duration = frame_count / fps
+    frame_interval = frame_count // num_frames
+    if frame_interval == 0 and frame_count <= 1:
+        print(f"frame_interval is equal to 0. return empty image. {video_file_name}")
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * num_frames, 0
+    # print("duration:", duration, "frames:", frame_count, "intervals:", frame_interval)
+
+    images = []
+    count = 0
+    success = True
+    frame_indices = np.linspace(0, frame_count - 1, num_frames, dtype=int)
+    while success:
+        # print("frame_count:", frame_count, "count:", count, "num_frames:", num_frames, "frame_interval:", frame_interval)
+        if frame_count >= num_frames:
+            success, frame = vidcap.read()
+            if count in frame_indices:
+                try:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                except BaseException:
+                    continue
+                if len(images) >= num_frames:
+                    return images, num_frames
+            count += 1
+        else:
+            # Left padding frames if the video is not long enough
+            success, frame = vidcap.read()
+            if success:
+                try:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                except BaseException:
+                    continue
+                count += 1
+            else:
+                break
+    if len(images) == 0:
+        raise ValueError("Did not find enough frames in the video. return empty image.")
+
+    return images, len(images)
+
+
+def get_frame_from_vcap_with_fps(vidcap, num_frames=10, max_fps=0.0, fps=None, frame_count=None, video_file_name=None):
+    """
+    num_frames is the max number of frames the model can support.
+    frame_count is the number of frames in the input video.
+    max_fps is the max FPS of the model can support.
+    fps is the fps of the input video.
+    """
+
+    import random
+
+    import cv2
+
+    if fps == None or frame_count == None:
+        # if one of fps or frame_count is None, still recompute
+        fps = vidcap.get(cv2.CAP_PROP_FPS)
+        frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    if fps == 0 or frame_count == 0:
+        print(f"Video file not found. return empty images. {video_file_name}")
+        empty_video_frames = int(random.uniform(2, 8 * max_fps))
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * empty_video_frames, 0
+
+    duration = frame_count / fps
+    # print("duration:", duration, "frames:", frame_count, "fps:", fps, "num_frames:", num_frames, "max_fps:", max_fps)
+    # If the video is too long (longer than max_fps and num_frames can support),
+    # we will use lower fps to sample frames.
+    if duration >= num_frames / max_fps:
+        frame_interval = frame_count // num_frames
+
+        # If the video is too short, we will skip the video if there is only one frame.
+        if frame_interval == 0 and frame_count <= 1:
+            print(f"frame_interval is equal to 0. return empty image. {video_file_name}")
+            empty_video_frames = int(random.uniform(2, 8 * max_fps))
+            return [
+                Image.new("RGB", (720, 720)),
+            ] * empty_video_frames, 0
+
+        images = []
+        count = 0
+        success = True
+        frame_indices = np.linspace(0, frame_count - 1, num_frames, dtype=int)
+
+        while success:
+            if frame_count >= num_frames:
+                # success, frame = vidcap.read()
+                if count in frame_indices:
+                    success, frame = vidcap.read()
+                    try:
+                        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                        im_pil = Image.fromarray(img)
+                        images.append(im_pil)
+                    except:
+                        # print("Failed to read frame:", count)
+                        continue
+                    if len(images) >= num_frames:
+                        return images, num_frames
+                else:
+                    success = vidcap.grab()
+                count += 1
+            else:
+                # Left padding frames if the video is not long enough
+                success, frame = vidcap.read()
+                if success:
+                    try:
+                        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                        im_pil = Image.fromarray(img)
+                        images.append(im_pil)
+                    except:
+                        # print("Failed to read frame:", count)
+                        continue
+                    count += 1
+                else:
+                    break
+    else:
+        frames_required = int(duration * max_fps)
+        frame_indices = np.linspace(0, frame_count - 1, frames_required, dtype=int)
+        if frames_required == 0:
+            print(f"frames_required is fewer than 2. Duration {duration}, return empty image.")
+            empty_video_frames = int(random.uniform(2, 8 * max_fps))
+            return [
+                Image.new("RGB", (720, 720)),
+            ] * empty_video_frames, 0
+        elif frames_required == 1:
+            frame_indices = np.linspace(0, frame_count - 1, 2, dtype=int)
+        images = []
+        count = 0
+        looked = 0
+        success = True
+
+        while success:
+            success, frame = vidcap.read()
+            if success and (looked in frame_indices):
+                try:
+                    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                    im_pil = Image.fromarray(img)
+                    images.append(im_pil)
+                except:
+                    continue
+                count += 1
+            looked += 1
+
+    if len(images) == 0:
+        empty_video_frames = int(random.uniform(2, 8 * max_fps))
+        return [
+            Image.new("RGB", (720, 720)),
+        ] * empty_video_frames, 0
+    else:
+        return images, len(images)
+
+
+def opencv_extract_frames(vpath_or_bytesio, frames=6, max_fps=0.0, fps=None, frame_count=None):
+    """
+    Extract frames from a video using OpenCV.
+
+    Args:
+        vpath_or_bytesio (str or BytesIO): Path to the video file or BytesIO object containing the video.
+        frames (int): Number of frames to extract from the video.
+        fps (float): Frames per second of the video. If 0.0, the function will extract frames at equal intervals.
+
+    Returns:
+        list: List of PIL Images extracted from the video.
+
+    Raises:
+        NotImplementedError: If the type of `vpath_or_bytesio` is not supported.
+    """
+    import cv2
+
+    if isinstance(vpath_or_bytesio, str):
+        vidcap = cv2.VideoCapture(vpath_or_bytesio)
+        if max_fps > 0.0:
+            return get_frame_from_vcap_with_fps(
+                vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=vpath_or_bytesio
+            )
+        return get_frame_from_vcap(
+            vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=vpath_or_bytesio
+        )
+    elif isinstance(vpath_or_bytesio, (BytesIO,)):
+        # assuming mp4
+        with tempfile.NamedTemporaryFile(delete=True, suffix=".mp4") as temp_video:
+            temp_video.write(vpath_or_bytesio.read())
+            temp_video_name = temp_video.name
+            vidcap = cv2.VideoCapture(temp_video_name)
+            if max_fps > 0.0:
+                return get_frame_from_vcap_with_fps(
+                    vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=temp_video_name
+                )
+            return get_frame_from_vcap(
+                vidcap, frames, max_fps, fps=fps, frame_count=frame_count, video_file_name=temp_video_name
+            )
+    else:
+        raise NotImplementedError(type(vpath_or_bytesio))
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+    """
+    Expand the given PIL image to a square shape by adding padding.
+
+    Parameters:
+    - pil_img: The PIL image to be expanded.
+    - background_color: The color of the padding to be added.
+
+    Returns:
+    - The expanded PIL image.
+
+    If the image is already square, it is returned as is.
+    If the image is wider than it is tall, padding is added to the top and bottom.
+    If the image is taller than it is wide, padding is added to the left and right.
+    """
+    width, height = pil_img.size
+    if pil_img.mode == "L":
+        background_color = background_color[0]
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float("inf")
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+
+def dynamic_preprocess(image, min_num=1, max_num=12, image_size=384, use_thumbnail=True):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    # calculate the existing image aspect ratio
+    target_ratios = {
+        (i, j)
+        for n in range(min_num, max_num + 1)
+        for i in range(1, n + 1)
+        for j in range(1, n + 1)
+        if i * j <= max_num and i * j >= min_num
+    }
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size,
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
+
+def dynamic_s2_preprocess(image, s2_scales=[384, 768, 1152], max_num=12, image_size=384):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+    min_num = (s2_scales[-1] // s2_scales[0]) ** 2  # at least use number of tiles as the largest scale
+
+    processed_images = []
+
+    ##########################################################################################
+    ############# Add tiles for all but the last scale using fixed squre ratio ###############
+    ##########################################################################################
+
+    for scale in s2_scales[:-1]:
+        target_width = image_size * (scale // s2_scales[0])
+        target_height = image_size * (scale // s2_scales[0])
+        blocks = (scale // s2_scales[0]) ** 2
+
+        # resize the image
+        resized_img = image.resize((target_width, target_height))
+        for i in range(blocks):
+            box = (
+                (i % (target_width // image_size)) * image_size,
+                (i // (target_width // image_size)) * image_size,
+                ((i % (target_width // image_size)) + 1) * image_size,
+                ((i // (target_width // image_size)) + 1) * image_size,
+            )
+            # split the image
+            split_img = resized_img.crop(box)
+            processed_images.append(split_img)
+
+    ##########################################################################################
+    ################ Add tiles for the last scale using dynamic aspect ratio #################
+    ##########################################################################################
+
+    # calculate the existing image aspect ratio
+    target_ratios = {
+        (i, j)
+        for n in range(min_num, max_num + 1)
+        for i in range(1, n + 1)
+        for j in range(1, n + 1)
+        if i * j <= max_num and i * j >= min_num
+    }
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size,
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+
+    return processed_images, (target_aspect_ratio[1], target_aspect_ratio[0])
+
+
+def dynamic_process_images_and_prompt(images, prompt, data_args, image_folder=None, max_tiles=None):
+    prompt = prompt.split(DEFAULT_IMAGE_TOKEN)
+    idx = 0
+    all_images = []
+    for img in images:
+        processed_images = process_image(img, data_args, image_folder, enable_dynamic_res=True, max_tiles=max_tiles)
+        all_images.append(processed_images)
+        prompt.insert(idx + 1, f"{DEFAULT_IMAGE_TOKEN}\n" * processed_images.shape[0])
+        idx += 2
+    prompt = "".join(prompt)
+    if all_images:
+        all_images = torch.cat(all_images)
+    else:
+        all_images = None
+        prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, "")
+    return all_images, prompt
+
+
+def dynamic_s2_process_images_and_prompt(images, prompt, data_args, image_folder=None):
+    idx = 0
+    all_images = []
+    all_block_size = []
+    for img in images:
+        processed_images, block_size = process_image(img, data_args, image_folder, enable_dynamic_s2=True)
+        all_images.append(processed_images)
+        all_block_size.append(block_size)
+        idx += 2
+    if all_images:
+        all_images = torch.cat(all_images)
+    else:
+        all_images = None
+    return all_images, all_block_size
+
+
+def process_image(
+    image_file, data_args, image_folder, enable_dynamic_res=False, enable_dynamic_s2=False, max_tiles=None
+):
+    processor = data_args.image_processor
+    if isinstance(image_file, str):
+        if image_folder is not None:
+            image = Image.open(os.path.join(image_folder, image_file)).convert("RGB")
+        else:
+            image = Image.open(image_file).convert("RGB")
+    else:
+        # image is stored in bytearray
+        image = image_file
+    image = image.convert("RGB")
+    if hasattr(data_args.image_processor, "crop_size"):
+        # CLIP vision tower
+        crop_size = data_args.image_processor.crop_size
+    else:
+        # SIGLIP vision tower
+        assert hasattr(data_args.image_processor, "size")
+        crop_size = data_args.image_processor.size
+    if "dynamic_s2" in data_args.image_aspect_ratio and enable_dynamic_s2:
+        assert crop_size["height"] == crop_size["width"]
+        images, block_size = dynamic_s2_preprocess(
+            image, s2_scales=data_args.s2_scales, max_num=data_args.max_tiles, image_size=crop_size["height"]
+        )
+        images = [processor.preprocess(image, return_tensors="pt")["pixel_values"][0] for image in images]
+        return torch.stack(images), block_size
+    if "dynamic" in data_args.image_aspect_ratio and enable_dynamic_res:
+        assert crop_size["height"] == crop_size["width"]
+        if max_tiles is not None:
+            max_num = max_tiles
+        else:
+            max_num = data_args.max_tiles
+        images = dynamic_preprocess(image, min_num=data_args.min_tiles, max_num=max_num, image_size=crop_size["height"])
+        images = [processor.preprocess(image, return_tensors="pt")["pixel_values"][0] for image in images]
+        return torch.stack(images)
+
+    if data_args.image_aspect_ratio == "resize":
+        image = image.resize((crop_size["width"], crop_size["height"]))
+    if data_args.image_aspect_ratio == "pad":
+
+        def expand2square(pil_img, background_color):
+            width, height = pil_img.size
+            if width == height:
+                return pil_img
+            elif width > height:
+                result = Image.new(pil_img.mode, (width, width), background_color)
+                result.paste(pil_img, (0, (width - height) // 2))
+                return result
+            else:
+                result = Image.new(pil_img.mode, (height, height), background_color)
+                result.paste(pil_img, ((height - width) // 2, 0))
+                return result
+
+        image = expand2square(image, tuple(int(x * 255) for x in processor.image_mean))
+        image = processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
+    else:
+        # Using default behavior of the vision encoder
+        # For CLIP, default is central crop
+        # For Radio, default is central crop
+        # For Siglip, default is resize
+        # For InternVIT, default is resize
+        image = processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
+    return image
+
+
+def process_images(images, image_processor, model_cfg, enable_dynamic_res=False, max_tiles=None):
+    model_cfg.image_processor = image_processor
+    new_images = [
+        process_image(image, model_cfg, None, enable_dynamic_res=enable_dynamic_res, max_tiles=max_tiles)
+        for image in images
+    ]
+
+    if all(x.shape == new_images[0].shape for x in new_images):
+        if len(new_images[0].shape) == 4:
+            new_images = torch.cat(new_images, dim=0)
+        elif len(new_images[0].shape) == 3:
+            new_images = torch.stack(new_images, dim=0)
+        else:
+            raise ValueError(f"new_images rank does not equal to 4, rank: {len(new_images[0].shape)}")
+    else:
+        raise ValueError("The shape of images in new_images is different!")
+    return new_images
+
+
+def tokenizer_image_token(prompt, tokenizer, return_tensors=None):
+    return tokenizer(prompt, return_tensors=return_tensors).input_ids[0]
+
+
+def is_gemma_tokenizer(tokenizer):
+    return "gemma" in tokenizer.__class__.__name__.lower()
+
+
+def get_model_name_from_path(model_path):
+    model_path = model_path.strip("/")
+    model_paths = model_path.split("/")
+    if model_paths[-1].startswith("checkpoint-"):
+        return model_paths[-2] + "_" + model_paths[-1]
+    else:
+        return model_paths[-1]
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.keyword_ids = []
+        self.max_keyword_len = 0
+        for keyword in keywords:
+            cur_keyword_ids = tokenizer(keyword).input_ids
+            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+                cur_keyword_ids = cur_keyword_ids[1:]
+            if len(cur_keyword_ids) > self.max_keyword_len:
+                self.max_keyword_len = len(cur_keyword_ids)
+            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+        self.tokenizer = tokenizer
+        self.start_len = input_ids.shape[1]
+
+    def call_for_batch(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+        for keyword_id in self.keyword_ids:
+            if (output_ids[0, -keyword_id.shape[0] :] == keyword_id).all():
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        for keyword in self.keywords:
+            if keyword in outputs:
+                return True
+        return False
+
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        outputs = []
+        for i in range(output_ids.shape[0]):
+            outputs.append(self.call_for_batch(output_ids[i].unsqueeze(0), scores))
+        return all(outputs)

model_utils_packing.py

@@ -0,0 +1,35 @@
+from importlib import import_module
+from typing import Tuple
+
+import torch
+import transformers
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ["patch"]
+
+
+def _get_unpad_data(attention_mask: torch.Tensor, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, int]:
+    if hasattr(_get_unpad_data, "seqlens_in_batch"):
+        seqlens_in_batch = _get_unpad_data.seqlens_in_batch
+    else:
+        seqlens_in_batch = torch.sum(attention_mask, dim=1)
+
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return indices, cu_seqlens, max_seqlen_in_batch
+
+
+def set_seqlens_in_batch(seqlens_in_batch: torch.Tensor) -> None:
+    _get_unpad_data.seqlens_in_batch = seqlens_in_batch
+
+
+def patch(model: nn.Module) -> None:
+    if transformers.__version__ < "4.43.0":
+        m = import_module(model.__module__)
+        if not hasattr(m, "_get_unpad_data"):
+            raise ValueError(f"Module {m} does not have function '_get_unpad_data' for packing")
+        m._get_unpad_data = _get_unpad_data
+    else:
+        transformers.modeling_flash_attention_utils._get_unpad_data = _get_unpad_data

modeling_vila.py

@@ -0,0 +1,1228 @@
+import copy
+import json
+import logging
+import math
+import os
+import os.path
+import os.path as osp
+import shutil
+import warnings
+from abc import ABC
+from collections import OrderedDict, defaultdict, deque
+from copy import deepcopy
+from itertools import chain
+from threading import Thread
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+from einops import rearrange
+from PIL import Image
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoProcessor,
+    AutoTokenizer,
+    GenerationConfig,
+    LogitsProcessor,
+    PretrainedConfig,
+    PreTrainedModel,
+    Qwen2Config,
+    Qwen2ForCausalLM,
+    Qwen2PreTrainedModel,
+    TextIteratorStreamer,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import ContextManagers, no_init_weights
+
+from .auto_processor import VILAProcessor
+from .base_projector import MultimodalProjector, MultimodalProjectorConfig
+from .builder import build_llm_and_tokenizer
+from .configuration_vila import VILAConfig
+from .constants import *
+from .conversation import SeparatorStyle, default_conversation
+from .distributed import all_gather as vila_all_gather
+from .loss import soft_cross_entropy
+from .media import extract_media
+from .media_encoder import BasicImageEncoder, BasicVideoEncoder
+from .mm_utils import process_image, process_images
+from .model_utils_packing import set_seqlens_in_batch
+from .siglip_encoder import SiglipVisionTower, SiglipVisionTowerDynamicS2, SiglipVisionTowerS2
+from .tokenizer_utils import tokenize_conversation
+from .utils import get_model_config, load_tokenizer_then_handle_media_tokens_and_chat_template
+
+# from llava.constants import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX, NUM_EXTRA_TOKENS
+
+# ease debugging
+python_input = input
+
+# quick hack for remote code
+def get_pg_manager():
+    return None
+
+
+def get_model_weights_dtype(model: nn.Module):
+    pass
+
+
+def build_mm_projector(model_type_or_path: str, config: PretrainedConfig) -> PreTrainedModel:
+    if model_type_or_path is None:
+        return None
+    ## load from pretrained model
+    if config.resume_path:
+        assert os.path.exists(model_type_or_path), f"Resume mm projector path {model_type_or_path} does not exist!"
+        return MultimodalProjector.from_pretrained(model_type_or_path, config)
+    ## build from scratch
+    else:
+        mm_projector_cfg = MultimodalProjectorConfig(model_type_or_path)
+        mm_projector = MultimodalProjector(mm_projector_cfg, config)
+        return mm_projector
+
+
+def check_dot_in_model_path(model_path: str):
+    """Check if the model path contains dot, which will affect the remote code loading."""
+    if osp.isdir(model_path):  # local model
+        if "." in osp.abspath(model_path):
+            return True
+    else:  # remote model
+        if "." in model_path:
+            return True
+    return False
+
+
+def get_vila_version(model_path: str) -> str:
+    VERSIONS = ["vila1.5", "vila-u", "longvila", "nvila", "vila-m3"]
+    for version in VERSIONS:
+        if version in model_path.lower():
+            return version
+    return None
+
+
+def generate_jinja_template(conv_mode: str) -> str:
+    if conv_mode == "vicuna_v1":
+        return """{% set system_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. " %}
+{% set roles = ["user", "assistant"] %}
+{% set sep = " " %}
+
+{{ system_prompt }}
+
+{% for message in messages %}
+    {% if message['role'] == roles[0] %}
+        {{ "USER: " }}{{ sep }}{{ message['content'] }}{{ sep }}
+    {% else %}
+        {{ "ASSISTANT: " }}{{ sep }}{{ message['content'] }}{{ sep }}
+    {% endif %}
+{% endfor %}
+{% if messages[-1]['role'] == 'user' %}
+    {{ "ASSISTANT:" }}
+{% endif %}
+"""
+    elif conv_mode == "llama_3":
+        return """{% set system_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\\n\\nYou are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.<|eot_id|>" %}
+{% set roles = ["<|start_header_id|>user<|end_header_id|>\\n\\n", "<|start_header_id|>assistant<|end_header_id|>\\n\\n"]%}
+{% set sep = "<|eot_id|>" %}
+
+{{ system_prompt }}
+{% for message in messages %}
+    {% if message['role'] == 'user' %}
+        {{ roles[0] }}{{ message['content'] }}{{ sep }}
+    {% else %}
+        {{ roles[1] }}{{ message['content'] }}{{ sep }}
+    {% endif %}
+{% endfor %}
+{% if messages[-1]['role'] == 'user' %}
+    {{ roles[1] }}
+{% endif %}
+"""
+    elif conv_mode == "hermes_2":
+        return """{% set system_prompt = "<|im_start|>system\nAnswer the questions." %}
+{% set roles = ["<|im_start|>user\n", "<|im_start|>assistant\n"] %}
+{% set sep = "<|im_end|>" %}
+
+{{ system_prompt }}{{ sep }}
+
+{% for message in messages %}
+    {% if message['role'] == 'user' %}
+        {{ roles[0] }}{{ message['content'] }}{{ sep }}
+    {% else %}
+        {{ roles[1] }}{{ message['content'] }}{{ sep }}
+    {% endif %}
+{% endfor %}"""
+    else:
+        raise NotImplementedError(f"Jinja template generation is not implemented for {conv_mode}.")
+
+
+def build_vision_tower(model_name_or_path: str, config: PretrainedConfig) -> PreTrainedModel:
+    ## skip vision tower instantiation
+    if model_name_or_path is None:
+        return None
+
+    vision_tower_arch = None
+    if config.resume_path and "radio" not in model_name_or_path:
+        assert os.path.exists(model_name_or_path), f"Resume vision tower path {model_name_or_path} does not exist!"
+        vision_tower_cfg = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
+        vision_tower_arch = vision_tower_cfg.architectures[0].lower()
+    vision_tower_name = vision_tower_arch if vision_tower_arch is not None else model_name_or_path
+
+    use_s2 = getattr(config, "s2", False)
+    use_dynamic_s2 = getattr(config, "dynamic_s2", False)
+
+    if "siglip" in vision_tower_name:
+        if use_dynamic_s2:
+            vision_tower = SiglipVisionTowerDynamicS2(model_name_or_path, config)
+        elif use_s2:
+            vision_tower = SiglipVisionTowerS2(model_name_or_path, config)
+        else:
+            vision_tower = SiglipVisionTower(model_name_or_path, config)
+    else:
+        raise NotImplementedError(f"Unknown vision tower: {model_name_or_path}")
+
+    config.mm_hidden_size = (
+        vision_tower.config.hidden_size if not (use_s2 or use_dynamic_s2) else vision_tower.hidden_size
+    )
+    return vision_tower
+
+
+class VILAPretrainedModel(PreTrainedModel):
+    config_class = VILAConfig
+    main_input_name = "input_embeds"
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config: VILAConfig, *args, **kwargs):
+        super().__init__(config)
+        self.config = config
+        cfgs = get_model_config(config)
+        if len(cfgs) == 3:
+            llm_cfg, vision_tower_cfg, mm_projector_cfg = cfgs
+        else:
+            raise ValueError("`llm_cfg` `mm_projector_cfg` `vision_tower_cfg` not found in the config.")
+
+        # loading on cpu by default
+        device_map = kwargs.get("device_map", "cpu")
+        self.mm_projector = build_mm_projector(mm_projector_cfg, config)
+        self.vision_tower = build_vision_tower(vision_tower_cfg, config)
+        if "auto" in device_map or "cuda" in device_map:
+            self.mm_projector = self.mm_projector.cuda()
+            self.vision_tower = self.vision_tower.cuda()
+        # set device_map auto can autoamtically shard llm to different devices
+        self.llm, self.tokenizer = self.init_llm(llm_cfg, config, device_map=device_map)
+
+        # NOTE(ligeng): need to add other decoders from config
+        self.encoders = {"image": BasicImageEncoder(self), "video": BasicVideoEncoder(self)}
+
+        self.post_config()
+        self.is_loaded = True
+
+        assert (
+            self.llm is not None or self.vision_tower is not None or self.mm_projector is not None
+        ), "At least one of the components must be instantiated."
+
+    @classmethod
+    def convert_vila_dev_ckpt_to_remote(
+        self,
+        model_path: str,
+        output_dir: str = None,
+        vila_version: str | None = None,
+        conv_mode: str | None = None,
+        copy: bool = False,
+        copy_weights: bool = True,
+        copy_code: bool = True,
+        *model_args,
+        **kwargs,
+    ):
+        # assert type(self) == VILAForCasualLM, "This method is only available for VILAForCasualLM."
+        assert model_path != output_dir, "model_path and output_dir cannot be the same"
+        if os.path.isdir(model_path):
+            model_path = model_path
+        else:
+            from huggingface_hub import HfApi, snapshot_download
+
+            model_path = snapshot_download(model_path)
+            print("downloading HF model to", model_path)
+
+        if check_dot_in_model_path(model_path) and output_dir is None:
+            raise ValueError(
+                f"Model path {model_path} contains a dot, which will affect the remote code loading. Please specify the output directory without dot in the path to fix this issue."
+            )
+        if output_dir is not None and "." in output_dir:
+            raise ValueError(
+                f"Output directory {output_dir} contains a dot, which will affect the remote code loading. Please specify a valid output directory without dots."
+            )
+
+        if copy:
+            print("copy is set to True, copying weights and code to output_dir")
+            copy_weights = copy_code = True
+        # copy weights and code to output_dir
+        self.copy_or_symlink_directory(model_path, output_dir, copy=copy_weights)
+        self.copy_remote_py_files(output_dir, copy=copy_code)
+
+        if vila_version is None:
+            vila_version = get_vila_version(output_dir)
+
+        cfg_path = os.path.join(output_dir, "config.json")
+        config = json.load(open(cfg_path))
+        config["version"] = "2.0"  # nvila tag
+        config["architectures"] = ["VILAForCasualLM"]
+        config["auto_map"] = {
+            "AutoProcessor": "auto_processor.VILAProcessor",
+            "AutoConfig": "modeling_vila.VILAConfig",
+            "AutoModel": "modeling_vila.VILAForCasualLM",
+            "AutoModelForCausalLM": "modeling_vila.VILAForCasualLM",
+        }
+        # vila1.5 legacy support
+        config["model_type"] = "vila"
+        if vila_version in ["vila1.5", "vila-m3"]:
+            if conv_mode is None:
+                raise ValueError(f"Please specify the conversation mode for {output_dir}.")
+            config["chat_template"] = conv_mode
+            jinja_template = generate_jinja_template(conv_mode)
+            jinja_path = os.path.join(output_dir, f"{conv_mode}.jinja")
+            with open(jinja_path, "w") as f:
+                f.write(jinja_template)
+        json.dump(config, open(cfg_path, "w"), indent=2)
+
+        ##########################################################################################
+        config = AutoConfig.from_pretrained(output_dir, trust_remote_code=True)
+        tokenizer = load_tokenizer_then_handle_media_tokens_and_chat_template(output_dir, config)
+        tokenizer.save_pretrained(osp.join(output_dir, "llm"))
+        ##########################################################################################
+
+    @classmethod
+    def copy_or_symlink_directory(cls, model_path, output_dir, copy=True):
+        # Create output directory if it doesn't exist
+        os.makedirs(output_dir, exist_ok=True)
+        # Create symlinks for all files in model_path to output_dir
+        for item in os.listdir(model_path):
+            src_path = os.path.join(model_path, item)
+            dst_path = os.path.join(output_dir, item)
+
+            # Remove existing file/directory at destination if it exists
+            if os.path.exists(dst_path):
+                if os.path.islink(dst_path):
+                    os.unlink(dst_path)
+                elif os.path.isdir(dst_path):
+                    shutil.rmtree(dst_path)
+                else:
+                    os.remove(dst_path)
+
+            # Create symlink
+            if copy:
+                if os.path.isdir(src_path):
+                    shutil.copytree(src_path, dst_path)
+                else:
+                    shutil.copy2(src_path, dst_path)
+                print(f"Copied {src_path} to {dst_path}")
+            else:
+                os.symlink(src_path, dst_path)
+                print(f"Created symlink from {src_path} to {dst_path}")
+
+    @classmethod
+    def copy_remote_py_files(cls, output_dir, copy=True):
+        ## copy .py and REAMDE for next loading remote code
+        current_file_path = os.path.abspath(__file__)
+        current_folder = os.path.dirname(current_file_path)
+        for file_name in os.listdir(current_folder):
+            if file_name == "INSTRUCTIONS.md":
+                src_fname = os.path.join(current_folder, file_name)
+                dst_fname = os.path.join(output_dir, "README.md")
+                if os.path.exists(dst_fname):
+                    old_reamde = open(dst_fname).read()
+                else:
+                    old_reamde = ""
+                with open(src_fname) as src, open(dst_fname, "w") as dst:
+                    dst.write(src.read())
+                    dst.write(old_reamde)
+                print("[HF remote code] REAMDE ", src_fname, "to", dst_fname)
+            if file_name.endswith(".py") or file_name.endswith(".jinja"):
+                full_file_name = os.path.join(current_folder, file_name)
+                if os.path.isfile(full_file_name):
+                    if copy:
+                        shutil.copy(full_file_name, output_dir)
+                        print("[HF remote code] copying", full_file_name, "to", output_dir)
+                    else:
+                        # symlink to ease development
+                        if os.path.exists(os.path.join(output_dir, file_name)):
+                            os.remove(os.path.join(output_dir, file_name))
+                        os.symlink(full_file_name, os.path.join(output_dir, file_name))
+                        print("[HF remote code] linking", full_file_name, "to", output_dir)
+
+    def save_pretrained(self, output_dir, state_dict=None, **kwargs):
+        if state_dict is None:
+            # other wise fetch from deepspeed
+            # state_dict = accelerator.get_state_dict(is_deepspeed_enabled)
+            state_dict = self.state_dict()
+
+        if getattr(self, "tokenizer", None):
+            self.tokenizer.save_pretrained(osp.join(output_dir, "llm"))
+
+        if self.get_llm():
+            print(f"saving llm to {osp.join(output_dir, 'llm')}")
+            self.llm.config._name_or_path = osp.join(output_dir, "llm")
+            llm_state_dict = OrderedDict({k.split("llm.")[-1]: v for k, v in state_dict.items() if "llm" in k})
+            self.llm.save_pretrained(os.path.join(output_dir, "llm"), state_dict=llm_state_dict)
+            self.config.llm_cfg = self.llm.config
+
+        if self.get_vision_tower():
+            print(f"saving vision_tower to {osp.join(output_dir, 'vision_tower')}")
+            self.vision_tower.config._name_or_path = osp.join(output_dir, "vision_tower")
+            vision_tower_state_dict = OrderedDict(
+                {k.split("vision_tower.vision_tower.")[-1]: v for k, v in state_dict.items() if "vision_tower" in k}
+            )
+            self.vision_tower.vision_tower.save_pretrained(
+                os.path.join(output_dir, "vision_tower"),
+                state_dict=vision_tower_state_dict,
+            )
+            self.vision_tower.image_processor.save_pretrained(os.path.join(output_dir, "vision_tower"))
+            self.config.vision_tower_cfg = self.vision_tower.config
+            if hasattr(self.config.vision_tower_cfg, "auto_map"):
+                if "radio" not in self.get_vision_tower().__class__.__name__.lower():
+                    delattr(self.config.vision_tower_cfg, "auto_map")
+
+        if self.get_mm_projector():
+            print(f"saving mm_projector to {osp.join(output_dir, 'mm_projector')}")
+            self.mm_projector.config._name_or_path = osp.join(output_dir, "mm_projector")
+            mm_projector_state_dict = OrderedDict(
+                {k.split("mm_projector.")[-1]: v for k, v in state_dict.items() if "mm_projector" in k}
+            )
+            self.mm_projector.save_pretrained(
+                os.path.join(output_dir, "mm_projector"),
+                state_dict=mm_projector_state_dict,
+            )
+            self.config.mm_projector_cfg = self.mm_projector.config
+
+        ## update and save top-level config
+        self.config._name_or_path = output_dir
+        self.config.architectures = [self.__class__.__name__]
+        self.config.save_pretrained(output_dir)
+
+        ## copy .py and REAMDE for next loading remote code
+        self.copy_remote_py_files(output_dir)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Optional[str] = None,
+        *model_args,
+        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
+        cache_dir: Optional[Union[str, os.PathLike]] = None,
+        ignore_mismatched_sizes: bool = False,
+        force_download: bool = False,
+        local_files_only: bool = False,
+        token: Optional[Union[str, bool]] = None,
+        revision: str = "main",
+        use_safetensors: Optional[bool] = None,
+        weights_only: bool = True,
+        **kwargs,
+    ):
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        return cls._from_config(config, **kwargs)
+
+    def init_llm(self, llm_config, config, *args, **kwargs):
+        self.llm, self.tokenizer = build_llm_and_tokenizer(llm_config, config, *args, **kwargs)
+        # hard coded for NVILA
+        # variables for XGrammar
+        # print("DEBUG", len(self.tokenizer.added_tokens_encoder.keys()), self.tokenizer.added_tokens_encoder.keys())
+        NUM_EXTRA_TOKENS = len(self.tokenizer.added_tokens_encoder.keys())
+
+        # TODO: SENTINEL_TOKEN is not added, need to check with Zhijian
+        self.vocab_size = self.tokenizer.vocab_size + NUM_EXTRA_TOKENS
+        # XGrammar tokenizer and grammar compiler
+        # lazy init only when specified json output during inference
+        self.grammar_compiler = None
+        self.llm.resize_token_embeddings(len(self.tokenizer))
+        return self.llm, self.tokenizer
+
+    def post_config(self):
+        ######################################################################
+        # TODO: need to check dtype with jason
+        self.llm = self.llm.to(torch.float16)
+        self.mm_projector = self.mm_projector.to(torch.float16)
+        self.vision_tower = self.vision_tower.to(torch.float16)
+        ######################################################################
+        self.training = self.llm.training
+        ## configuration
+        if getattr(self.config, "llm_cfg", None) is None:
+            self.config.llm_cfg = self.llm.config
+        if getattr(self.config, "vision_tower_cfg", None) is None:
+            self.config.vision_tower_cfg = self.vision_tower.config
+        if getattr(self.config, "mm_projector_cfg", None) is None:
+            self.config.mm_projector_cfg = self.mm_projector.config
+
+    def get_llm(self):
+        llm = getattr(self, "llm", None)
+        if type(llm) is list:
+            llm = llm[0]
+        return llm
+
+    def get_lm_head(self):
+        lm_head = getattr(self.get_llm(), "lm_head", None)
+        return lm_head
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, "vision_tower", None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def get_mm_projector(self):
+        mm_projector = getattr(self, "mm_projector", None)
+        if type(mm_projector) is list:
+            mm_projector = mm_projector[0]
+        return mm_projector
+
+    def freezed_module_patch(self):
+        """
+        Huggingface will call model.train() at each training_step. To ensure the expected behaviors for modules like dropout, batchnorm, etc., we need to call model.eval() for the freezed modules.
+        """
+        if self.training:
+            if self.get_llm() and not getattr(self.config, "tune_language_model", False):
+                pass
+                # logging.warning("Caution: Your LLM is currently in training mode, ensuring accurate gradient computation. Please be vigilant, particularly regarding BatchNorm and Dropout operations.")
+            if self.get_vision_tower() and not getattr(self.config, "tune_vision_tower", False):
+                self.get_vision_tower().eval()
+            if self.get_mm_projector() and not getattr(self.config, "tune_mm_projector", False):
+                self.get_mm_projector().eval()
+
+
+class VILAForCasualLM(VILAPretrainedModel):
+    def __init__(self, config: VILAConfig, *args, **kwargs):
+        super().__init__(config, *args, **kwargs)
+
+    def merge_features_for_dynamic_s2(self, image_features, block_sizes):
+        scales = self.get_vision_tower().scales
+        resize_output_to_scale_idx = self.get_vision_tower().resize_output_to_scale_idx
+
+        image_features_each_image = []
+        new_block_sizes = []
+        block_cnt = 0
+        for block_size_each_image in block_sizes:
+            if block_size_each_image is None:
+                cur_features = image_features[block_cnt : block_cnt + 1]
+                cur_features = rearrange(cur_features, "1 (h w) c -> 1 c h w", h=int(cur_features.shape[1] ** 0.5))
+                cur_features = cur_features.repeat(1, len(scales), 1, 1)
+                image_features_each_image.append(cur_features)
+                new_block_sizes.append((1, 1))
+                block_cnt += 1
+            else:
+                cur_features_each_scale = []
+                for scale in scales[:-1]:
+                    num_blocks_this_scale = (scale // scales[0]) ** 2
+                    cur_features_each_scale.append(
+                        self.merge_chessboard(
+                            image_features[block_cnt : block_cnt + num_blocks_this_scale],
+                            num_split_h=scale // scales[0],
+                            num_split_w=scale // scales[0],
+                        )
+                    )  # 1 * C * H * W
+                    block_cnt += num_blocks_this_scale
+                num_blocks_last_scale = block_size_each_image[0] * block_size_each_image[1]
+                cur_features_each_scale.append(
+                    self.merge_chessboard(
+                        image_features[block_cnt : block_cnt + num_blocks_last_scale],
+                        num_split_h=block_size_each_image[0],
+                        num_split_w=block_size_each_image[1],
+                    )
+                )  # 1 * C * H * W
+                block_cnt += num_blocks_last_scale
+
+                # resize and concat features from different scales
+                output_size = cur_features_each_scale[resize_output_to_scale_idx].shape[-2:]
+                cur_features = torch.cat(
+                    [
+                        F.interpolate(cur_features_each_scale[i].to(torch.float32), size=output_size, mode="area").to(
+                            cur_features_each_scale[i].dtype
+                        )
+                        for i in range(len(cur_features_each_scale))
+                    ],
+                    dim=1,
+                )
+                # cur_features = rearrange(cur_features, "1 c h w -> (h w) c")
+
+                image_features_each_image.append(cur_features)
+
+                if resize_output_to_scale_idx == len(scales) - 1 or resize_output_to_scale_idx == -1:
+                    new_block_sizes.append(block_size_each_image)
+                else:
+                    new_block_sizes.append(
+                        (
+                            scales[resize_output_to_scale_idx] // scales[0],
+                            scales[resize_output_to_scale_idx] // scales[0],
+                        )
+                    )
+
+        assert block_cnt == len(image_features)
+
+        return image_features_each_image, new_block_sizes
+
+    def encode_images(self, images, block_sizes: Optional[Optional[Tuple[int, ...]]] = None):
+        if block_sizes is None:
+            block_sizes = [None] * len(images)
+        if getattr(self.config, "dynamic_s2", False):
+            image_features = self.get_vision_tower()(images)
+            image_features, new_block_sizes = self.merge_features_for_dynamic_s2(image_features, block_sizes)
+
+            image_features = [
+                self.split_chessboard(x, block_size[0], block_size[1])
+                for x, block_size in zip(image_features, new_block_sizes)
+            ]  # list of B * C * H * W tensors
+            image_features = torch.cat(
+                [rearrange(x, "b c h w -> b (h w) c") for x in image_features], dim=0
+            )  # B * N * C
+            image_features = self.get_mm_projector()(image_features)
+            image_features = list(
+                image_features.split([block_size[0] * block_size[1] for block_size in new_block_sizes], dim=0)
+            )
+            image_features = [
+                self.merge_chessboard(x, block_size[0], block_size[1])
+                for x, block_size in zip(image_features, new_block_sizes)
+            ]  # list of 1 * C * H * W tensors
+            image_features = [rearrange(x, "1 c h w -> (h w) c") for x in image_features]  # list of N * C tensors
+            if all([feature.shape[0] == image_features[0].shape[0] for feature in image_features]):
+                image_features = torch.stack(image_features, dim=0)
+        else:
+            image_features = self.get_vision_tower()(images)
+            image_features = self.get_mm_projector()(image_features)
+        return image_features
+
+    def train(self, mode: bool = True):
+        if mode:
+            self.tokenizer.padding_side = "right"
+        else:
+            self.tokenizer.padding_side = "left"
+        super().train(mode)
+        return self
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+        labels: Optional[torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # NOTE(ligeng): deep copy to avoid modifying the original media and media_config
+        media = copy.deepcopy(media)
+        media_config = copy.deepcopy(media_config)
+
+        labels = labels if labels is not None else torch.full_like(input_ids, IGNORE_INDEX)
+        attention_mask = attention_mask if attention_mask is not None else torch.ones_like(input_ids, dtype=torch.bool)
+
+        PROCESS_GROUP_MANAGER = get_pg_manager()
+        if PROCESS_GROUP_MANAGER is not None:
+            for name in media:
+                self.encoders[name].end_tokens = None
+
+        # Extract text and media embeddings
+        text_embeds = self.llm.model.embed_tokens(input_ids)
+        if media is not None:
+            media_embeds = self.__embed_media_tokens(media, media_config)
+        else:
+            # no media was provided, so we just return an empty dict
+            media_embeds = {}
+
+        # This is a workaround to make sure the dummy embeddings are consumed
+        while media_embeds.get("dummy"):
+            dummy_embed = media_embeds["dummy"].popleft()
+            text_embeds += torch.sum(dummy_embed) * 0
+
+        # Remove padding
+        batch_size = labels.shape[0]
+        text_embeds = [text_embeds[k][attention_mask[k]] for k in range(batch_size)]
+        labels = [labels[k][attention_mask[k]] for k in range(batch_size)]
+
+        # Build inverse mapping from token ID to media name
+        media_tokens = {}
+        for name, token_id in self.tokenizer.media_token_ids.items():
+            media_tokens[token_id] = name
+
+        # Fuse text and media embeddings
+        inputs_m, labels_m = [], []
+        for k in range(batch_size):
+            inputs_mk, labels_mk = [], []
+            pos = 0
+            while pos < len(labels[k]):
+                if input_ids[k][pos].item() in media_tokens:
+                    end = pos + 1
+                    name = media_tokens[input_ids[k][pos].item()]
+                    input = media_embeds[name].popleft()
+                    label = torch.full([input.shape[0]], IGNORE_INDEX, device=labels[k].device, dtype=labels[k].dtype)
+                    # print(f"{self.tokenizer.padding_side} [media] {k=} {pos=}, {self.tokenizer.batch_decode(input_ids[k][pos:pos+1])}"); python_input()
+                elif input_ids[k][pos].item() in (self.tokenizer.pad_token_id, self.tokenizer.eos_token_id):
+                    end = pos + 1
+                    pos = end
+                    # print(f"[skip PAD/EOS] {k=} {pos=}, {self.tokenizer.batch_decode(input_ids[k][pos:end])}"); python_input()
+                    continue
+                else:
+                    end = pos
+                    while end < len(labels[k]) and input_ids[k][end].item() not in media_tokens:
+                        end += 1
+                    input = text_embeds[k][pos:end]
+                    label = labels[k][pos:end]
+                    # print(f"[text] {k=} {pos=}, {self.tokenizer.batch_decode(input_ids[k][pos:end])}"); python_input()
+
+                inputs_mk.append(input)
+                labels_mk.append(label)
+                pos = end
+            inputs_m.append(torch.cat(inputs_mk, dim=0))
+            labels_m.append(torch.cat(labels_mk, dim=0))
+        inputs, labels = inputs_m, labels_m
+
+        # Check if all media embeddings are consumed
+        for name in media_embeds:
+            if media_embeds[name]:
+                raise ValueError(f"Not all {name} embeddings are consumed! Still {len(media_embeds[name])} left.")
+
+        # Truncate sequences to `model_max_length` as media embeddings are inserted
+        inputs, labels = self.__truncate_sequence(inputs, labels)
+
+        # Pad sequences to the longest one in the batch
+        return self.__batchify_sequence(inputs, labels)
+
+    def __embed_media_tokens(
+        self,
+        media: Dict[str, List[torch.Tensor]],
+        media_config: Dict[str, Dict[str, Any]],
+    ) -> Dict[str, List[torch.Tensor]]:
+        embeds = defaultdict(deque)
+        for name in media:
+            if self.training:
+                # Gather metainfo of media objects from all ranks
+                info = [{"shape": tensor.shape, "dtype": tensor.dtype} for tensor in media.get(name, [])]
+                infos = list(chain(vila_all_gather(info)))
+
+                # The entire batch does not contain any media objects of this type.
+                if not infos:
+                    continue
+
+                # Create a dummy tensor to ensure the encoder is called, otherwise the training will hang.
+                if media.get(name) is None or len(media[name]) == 0:
+                    dummy = torch.zeros(infos[0]["shape"], dtype=infos[0]["dtype"], device=self.device)
+                    embeds["dummy"].extend(self.encoders[name]([dummy], media_config[name]))
+                    continue
+            embeds[name] = deque(self.encoders[name](media[name], media_config[name]))
+        return embeds
+
+    def __truncate_sequence(
+        self, inputs: List[torch.Tensor], labels: List[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.training and any(len(input) > self.tokenizer.model_max_length for input in inputs):
+            warnings.warn(f"Truncating sequences to `model_max_length` ({self.tokenizer.model_max_length}).")
+            inputs = [input[: self.tokenizer.model_max_length] for input in inputs]
+            labels = [label[: self.tokenizer.model_max_length] for label in labels]
+        return inputs, labels
+
+    def __batchify_sequence(
+        self, inputs: List[torch.Tensor], labels: List[torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size = len(inputs)
+        device = inputs[0].device
+        hidden_size = inputs[0].shape[1]
+        max_length = max(inputs[k].shape[0] for k in range(batch_size))
+        attention_mask = torch.ones((batch_size, max_length), dtype=torch.bool, device=device)
+
+        inputs_p, labels_p = [], []
+        for k in range(batch_size):
+            size_pk = max_length - inputs[k].shape[0]
+            inputs_pk = torch.zeros((size_pk, hidden_size), dtype=inputs[k].dtype, device=device)
+            labels_pk = torch.full((size_pk,), IGNORE_INDEX, dtype=labels[k].dtype, device=device)
+            if self.tokenizer.padding_side == "right":
+                attention_mask[k, inputs[k].shape[0] :] = False
+                inputs_pk = torch.cat([inputs[k], inputs_pk], dim=0)
+                labels_pk = torch.cat([labels[k], labels_pk], dim=0)
+            else:
+                attention_mask[k, : -inputs[k].shape[0]] = False
+                inputs_pk = torch.cat([inputs_pk, inputs[k]], dim=0)
+                labels_pk = torch.cat([labels_pk, labels[k]], dim=0)
+            inputs_p.append(inputs_pk)
+            labels_p.append(labels_pk)
+
+        inputs = torch.stack(inputs_p, dim=0)
+        labels = torch.stack(labels_p, dim=0)
+        return inputs, labels, attention_mask
+
+    def repack_multimodal_data(self, inputs_embeds, attention_mask, position_ids, labels):
+        # Handle sequence parallelism
+        PROCESS_GROUP_MANAGER = get_pg_manager()
+
+        # We do re-sharding instead of packing here to ensure the sequence length is the same across all ranks.
+        if PROCESS_GROUP_MANAGER is not None:
+            sp_degree = PROCESS_GROUP_MANAGER.sp_degree
+            sp_rank = PROCESS_GROUP_MANAGER.sp_rank
+            sp_group = PROCESS_GROUP_MANAGER.sp_pg
+            ring_degree = PROCESS_GROUP_MANAGER.ring_degree
+            ring_rank = PROCESS_GROUP_MANAGER.ring_rank
+            ring_type = PROCESS_GROUP_MANAGER.ring_type
+            ulysses_degree = PROCESS_GROUP_MANAGER.ulysses_degree
+            ulysses_rank = PROCESS_GROUP_MANAGER.ulysses_rank
+
+            bs, shard_seqlen = position_ids.shape
+            sp_seq_len = [torch.zeros(1, dtype=torch.int64, device=position_ids.device) for _ in range(sp_degree)]
+            dist.all_gather(sp_seq_len, torch.tensor(shard_seqlen, device=position_ids.device), group=sp_group)
+            sp_seq_len_cat = torch.cat(sp_seq_len, dim=0)
+
+            if sp_rank == 0:
+                original_start_id = 0
+            else:
+                original_start_id = torch.sum(sp_seq_len_cat[:sp_rank]).item()
+            original_end_id = torch.sum(sp_seq_len_cat[: sp_rank + 1]).item()
+
+            # Gather attention_mask, position_ids, labels and input_embeds
+            all_inputs_embeds = torch.zeros(
+                bs,
+                torch.sum(sp_seq_len_cat),
+                inputs_embeds.shape[-1],
+                dtype=inputs_embeds.dtype,
+                device=inputs_embeds.device,
+            ).contiguous()
+            all_inputs_embeds[:, original_start_id:original_end_id, :] += inputs_embeds
+            dist.barrier(group=sp_group)
+            dist.all_reduce(all_inputs_embeds, group=sp_group)
+            dist.barrier(group=sp_group)
+
+            attention_mask_list = [
+                torch.zeros((bs, sp_seq_len[i]), dtype=attention_mask.dtype, device=attention_mask.device)
+                for i in range(sp_degree)
+            ]
+            position_ids_list = [
+                torch.zeros((bs, sp_seq_len[i]), dtype=position_ids.dtype, device=position_ids.device)
+                for i in range(sp_degree)
+            ]
+            labels_list = [
+                torch.zeros((bs, sp_seq_len[i]), dtype=labels.dtype, device=labels.device) for i in range(sp_degree)
+            ]
+
+            dist.all_gather(attention_mask_list, attention_mask, group=sp_group)
+            dist.all_gather(position_ids_list, position_ids, group=sp_group)
+            dist.all_gather(labels_list, labels, group=sp_group)
+
+            effective_seqlen_list = [attention_mask_list[i].sum(dim=-1) for i in range(sp_degree)]
+            effective_seqlen = torch.stack(effective_seqlen_list, dim=-1)
+            effective_seqlen_batch_list = torch.unbind(effective_seqlen, dim=0)
+
+            global_attention_mask_list = []
+            global_position_ids_list = []
+            global_labels_list = []
+            global_inputs_embeds_list = []
+            for i in range(bs):
+                global_attention_mask_batch_list = []
+                global_position_ids_batch_list = []
+                global_labels_batch_list = []
+                global_inputs_embeds_batch_list = []
+                for j in range(sp_degree):
+                    eff_len = effective_seqlen_batch_list[i][j]
+                    prev_len = torch.sum(sp_seq_len_cat[:j]).item() if j > 0 else 0
+
+                    global_attention_mask_batch_list.append(attention_mask_list[j][i, :eff_len])
+                    global_position_ids_batch_list.append(position_ids_list[j][i, :eff_len])
+                    global_labels_batch_list.append(labels_list[j][i, :eff_len])
+                    global_inputs_embeds_batch_list.append(all_inputs_embeds[i, prev_len : prev_len + eff_len, :])
+                global_attention_mask_list.append(torch.cat(global_attention_mask_batch_list, dim=0))
+                global_position_ids_list.append(torch.cat(global_position_ids_batch_list, dim=0))
+                global_labels_list.append(torch.cat(global_labels_batch_list, dim=0))
+                global_inputs_embeds_list.append(torch.cat(global_inputs_embeds_batch_list, dim=0))
+
+                global_attention_mask = torch.nn.utils.rnn.pad_sequence(
+                    global_attention_mask_list, batch_first=True, padding_value=False
+                )
+                global_position_ids = torch.nn.utils.rnn.pad_sequence(
+                    global_position_ids_list, batch_first=True, padding_value=-1
+                )
+                global_labels = torch.nn.utils.rnn.pad_sequence(
+                    global_labels_list, batch_first=True, padding_value=IGNORE_INDEX
+                )
+                global_inputs_embeds = torch.nn.utils.rnn.pad_sequence(
+                    global_inputs_embeds_list, batch_first=True, padding_value=0
+                )
+
+            # Re-shard the inputs
+            if ring_degree > 1:
+                total_effective_seqlen = torch.sum(effective_seqlen, dim=1)
+                new_seqlen_per_rank = total_effective_seqlen // sp_degree
+                assert torch.all(
+                    total_effective_seqlen % sp_degree == 0
+                ), "total_effective_seqlen must be divisible by sp_degree"
+
+                max_new_seqlen = torch.max(new_seqlen_per_rank).item()
+
+                new_attention_mask = torch.zeros(
+                    (bs, max_new_seqlen), dtype=global_attention_mask.dtype, device=global_attention_mask.device
+                )
+                new_position_ids = torch.zeros(
+                    (bs, max_new_seqlen), dtype=global_position_ids.dtype, device=global_position_ids.device
+                )
+                new_labels = torch.full(
+                    (bs, max_new_seqlen), IGNORE_INDEX, dtype=global_labels.dtype, device=global_labels.device
+                )
+                new_inputs_embeds = torch.zeros(
+                    (bs, max_new_seqlen, global_inputs_embeds.shape[-1]),
+                    dtype=global_inputs_embeds.dtype,
+                    device=global_inputs_embeds.device,
+                )
+
+                if ring_type == "ring_varlen":
+                    for i in range(bs):
+                        start_idx = new_seqlen_per_rank[i] * sp_rank
+                        end_idx = start_idx + new_seqlen_per_rank[i]
+                        new_attention_mask[i, : new_seqlen_per_rank[i]] = global_attention_mask[i, start_idx:end_idx]
+                        new_position_ids[i, : new_seqlen_per_rank[i]] = global_position_ids[i, start_idx:end_idx]
+                        new_labels[i, : new_seqlen_per_rank[i]] = global_labels[i, start_idx:end_idx]
+                        new_inputs_embeds[i, : new_seqlen_per_rank[i], :] = global_inputs_embeds[
+                            i, start_idx:end_idx, :
+                        ]
+                elif ring_type == "zigzag_ring_varlen":
+                    chunk_size = total_effective_seqlen // (2 * sp_degree)
+                    for i in range(bs):
+                        # Zigzag pattern indices
+                        if sp_degree == ring_degree:
+                            forward_rank_idx = sp_rank
+                            backward_rank_idx = 2 * sp_degree - sp_rank - 1
+                        else:
+                            ulysses_offset = ulysses_rank * ring_degree * 2
+                            forward_rank_idx = ring_rank + ulysses_offset
+                            backward_rank_idx = sp_degree - ring_rank - 1 + ulysses_offset
+
+                        # Calculate start and end indices for the forward and backward zigzag
+                        start_idx_fwd = forward_rank_idx * chunk_size[i]
+                        end_idx_fwd = start_idx_fwd + chunk_size[i]
+
+                        start_idx_bwd = backward_rank_idx * chunk_size[i]
+                        end_idx_bwd = start_idx_bwd + chunk_size[i]
+
+                        # Fill new tensors with zigzag data
+                        new_attention_mask[i, : chunk_size[i]] = global_attention_mask[i, start_idx_fwd:end_idx_fwd]
+                        new_attention_mask[i, chunk_size[i] : 2 * chunk_size[i]] = global_attention_mask[
+                            i, start_idx_bwd:end_idx_bwd
+                        ]
+
+                        new_position_ids[i, : chunk_size[i]] = global_position_ids[i, start_idx_fwd:end_idx_fwd]
+                        new_position_ids[i, chunk_size[i] : 2 * chunk_size[i]] = global_position_ids[
+                            i, start_idx_bwd:end_idx_bwd
+                        ]
+
+                        new_labels[i, : chunk_size[i]] = global_labels[i, start_idx_fwd:end_idx_fwd]
+                        new_labels[i, chunk_size[i] : 2 * chunk_size[i]] = global_labels[i, start_idx_bwd:end_idx_bwd]
+
+                        new_inputs_embeds[i, : chunk_size[i], :] = global_inputs_embeds[i, start_idx_fwd:end_idx_fwd, :]
+                        new_inputs_embeds[i, chunk_size[i] : 2 * chunk_size[i], :] = global_inputs_embeds[
+                            i, start_idx_bwd:end_idx_bwd, :
+                        ]
+                else:
+                    raise ValueError(f"Invalid ring_type: {ring_type}")
+            else:
+                global_seq_len = global_attention_mask.shape[-1]
+                seq_len_sharded = global_seq_len // sp_degree
+                start_idx_reshard = seq_len_sharded * sp_rank
+                end_idx_reshard = start_idx_reshard + seq_len_sharded if sp_rank < sp_degree - 1 else global_seq_len
+
+                new_attention_mask = torch.narrow(
+                    global_attention_mask, 1, start_idx_reshard, end_idx_reshard - start_idx_reshard
+                )
+                new_position_ids = torch.narrow(
+                    global_position_ids, 1, start_idx_reshard, end_idx_reshard - start_idx_reshard
+                )
+                new_labels = torch.narrow(global_labels, 1, start_idx_reshard, end_idx_reshard - start_idx_reshard)
+                new_inputs_embeds = torch.narrow(
+                    global_inputs_embeds, 1, start_idx_reshard, end_idx_reshard - start_idx_reshard
+                )
+
+            return new_inputs_embeds, new_attention_mask, new_position_ids, new_labels
+
+        device = inputs_embeds.device
+        batch_size = inputs_embeds.shape[0]
+        seqlens = [attention_mask[k].sum().item() for k in range(batch_size)]
+
+        # Pack all sequences together
+        inputs_embeds_p = [inputs_embeds[k][attention_mask[k]] for k in range(batch_size)]
+        attention_mask_p = [torch.ones(seqlens[k], dtype=torch.int, device=device) for k in range(batch_size)]
+        position_ids_p = [torch.arange(seqlens[k], dtype=torch.int, device=device) for k in range(batch_size)]
+        labels_p = [labels[k][attention_mask[k]] for k in range(batch_size)]
+
+        # Add one dummy token at the end of the packed sequence to ensure that `_get_unpacked_data` will be called
+        inputs_embeds_p.append(torch.zeros(1, inputs_embeds.shape[-1], dtype=inputs_embeds.dtype, device=device))
+        attention_mask_p.append(torch.tensor([0], dtype=torch.int, device=device))
+        position_ids_p.append(torch.tensor([0], dtype=torch.int, device=device))
+        labels_p.append(torch.tensor([IGNORE_INDEX], dtype=torch.int, device=device))
+
+        # Mask the first token of each sequence to avoid contamination
+        for label in labels_p:
+            label[0] = IGNORE_INDEX
+
+        # Batch the data
+        inputs_embeds_p = torch.cat(inputs_embeds_p, dim=0).unsqueeze(0)
+        attention_mask_p = torch.cat(attention_mask_p, dim=0).unsqueeze(0)
+        position_ids_p = torch.cat(position_ids_p, dim=0).unsqueeze(0)
+        labels_p = torch.cat(labels_p, dim=0).unsqueeze(0)
+
+        if hasattr(
+            self, "pad_to_multiple_of"
+        ):  # related to quantization, please refer to ModelArguments for more information.
+            assert len(labels_p.shape) == 2
+            batch_size, max_length, cur_length = labels_p.shape[0], labels_p.shape[1], labels_p.shape[1]
+            hidden_size = inputs_embeds_p.shape[-1]
+
+            if max_length % self.pad_to_multiple_of != 0:
+                max_length = ((max_length // self.pad_to_multiple_of) + 1) * self.pad_to_multiple_of
+                difference = max_length - cur_length
+
+                inputs_embeds_p = torch.cat(
+                    (
+                        inputs_embeds_p,
+                        torch.full((batch_size, difference, hidden_size), self.llm.pad_token_id).to(inputs_embeds_p),
+                    ),
+                    dim=1,
+                )
+                labels_p = torch.cat((labels_p, torch.full((batch_size, difference), IGNORE_INDEX).to(labels_p)), dim=1)
+                attention_mask_p = torch.cat(
+                    (
+                        attention_mask_p,
+                        torch.zeros((batch_size, difference), dtype=torch.bool).to(attention_mask_p),
+                    ),
+                    dim=1,
+                )
+                position_ids_p = torch.cat(
+                    (position_ids_p, torch.full((batch_size, difference), -1).to(position_ids_p)), dim=1
+                )
+
+        return inputs_embeds_p, attention_mask_p, position_ids_p, labels_p
+
+    def get_xgr_logits_processor(self, response_format) -> List[LogitsProcessor]:
+        raise NotImplementedError("This method is not implemented for VILA model.")
+        # Convert response format to logits processor
+        import xgrammar as xgr
+
+        logging.info("[XGrammar] Compiling grammar for contrained output")
+
+        if self.grammar_compiler is None:
+            # logging.info(f"[XGrammar] {self.tokenizer}, {self.tokenizer.vocab_size}, {self.vocab_size}")
+            self.grammar_compiler = xgr.GrammarCompiler(
+                xgr.TokenizerInfo.from_huggingface(self.tokenizer, vocab_size=self.vocab_size)
+            )
+
+        if response_format.type == "json_schema":
+            compiled_grammar = self.grammar_compiler.compile_json_schema(
+                response_format.json_schema.schema_,
+                indent=2,
+            )
+        else:
+            compiled_grammar = self.grammar_compiler.compile_builtin_json_grammar()
+
+        return [xgr.contrib.hf.LogitsProcessor(compiled_grammar)]
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        media: Optional[Dict[str, List[torch.Tensor]]] = None,
+        images: Optional[torch.FloatTensor] = None,
+        media_config: Optional[List] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        packing: bool = True,
+        force_packing: bool = False,
+        seqlens_in_batch: Optional[torch.LongTensor] = None,
+        dpo_forward: bool = False,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        self.freezed_module_patch()
+
+        if images is not None:
+            if media is not None:
+                raise ValueError("Both 'media' and 'images' are provided. Please provide only one.")
+            print("The 'images' argument is deprecated. Please use 'media' instead.")
+            media = {"image": images}
+
+        if media_config is None:
+            media_config = defaultdict(dict)
+
+        if inputs_embeds is None:
+            inputs_embeds, labels, attention_mask = self._embed(input_ids, media, media_config, labels, attention_mask)
+
+        if force_packing or (packing and self.training and not dpo_forward):
+            if seqlens_in_batch is None:
+                seqlens_in_batch = torch.sum(attention_mask, dim=1)
+            set_seqlens_in_batch(seqlens_in_batch)
+
+            (inputs_embeds, attention_mask, position_ids, labels) = self.repack_multimodal_data(
+                inputs_embeds, attention_mask, position_ids, labels
+            )
+
+        outputs = self.llm(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            labels=labels,
+            **kwargs,
+        )
+
+        if self.training and getattr(self.config, "time_token_ids", []):
+            outputs.loss = soft_cross_entropy(
+                outputs.logits,
+                labels,
+                soft_tokens=self.config.time_token_ids,
+                std=self.config.soft_ce_std,
+            )
+
+        if dpo_forward:
+            return outputs.logits, labels
+
+        return outputs
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: Optional[torch.FloatTensor] = None,
+        media: Optional[Dict[str, List[torch.Tensor]]] = None,
+        media_config: Dict[str, Dict[str, Any]] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        **generation_kwargs,
+    ):
+        if self.training:
+            warnings.warn(
+                "Model is in training mode, using default padding strategy to right. This is not recommended for generation."
+            )
+        inputs_embeds, _, attention_mask = self._embed(input_ids, media, media_config, None, attention_mask)
+        return self.llm.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, **generation_kwargs)
+
+    @torch.inference_mode()
+    def generate_content(
+        self,
+        prompt: Union[str, List],
+        generation_config: Optional[GenerationConfig] = None,
+        response_format=None,
+    ) -> str:
+        # TODO(zhijianl): Support directly taking conversation as input
+        conversation = [{"from": "human", "value": prompt}]
+
+        # Convert response format to logits processor
+        if response_format:
+            xgr_logits_processor = self.get_xgr_logits_processor(response_format)
+        else:
+            xgr_logits_processor = None
+
+        # Extract media from the conversation
+
+        # TODO (extract and preprocess should be done together, as the preprocess of image and video can be different, i.e. when dynamic res is used)
+        media = extract_media(conversation, self.config)
+
+        # Process media
+        media_config = defaultdict(dict)
+        for name in media:
+            if name == "image":
+                if len(media["image"]) == 1 and self.config.image_aspect_ratio in ["dynamic", "dynamic_s2"]:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if self.config.image_aspect_ratio == "dynamic":
+                        images = process_image(media["image"][0], self.config, None, enable_dynamic_res=True).half()
+                        conversation[0]["value"] = conversation[0]["value"].replace(
+                            DEFAULT_IMAGE_TOKEN, f"{DEFAULT_IMAGE_TOKEN}\n" * images.shape[0]
+                        )
+                    else:
+                        if type(self.config.s2_scales) is str:
+                            self.config.s2_scales = list(map(int, self.config.s2_scales.split(",")))
+                        images, block_sizes = process_image(
+                            media["image"][0], self.config, None, enable_dynamic_s2=True
+                        )
+                        images = images.half()
+                        media_config[name]["block_sizes"] = [block_sizes]
+                else:
+                    images = process_images(media["image"], self.vision_tower.image_processor, self.config).half()
+                media[name] = [image for image in images]
+            elif name == "video":
+                if self.config.image_aspect_ratio == "dynamic" and self.config.video_max_tiles > 1:
+                    media[name] = [
+                        process_images(
+                            images,
+                            self.vision_tower.image_processor,
+                            self.config,
+                            enable_dynamic_res=True,
+                            max_tiles=self.config.video_max_tiles,
+                        ).half()
+                        for images in media[name]
+                    ]
+                elif self.config.image_aspect_ratio == "dynamic_s2" and self.config.video_max_tiles > 1:
+                    self.config.image_processor = self.vision_tower.image_processor
+                    if type(self.config.s2_scales) is str:
+                        self.config.s2_scales = list(map(int, self.config.s2_scales.split(",")))
+                    media[name] = [
+                        torch.cat(
+                            [
+                                process_image(
+                                    image,
+                                    self.config,
+                                    None,
+                                    enable_dynamic_s2=True,
+                                    max_tiles=self.config.video_max_tiles,
+                                )[0].half()
+                                for image in images
+                            ]
+                        )
+                        for images in media[name]
+                    ]
+                else:
+                    media[name] = [
+                        process_images(images, self.vision_tower.image_processor, self.config).half()
+                        for images in media[name]
+                    ]
+            else:
+                raise ValueError(f"Unsupported media type: {name}")
+
+        # Tokenize the conversation
+        input_ids = tokenize_conversation(conversation, self.tokenizer, add_generation_prompt=True).cuda().unsqueeze(0)
+
+        # Set up the generation config
+        generation_config = generation_config or self.default_generation_config
+
+        # print("input_ids", input_ids.shape)
+        # print(input_ids)
+        # print(self.tokenizer.batch_decode(input_ids))
+        # print("media", {k: len(v) for k, v in media.items()})
+        # print("media_config", media_config)
+        # print("generation_config", generation_config)
+        # input("wait for debug")
+        # Generate the response
+        try:
+            output_ids = self.generate(
+                input_ids=input_ids,
+                media=media,
+                media_config=media_config,
+                generation_config=generation_config,
+                logits_processor=xgr_logits_processor,  # structured generation
+            )
+        except ValueError:
+            if not generation_config.do_sample:
+                raise
+            # FIXME(zhijianl): This is a temporary workaround for the sampling issue
+            logging.warning("Generation failed with sampling, retrying with greedy decoding.")
+            generation_config.do_sample = False
+            output_ids = self.generate(
+                input_ids=input_ids,
+                media=media,
+                media_config=media_config,
+                generation_config=generation_config,
+                logits_processor=xgr_logits_processor,
+            )
+
+        # Decode the response
+        response = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+        return response
+
+    @property
+    def default_generation_config(self) -> GenerationConfig:
+        generation_config = copy.deepcopy(self.generation_config or GenerationConfig())
+        if self.tokenizer.eos_token_id is None:
+            raise ValueError("Tokenizer must have an EOS token")
+        if generation_config.max_length == GenerationConfig().max_length:
+            generation_config.max_length = self.tokenizer.model_max_length
+        if generation_config.pad_token_id is None:
+            generation_config.pad_token_id = self.tokenizer.pad_token_id or self.tokenizer.eos_token_id
+        if generation_config.bos_token_id is None:
+            generation_config.bos_token_id = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id
+        if generation_config.eos_token_id is None:
+            generation_config.eos_token_id = self.tokenizer.eos_token_id
+        return generation_config

prediction.jsonl

@@ -0,0 +1,103 @@
+{"question_id": "0", "text": "このイラストは、著名な映画監督である宮崎駿の作品であると考えられます。彼の作品は、独特のキャラクターデザインや細部へのこだわりが特徴です。このイラストには、手に持った羽根や、周囲に散らばった本、そして作業に集中している様子が描かれており、宮崎駿の作品に見られるような幻想的で創造的な要素が感じられます。また、彼の作品はしばしば自然や動物をテーマにしているため、羽根や本といったアイテムも彼のスタイルに合致しています。このように、イラストの特徴やテーマから、宮崎駿の作品である可能性が高いと推測できます。", "answer": "この作品は宮崎駿監督によるものだと考えられます。画像の説明から、細部への注意と物語の中での小さなアイテムの重要性を示すスタイルが宮崎駿の作品に特徴的です。また、「君たちはどう生きるか」というタイトルも、宮崎駿の作品のテーマに合致しています。", "input_text": "この作品は著名な映画監督により制作されました。イラストから誰の作品だと考えられますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "1", "text": "映像に写っている少年は、羽根を切ったり、羽根の先を削ったりしています。", "answer": "映像に写っている少年は、机に置かれた青色の鳥の羽のうちの一つを小型ナイフで加工しています。机の左側には竹の棒、赤いワイヤー、2本の鉛筆が置かれており、右側には数冊の本といくつかの書類が配置されています。", "input_text": "映像に写っている少年は何をしていますか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "2", "text": "鳥の羽の色は青色です。", "answer": "鳥の羽の色は青色です。", "input_text": "鳥の羽の色は何色でしょうか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "3", "text": "机の上には2本の鉛筆が置かれています。", "answer": "机の上には2本の鉛筆が置かれています。これらは机の左側に配置されています。", "input_text": "机の上には何本の鉛筆が置かれていますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "4", "text": "このイラストには、男性がテーブルの上で作業をしている様子が描かれています。彼は手にナイフを持ち、青い羽を切っています。テーブルの上には、いくつかの本が散らばっており、作業のための資料として使われているようです。\n\n男性の周りには、羽の他にもいくつかの羽が見え、彼がそれらを整理している様子が伺えます。テーブルの一部には、羽を切るための道具や、作業に必要なアイテムが置かれているようです。全体的に、彼は集中して作業に取り組んでいる印象を与えます。", "answer": "このイラストでは、少年が机に座り、一枚の青い鳥の羽を小型ナイフで加工している様子が描かれています。机の左側には竹の棒、赤いワイヤー、2本の鉛筆が配置されており、右側には数冊の本といくつかの書類が整然と置かれています。また、机の上にはもう一枚の青い鳥の羽が置かれています。", "input_text": "このイラストについて詳しく説明してください。", "heron-bench": {"score": 8, "score_gpt": 10}}
+{"question_id": "5", "text": "このイラストは、アニメーション映画『千と千尋の神隠し』のシーンを描いています。物語は、主人公の千尋が不思議な世界に迷い込むところから始まります。千尋は、両親と共に車で移動中に、突然のトンネルを通り抜け、異世界にたどり着きます。この世界では、両親が豚に変身してしまうという驚くべき出来事が起こります。\n\n千尋は、この異世界で働くために、湯婆婆の経営する「油屋」という温泉施設に雇われます。そこで、彼女はさまざまなキャラクターと出会い、成長していきます。このシーンでは、千尋が湯婆婆や他のキャラクターたちと一緒に食事をしている様子が描かれています。テーブルにはケーキやカップが並び、和やかな雰囲気が漂っています。この場面は、千尋が新しい環境に適応し、仲間たちとの絆を深める重要な瞬間を象徴しています。", "answer": "このイラストは、「千と千尋の神隠し」という映画の一場面を描いています。この作品では、主人公の少女・千尋が不思議な世界に迷い込み、両親を豚に変えられた後、彼らを元に戻すために奮闘します。千尋はこの世界で様々な個性的なキャラクターに出会い、成長していく物語です。", "input_text": "このイラストが示す作品のあらすじを詳細に教えてください。", "heron-bench": {"score": 10, "score_gpt": 8}}
+{"question_id": "6", "text": "この映像の中で明らかに人間ではないのは、白いマスクをかぶったキャラクターです。彼は他のキャラクターたちと一緒にテーブルに座っていますが、その外見から人間ではないことがわかります。", "answer": "この映像の中で明らかに人間ではないのはカオナシです。カオナシは化け物であり、千尋の左隣に位置しています。また、湯婆婆も人間とは異なる特徴を持つキャラクターですが、カオナシが最も人間ではない外見をしています。", "input_text": "この映像の中で明らかに人間ではないのはどれでしょうか？", "heron-bench": {"score": 8, "score_gpt": 9}}
+{"question_id": "7", "text": "このイラストの少女の名前は「千尋」です。", "answer": "イラストの少女の名前は千尋です。", "input_text": "イラストの少女の名前はなんですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "8", "text": "いいえ、一番左にいるのは人間ではなく、白い仮面をかぶったキャラクターです。", "answer": "一番左にいるのは人間ではありません。それはカオナシという化け物です。", "input_text": "一番左にいるのは人間でしょうか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "9", "text": "このイラストには、テーブルを囲んで座っている3人のキャラクターが描かれています。テーブルの上には、ケーキやカップ、フォークが並んでおり、食事を楽しんでいる様子が伺えます。テーブルの中央には、大きなケーキがあり、その周りにはいくつかのカップが配置されています。\n\nキャラクターの一人は、テーブルの左側に座っており、もう一人は右側にいます。さらに、テーブルの奥には、もう一人のキャラクターが見えます。テーブルの上には、フォークが2本あり、それぞれのカップの近くに置かれています。また、テーブルの上には花瓶も見え、部屋の雰囲気を和らげています。全体的に、和やかな食事のシーンが表現されています。", "answer": "このイラストでは、中心に千尋が大きなケーキ、紅茶、小菓子が並ぶテーブルに座っており、彼女の左隣には無表情のカオナシが立っています。千尋の右側には、特徴的な大きな鼻と目、そして独特な髪型を持つ老婦人の湯婆婆が微笑みながら座っています。画像の背景には、食器が並べられた食器棚が配置されています。", "input_text": "このイラストについて詳しく説明してください。", "heron-bench": {"score": 7, "score_gpt": 10}}
+{"question_id": "10", "text": "この作品は、水墨画の技法を用いて描かれており、その独特な表現力が評価されています。水墨画は、墨の濃淡や筆の動きによって、自然の風景や物体の質感を表現する技法です。この作品では、山や木々、建物が描かれており、それぞれの要素が墨の濃淡によって強調されています。特に、山の輪郭や木の枝の細部が、筆のタッチによって生き生きと表現されており、観る者に深い印象を与えます。\n\nまた、水墨画はしばしばシンプルな構図と余白を重視し、その中で物語や感情を伝えることが求められます。この作品も、余白を巧みに使い、自然の美しさと静けさを強調しています。さらに、作品のタイトルや署名が含まれていることで、作者の意図や背景が感じられ、鑑賞者に深い理解を促します。このように、技法と構成が相まって、この作品は高い評価を得ているのです。", "answer": "この作品は、雪舟の独特な技法と空間構成によって高く評価されています。彼の強調された輪郭線と細い線による簡略化された表現は、日本の山岳風景を独自の視点で捉え、深い情感と静謐な美を伝えています。また、構築的な空間構成は、見る者をその風景の中に引き込み、自然との一体感を感じさせる効果を持っています。", "input_text": "この作品がなぜ評価されているのか、技法に触れて論じてください。", "heron-bench": {"score": 10, "score_gpt": 9}}
+{"question_id": "11", "text": "作品に写っている人は、山の斜面を登っており、山の頂上に向かっています。", "answer": "作品に写っている人は、画像の下方にある坂道から家に向かって登っていま��。この人物は右下に位置する川の近くから出発し、左の中央にある家を目指している様子が描かれています。", "input_text": "作品に写っている人はどこからどこに向かっていますか？", "heron-bench": {"score": 6, "score_gpt": 9}}
+{"question_id": "12", "text": "この水墨画の作品名は「雪峰禅寺」です。", "answer": "この水墨画の作品名は「秋冬山水図」です。雪舟筆による日本の山岳風景を描いた作品で、構築的な空間構成と強調された輪郭線、細い線による簡略化された技法が特徴です。", "input_text": "この水墨画の作品名はなんですか？", "heron-bench": {"score": 1, "score_gpt": 10}}
+{"question_id": "13", "text": "「雪峰禅」と書かれています。", "answer": "左の真ん中あたりには家が描かれています。この家は山岳風景の中に位置しており、周囲は自然に囲まれています。家の構造は簡略化された線で表現されており、細かいディテールは省略されています。", "input_text": "左の真ん中あたりには何が書かれていますか？", "heron-bench": {"score": 1, "score_gpt": 10}}
+{"question_id": "14", "text": "この水墨画は、山々とその間に広がる谷を描いた風景画です。画面の中央には、山の斜面に沿って建物が見え、その周囲には木々が点在しています。山の頂上には雪が積もっており、冬の厳しさを感じさせます。\n\n画面の左側には、山の斜面が広がり、その下には谷が見えます。右側には、さらに高い山がそびえ立ち、その山肌には木々が生い茂っています。全体的に、水墨画特有の淡い色合いと筆のタッチが、自然の美しさを引き立てています。この作品は、静寂で荘厳な雰囲気を醸し出しており、自然の力強さと美しさを表現しています。", "answer": "この水墨画は、秋冬の山岳風景を描いたもので、左中央には小さな家が、右側にはいくつかの木々が配置されています。画面の右下部には川が流れ、その川沿いの坂道を一人の人物が家に向かって登っている様子が描かれています。この作品は、構築的な空間構成と強調された輪郭線、細い線による簡略化された技法が特徴です。", "input_text": "この水墨画について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "15", "text": "この作品では、白い妖怪と緑の妖怪が対比されています。白い妖怪は、円形の輪の中を跳び跳ねている姿が描かれており、その動きからは軽やかさや楽しさが感じられます。一方、緑の妖怪は、より力強く、力強い印象を与えています。この対比は、色の違いだけでなく、動きや表情にも表れており、白い妖怪の軽快さと緑の妖怪の力強さが対照的に描かれています。このような対比は、視覚的な興味を引き、観る者に異なる感情や印象を与える効果があります。", "answer": "この作品では、風を操る風神と雷を起こす雷神を対比させています。風神は右側に描かれ、風を吹き出しているのに対し、雷神は左側に描かれ、太鼓を打って雷を起こしています。これにより、自然界の二つの強大な力、風と雷をユーモラスに表現しています。", "input_text": "この作品では、何と何を対比させていますか？", "heron-bench": {"score": 8, "score_gpt": 10}}
+{"question_id": "16", "text": "この作品は、日本の伝統的な妖怪や神話のキャラクターを描いたもので、視覚的に非常に魅力的です。二つの異なるキャラクターが描かれており、それぞれが異なる動きや表情を持っています。このような作品は、日本の文化や歴史、そしてその背後にあるストーリーを伝える重要な役割を果たしています。また、妖怪や神話のキャラクターは、時には恐怖や神秘を象徴し、観る者に想像力をかき立てる効果があります。この作品を通じて、私たちは日本の伝統や美意識、そして物語の力について考えさせられるでしょう。さらに、これらのキャラクターは、現代のポップカルチャーやアニメにも影響を与えているため、その文化的な影響力も無視できません。", "answer": "この作品は、自然の力を人間がコントロールできないものとして、しかし同時にその力をユーモラスに捉えることで、恐怖や畏敬の念を和らげる方法を示唆しています。風神と雷神の表現を通じて、自然現象への敬意と、それに対する人間の創造的な対応を伝えています。また、風と雷という自然の力を擬人化することで、それらの力が持つ生命力や動きを視覚的に表現しています。", "input_text": "この作品は私達に何を伝えていますか？", "heron-bench": {"score": 8, "score_gpt": 10}}
+{"question_id": "17", "text": "この屏風は、江戸時代の浮世絵師である歌川国芳によって作成されました。", "answer": "この屏風は江戸時代中期を代表する画家、尾形光琳によって作成されました。", "input_text": "この屏風は誰によって作成されましたか？", "heron-bench": {"score": 1, "score_gpt": 10}}
+{"question_id": "18", "text": "この作品は江戸時代のもので、18世紀後半から19世紀初頭にかけて描かれたと考えられます。", "answer": "この作品は江戸時代中期の作品です。尾形光琳によって描かれた風神雷神図屛風は、その時代を代表する芸術作品の一つとされています。", "input_text": "これはいつの時代の作品でしょうか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "19", "text": "この屏風絵は、日本の伝統的な美術様式で描かれたもので、2つの異なる妖怪が描かれています。左側には、白い毛皮を持つ妖怪が、右側には緑色の肌を持つ妖怪が描かれています。両方の妖怪は、それぞれ異なるポーズをとっており、動きのある表現がされています。\n\n左側の妖怪は、まるで何かを捕まえようとしているかのように、手を伸ばしている姿が印象的です。一方、右側の妖怪は、より力強い姿勢で、まるで何かを追いかけているかのようです。背景には、自然の要素が描かれており、妖怪たちの存在感を引き立てています。この屏風絵は、日本の妖怪文化や伝説を表現しており、見る者に強い印象を与えます。", "answer": "この屏風絵は、尾形光琳による風神雷神図屛風で、二枚一対の屏風絵です。右側の屏風には、力強く風を吹き出している風神が描かれており、左側の屏風には、太鼓を激しく打ち鳴らして雷を起こしている雷神が描かれています。両者は動きが大きく、表情もユーモラスで、自然の力を生き生きと表現しています。", "input_text": "この屏風絵について詳しく説明してください。", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "20", "text": "この作品は、動物たちが楽しそうに走り回る様子を描いたユーモラスな絵です。特に、カエルやウサギ、そして人間が一緒に競争しているシーンが描かれており、観る者に笑いを誘います。絵の中では、カエルがウサギを追いかけている姿が印象的で、彼らの動きからは生き生きとしたエネルギーが感じられます。また、背景には山々が広がり、自然の中での遊び心満載のシーンが展開されています。この作品は、動物たちの無邪気さや、競争の楽しさを表現しており、見る人に温かい気持ちをもたらします。", "answer": "この絵巻物は、京都の高山寺に伝わる鳥獣人物戯画の一部で、猿、兎、カエルが描かれた紙本墨画です。猿は首に麦わら帽子をかけ、一匹の兎と二匹のカエルに追いかけられる様子が描かれていますが、もう一匹の兎は猿とは反対方向を向いています。この作品は、動物たちを通じて人間社会の風刺を描いたものと解釈されており、平安時代のユーモアと洞察が感じられます。", "input_text": "あなたは美術館の学芸員です。この作品にどのような説明文をつけて展示しますか？", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "21", "text": "この作品には、カエルが3匹います。", "answer": "この作品の中にはカエルが二匹います。", "input_text": "この作品のなかに、カエルは何匹いますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "22", "text": "この作品の中には、兎が2匹います。", "answer": "この作品の中には兎が二匹います。一匹は猿を追いかけており、もう一匹は猿とは反対側に向かっています。", "input_text": "この作品の中に兎は何匹いますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "23", "text": "この墨画は、動物たちが丘を駆け下りる様子を描いています。画面には、さまざまな動物が描かれており、特にカエルやウサギが目立っています。カエルは丘の斜面を滑り降りる姿が描かれており、その動きからは勢いを感じることができます。\n\nウサギは、丘の上から下へと走り降りる姿が描かれており、その姿勢からは楽しそうな様子が伝わってきます。また、他の動物たちも丘を駆け下りており、全体的に活気に満ちたシーンが展開されています。この墨画は、動物たちの動きや表情を巧みに捉えており、自然の中での生き生きとした瞬間を表現しています。", "answer": "この墨画は、猿が麦わら帽子を首にかけており、一匹の兎と二匹のカエルに追いかけられている様子を描いています。猿は画像の右側に位置し、追いかける兎とカエルはその左側にいます。また、もう一匹の兎は猿とは反対方向、つまり画像の左側を向いています。", "input_text": "この墨画について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 9}}
+{"question_id": "24", "text": "この作品は、現在の東京都を対象に描かれたと考えられます。理由は、作品に描かれている風景や建物の特徴から、江戸時代の東京の様子を表現していると推測されるからです。特に、江戸時代の東京は「江戸」として知られ、多くの人々が集まり、商業や文化が栄えた場所でした。また、作品には多くの人々や舟が描かれており、当時の賑わいを示しています。さらに、背景に描かれた山は、富士山を連想させるものであり、江戸時代の東京の風景を象徴する要素の一つです。これらの要素から、この作品は東京都を対象に描かれたものであると推測できます。", "answer": "この作品は現在の東京都を対象に描かれています。理由は、画中に江戸日本橋が描かれており、江戸城の存在が示されているためです。江戸は現在の東京都にあたります。", "input_text": "この作品は現在のどの都道府県を対象に描かれたのか教えてください。また、なぜそのように考えたのか教えてください。", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "25", "text": "町並みの先に写っているお城の名前は「富士山」です。お城の向こう側には、富士山が写っています。", "answer": "町並みの先に写っているお城の名前は江戸城です。江戸城の向こう側には富士山が写っています。", "input_text": "町並みの先に写っているお城の名前は何でしょうか？また、お城の向こう側には何が写っていますか？", "heron-bench": {"score": 1, "score_gpt": 10}}
+{"question_id": "26", "text": "この作品は、江戸時代の日本、特に富士山が見える場所からの景色を描いています。", "answer": "この作品は江戸の日本橋からの景色を描いています。画中では、日本橋の賑わいとその上を行き交う人々が描かれており、背景には江戸城とその向こうに霞む富士山が見えます。", "input_text": "この作品はどこからの景色を描いていますか？", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "27", "text": "この作品の作者は歌川広重です。", "answer": "この作品の作者は葛飾北斎です。彼は「冨嶽三十六景」シリーズの中で《江戸日本橋》を描いており、この作品はその一部です。北斎は江戸時代後期の日本の浮世絵師として広く知られています。", "input_text": "この作品の作者は誰ですか？", "heron-bench": {"score": 1, "score_gpt": 10}}
+{"question_id": "28", "text": "この絵画は、江戸時代の日本の風景を描いたもので、川沿いの町並みが広がっています。川には多くのボートが浮かんでおり、その中には人々が乗っている様子が見受けられます。ボートは川の両側に並んでおり、賑やかな雰囲気を醸し出しています。\n\n町の背景には、伝統的な日本の建物が立ち並び、その中には富士山が遠くに見える風景が描かれています。富士山は日本の象徴的な山であり、この絵画の重要な要素となっています。また、町の中心には橋があり、その上を人々が行き交っている様子が描かれています。\n\n絵画全体は、色鮮やかで、特に青い空と川の水面が印象的です。人々は様々な服装をしており、当時の生活様式や文化を感じさせる要素が豊富に含まれています。この作品は、日本の歴史や文化を理解するための貴重な資料となっています。", "answer": "この絵画は、葛飾北斎の「冨嶽三十六景」シリーズの一部で、「江戸日本橋」を描いています。画面下部には、擬宝珠を中心に賑わう人々が描かれ、その背後には大胆な透視法を用いて描かれた江戸の街並みが広がり、さらにその向こうには江戸城が見えます。画面の最も遠い背景には、霞むような表現で富士山が描かれており、日本の自然と都市生活の調和を象徴しています。", "input_text": "この絵画に���いて詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "29", "text": "こたつと言えば、みかんが思い浮かびます。", "answer": "こたつと言えば、みかんが思い浮かびます。画像には、赤いこたつの上にみかんが盛られた皿が置かれています。この光景は、冬の暖かいひとときを象徴しています。", "input_text": "こたつと言えば、何の果物が思い浮かびますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "30", "text": "この部屋は日本のスタイルで、畳の床と布団が特徴的です。", "answer": "写真に写っている部屋は日本のスタイルです。畳の和室、こたつ、そしてみかんが盛られた皿は日本の伝統的な冬の風景を表しています。また、部屋の収納も日本の和室によく見られるデザインです。", "input_text": "写真に写っている部屋はどこの国のスタイルでしょうか？", "heron-bench": {"score": 8, "score_gpt": 10}}
+{"question_id": "31", "text": "こたつです。", "answer": "中央にあるテーブルの名称は「こたつ」です。", "input_text": "中央にあるテーブルの名称はなんですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "32", "text": "机の上にはオレンジが乗っています。", "answer": "机に乗っている果物はみかんです。みかんは皿に盛られており、こたつの上に置かれています。", "input_text": "机に乗っている果物は何ですか？", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "33", "text": "この画像には、和室の床に敷かれた布団の上に小さなテーブルが置かれています。テーブルの上には、オレンジが入ったボウルが2つあり、色とりどりのオレンジが並んでいます。また、テーブルの近くには、いくつかのリンゴも見えます。\n\n布団は部屋の中央に広がっており、その上にはテーブルがしっかりと置かれています。テーブルの周りには、オレンジやリンゴが散らばっており、和室の落ち着いた雰囲気を引き立てています。全体的に、リラックスした空間が感じられるシーンです。", "answer": "この写真は、和室の中心に設置された赤いこたつを捉えています。こたつの上には、みかんが数個盛られた皿と、いくつかの小菓子が置かれており、和室の雰囲気に温かみを加えています。部屋の側面には、扉付きの収納が見え、部屋の整理整頓に役立っている様子が伺えます。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "34", "text": "旅客機に描かれているキャラクターの中で最も有名なのは、ピカチュウです。", "answer": "旅客機に写っている最も有名なキャラクターの名前はピカチュウです。", "input_text": "旅客機に写っているキャラクターで最も有名なキャラクターの名前を教えてください。", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "35", "text": "ジェットエンジンにはポケモンの絵が描かれています。", "answer": "ジェットエンジンにはポケモンボールの絵が描かれています。", "input_text": "ジェットエンジンには何の絵が描かれていますか？", "heron-bench": {"score": 8, "score_gpt": 10}}
+{"question_id": "36", "text": "写真には2機の飛行機が写っています。", "answer": "写真には2機の飛行機が写っています。", "input_text": "写真には何機の飛行機が写っていますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "37", "text": "写真の奥の方には、海が広がっています。", "answer": "写真の奥の方には、海があり白波が立っています。", "input_text": "写真の奥の方には何が写っていますか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "38", "text": "画像には、空港の滑走路に停まっている2機のカラフルな飛行機が写っています。これらの飛行機は、アニメのキャラクターで装飾されており、非常に目を引くデザインです。1機は左側に位置し、もう1機はその右側にあります。\n\n滑走路の背景には、海が広がっており、飛行機の鮮やかな色合いが海の青さと対比を成しています。また、滑走路の近くには数隻のボートが見え、海の風景をさらに引き立てています。全体的に、このシーンは楽しい雰囲気を醸し出しており、旅行や冒険の期待感を感じさせます。", "answer": "この写真には、ポケモンのキャラクターがデザインされた2機のANAの旅客機が飛行場に停まっています。一方のジェット機の先端にはピカチュウが描かれており、ジェットエンジンにはポケモンボールのデザインが施されています。背景には海が広がり、白波が見えます。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 8, "score_gpt": 9}}
+{"question_id": "39", "text": "力士たちは土俵の上で相撲の試合を行っています。彼らは互いに競い合い、力と技術を駆使して勝利を目指しています。", "answer": "力士たちは土俵で相撲の儀式を行っています。中央の横綱は土俵入りの型を演じており、これは相撲の伝統的な儀式の一つです。他の二人の力士はおそらく横綱の土俵入りを見守っているか、次の取組の準備をしている状態です。", "input_text": "力士たちは何を行っているか説明してください。", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "40", "text": "水引幕から垂れている房の色は赤色です。", "answer": "水引幕から垂れている房の色は赤、緑、白、黒の4色です。", "input_text": "水引幕から垂れている房の色は何色ですか？", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "41", "text": "土俵には2人の力士がいます。", "answer": "土俵には3人の力士がいます。中央で横綱が土俵入りの型を演じており、他の2人はおそらく彼の土俵入りを見守っています。", "input_text": "土俵には何人の力士がいますか？", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "42", "text": "観客席は非常に混雑しており、多くの人々が座っています。", "answer": "観客席は満員で、数百人の観客が詰めかけています。多くの人がスマートフォンで撮影をしている様子が見受けられます。観客は土俵を囲むように座っており、全席が埋まっている状態です。", "input_text": "観客席に人はどのくらい入っていますか？", "heron-bench": {"score": 8, "score_gpt": 9}}
+{"question_id": "43", "text": "この画像には、日本の伝統的な相撲の試合が行われている様子が写っています。土俵の上で2人の力士が激しく戦っており、その周りには多くの観客が集まっています。観客は試合を見守りながら、興奮した様子で応援しています。\n\n土俵の周囲には、さまざまな人々が立っており、彼らの表情からは試合の緊張感が伝わってきます。観客の中には、携帯電話で写真を撮っている人も見受けられます。また、観客の中にはネクタイを締めた人もおり、フォーマルな服装で試合を楽しんでいる様子が伺えます。\n\n全体的に、この場面は日本の文化と伝統が色濃く反映されており、力士たちの力強い動きと観客の熱気が融合した、非常に活気のある雰囲気が感じられます。", "answer": "この写真は、相撲場所の土俵で撮影されたもので、中央には横綱が土俵入りの型を演じている様子が写っています。土俵の周りには他の2人の力士がおり、背景には水引幕から垂れる赤、緑、白、黒の4本の房が見えます。観客席は満員で、多くの観客がスマートフォンでこの瞬間を撮影している様子が確認できます。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 5, "score_gpt": 10}}
+{"question_id": "44", "text": "この写真に写っているのは日本の料理です。", "answer": "写真に写っているのは日本の料理です。天ぷらは日本の伝統的な揚げ物料理で、野菜や海鮮を衣で包んで揚げたものです。この写真には、特にエビの天ぷらといくつかの野菜の天ぷらが盛り付けられており、日本料理の一例として認識されます。", "input_text": "写真に写っているのはどこの国の料理が教えてください。", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "45", "text": "野菜の天ぷらはエビの天ぷらの隣に配置されています。", "answer": "野菜の天ぷらはエビの天ぷらの手前に置かれています。左からパプリカ、茄子、シシトウの順に並んでおり、その後ろに2尾のエビの天ぷらが配置されています。この配置により、野菜とエビの天ぷらが一緒に綺麗に盛り付けられている様子が見て取れます。", "input_text": "野菜の天ぷらとエビの天ぷらの位置関係はどのようになっていますか？", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "46", "text": "エビの天ぷらは2つあります。", "answer": "エビの天ぷらは2尾あります。", "input_text": "エビの天ぷらはいくつありますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "47", "text": "野菜の天ぷらは2つあります。", "answer": "野菜の天ぷらは3つあります。左からパプリカ、茄子、シシトウの順に盛り付けられています。", "input_text": "野菜の天ぷらはいくつありますか？", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "48", "text": "この画像には、白い紙の上に盛り付けられた天ぷらが写っています。天ぷらには、エビ、ナス、ピーマン、そして赤いトマトが含まれています。エビは特に目を引く存在で、他の食材とともにカラフルな盛り付けをしています。\n\nテーブルの上には、天ぷらの他にもカップが置かれており、飲み物が入っているようです。また、テーブルの一部には椅子も見え、食事を楽しむための準備が整っている様子が伺えます。全体的に、食欲をそそる美味しそうな料理が並んでいます。", "answer": "この写真は、美しく盛り付けられた天ぷらの一皿を捉えています。前方にはパプリカ、茄子、シシトウの野菜天ぷらが左から順に並び、その奥には2尾のエビの天ぷらが置かれています。写真の右奥には、グラスが1つ配置されており、全体的に料理の魅力を引き立てるシンプルながらも洗練された演出がなされています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "49", "text": "この料理は、肉や野菜がたっぷりと盛り付けられたボウルで、特に鶏肉やネギが目立っています。このような料理には、醤油やごま油、またはラー油などの調味料が非常に合います。醤油は、料理の風味を引き立て、塩味を加えるのに最適です。また、ごま油は香ばしさを加え、料理全体に豊かな風味をもたらします。さらに、ラー油を少し加えることで、ピリッとした辛味が加わり、食欲をそそる一品に仕上がります。これらの調味料を組み合わせることで、料理の味わいが一層引き立ち、より美味しく楽しむことができるでしょう。", "answer": "この料理にはラー油や酢がぴったりです。ラー油は辛味を加え、酢は麺にさっぱりとした風味をもたらします。これらの調味料は油そばの濃厚な味わいを引き立て、バランスの良い食感を提供します。", "input_text": "この料理を食べるときにぴったりの調味料を教えてください。", "heron-bench": {"score": 7, "score_gpt": 10}}
+{"question_id": "50", "text": "このラーメンは、一般的なラーメンとは異なる特徴を持っています。まず、ラーメンの上には、ネギや肉、卵などのトッピングが豊富に盛り付けられています。特に、ネギがたっぷりと使われているため、見た目にも鮮やかで食欲をそそります。また、ラーメンのスープは、通常のラーメンとは異なる風味を持っている可能性があります。さらに、ラーメンの器は黒いボウルで、全体的に高級感や特別感を演出しています。このような特徴から、このラーメンは一般的なラーメンよりも豪華で、特別な場面や特別なメニューとして提供されることが考えられます。", "answer": "油そばは普通のラーメンと異なり、スープがなく、特製の油と醤油ベースのソースで味付けされています。また、トッピングには温泉卵、角切りチャーシュー、刻みネギ、海苔、メンマが含まれ、ラー油や酢を好みで加えることが一般的です。これに対して、普通のラーメンはスープベースで提供され、具材や味付けにも違いがあります。", "input_text": "普通のラーメンとの違いを教えてください。", "heron-bench": {"score": 5, "score_gpt": 10}}
+{"question_id": "51", "text": "この写真の料理は、ラーメンの上に様々なトッピングが載ったもので、特にネギや肉が目立っています。", "answer": "この写真は東京油組の油そばの料理です。トッピングには温泉卵、角切りチャーシュー、刻みネギ、海苔、メンマがあり、ラー油や酢が調味料として添えられています。油そばはスープのない麺料理で、ラーメンに似ています。", "input_text": "この写真の料理はなんですか？", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "52", "text": "写真に写っている卵は半熟の状態で、麺の上に置かれています。", "answer": "写真に写っている卵は温泉卵です。黄身が半熟状態で、白身はやわらかく固まっています。この温泉卵は油そばの上に乗せられており、他のトッピングと一緒に中���に配置されています。", "input_text": "写真に写っている卵はどんな状態ですか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "53", "text": "画像には、テーブルの上に置かれた黒いボウルが写っています。このボウルには、麺と肉、そしてネギがたっぷりと盛り付けられています。ボウルの周りには、いくつかの調味料が置かれており、食事を楽しむ準備が整っています。\n\nテーブルの上には、ボウルの他にもカップが2つ見えます。1つはボウルの左側にあり、もう1つは右側に配置されています。また、ボウルの近くには、箸が置かれているのが見えます。全体的に、食事の準備が整った温かみのあるシーンが広がっています。", "answer": "この写真は、東京油組の油そばを映しており、麺の上には温泉卵、角切りチャーシュー、刻みネギ、海苔、メンマが美しくトッピングされています。麺は艶やかで、トッピングの配置がバランス良く、色彩豊かな見た目をしています。また、写真の端にはラー油と酢の小皿が置かれており、これらの調味料が油そばの味を引き立てることが予想されます。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 7, "score_gpt": 9}}
+{"question_id": "54", "text": "この写真に写っている彫像は、伝統的な日本の狛犬（こまいぬ）であり、通常は神社や寺院の入口に置かれています。狛犬は、邪悪な霊や悪運を防ぐ守護神としての役割を果たします。この彫像は、建物の前に立っており、訪れる人々に安心感を与え、神聖な場所を守る存在としての重要性を持っています。また、狛犬は通常、一対で置かれることが多く、それぞれが異なる表情や姿勢を持っていることが一般的です。この彫像は、その地域の文化や歴史を象徴するものでもあり、訪れる人々にその場所の特別な意味を伝える役割も果たしています。", "answer": "この写真に写っているシーサーの彫像は、琉球の伝統的な守り神としての役割を持っています。主に悪霊や災いを避け、家や村を守るために屋根の上や門の入り口に設置されます。沖縄県文化財指定の文字が示すように、このシーサーは文化的に重要な意味を持つアイテムであり、地域の伝統や信仰を象徴しています。", "input_text": "この写真に写っている彫像の役割について教えてください。", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "55", "text": "この写真は沖縄県で撮影されました。", "answer": "この写真は沖縄県で撮影されました。写真には「沖縄県文化財指定」という文字が見えるため、撮影場所が沖縄県であることが確認できます。また、シーサーは沖縄の伝統的な守り神であり、その存在も撮影地が沖縄県であることを支持しています。", "input_text": "この写真は何県で撮影されましたか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "56", "text": "写真に写っている獣像は「狛犬」です。", "answer": "写真に写っている獣像の名前は「シーサー」です。シーサーは琉球の伝統的な守り神で、魔除けの目的で屋根の上などに設置されます。", "input_text": "写真に写っている獣像の名前はなんですか？", "heron-bench": {"score": 1, "score_gpt": 10}}
+{"question_id": "57", "text": "画像にはシーサーが1体写っています。", "answer": "シーサーは1体写っています。", "input_text": "シーサーは何体写っていますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "58", "text": "この画像には、建物の前に立つ大きなシーサーの像が写っています。シーサーは沖縄の伝統的な守り神で、その特徴的な姿が目を引きます。像は建物の壁に寄りかかっており、その存在感が際立っています。\n\nシーサーの像の近くには、いくつかの鉢植えが配置されており、緑が建物の外観を引き立てています。鉢植えは、シーサーの像の左側と右側にそれぞれ置かれており、全体的に和やかな雰囲気を醸し出しています。また、建物の上部には電線が見え、現代的な要素も感じられます。", "answer": "この写真は、琉球の伝統的な守り神であるシーサーの彫像を特徴としています。彫像は、魔除けとして屋根の上に設置されている様子が写されており、背景には沖縄県文化財指定の文字が見える琉球の建物があります。また、左後方には緑豊かな木が生い茂っています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 9}}
+{"question_id": "59", "text": "この写真に写っている神社は、赤い鳥居とその背後に広がる山々が特徴的です。鳥居は水の上に立っており、その周囲には美しい自然が広がっています。この神社は、日本の伝統的な建築様式を反映しており、特に赤い色合いが目を引きます。また、鳥居の近くには木々が生い茂っており、自然との調和が感じられます。このような神社は、訪れる人々に静けさや神聖さを提供し、観光地としても人気があります。", "answer": "この写真に写っている厳島神社の特徴として、目立つのはその象徴的な赤い鳥居で、潮の満ち引きによって海上に浮かぶかのように見えます。背景には緑豊かな山々が広がり、その自然の美しさと調和しています。また、鳥居の左右には厳島神社の伝統的な建築様式を持つ建物が見え、その歴史的な雰囲気を一層際立たせています。", "input_text": "この写真に写っている神社の特徴について説明してください。", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "60", "text": "この写真に写っている鳥居は、通常の鳥居とは異なり、水の中に立っています。一般的に、鳥居は神社の入り口に設置され、地面にしっかりと固定されていますが、ここでは水の中に浮かんでいるように見えます。この特異な状況は、周囲の環境や季節によって引き起こされる可能性があります。例えば、洪水や高水位の影響で、鳥居が水に浸かっている状態が見られることがあります。このような状況は、地域の自然環境や気候条件を反映しており、訪れる人々にとって印象的な光景となるでしょう。また、この鳥居の存在は、神聖な場所であることを示す一方で、自然の力の強さや変化をも象徴しています。", "answer": "この写真に写っている鳥居の普通とは違うところは、水に浮かんでいる点です。通常、鳥居は地上に設置されますが、この鳥居は潮の満ち引きによって海上に現れる特徴があります。また、その背後に広がる緑豊かな山々とのコントラストが美しい景色を作り出しています。", "input_text": "この写真に写っている鳥居の普通とは違うところはどこですか？", "heron-bench": {"score": 8, "score_gpt": 9}}
+{"question_id": "61", "text": "この写真は厳島神社で撮影されました。厳島神社は広島県廿日市市宮島町に位置し、海に浮かぶ鳥居が特徴的です。", "answer": "この写真は日本の広島県にある厳島神社で撮影されました。", "input_text": "この写真はどこの神社で撮影されましたか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "62", "text": "写真の鳥居は水の中にあり、その背後には山々が広がっています。", "answer": "写真の鳥居は、緑豊かな山々を背景にして、厳島神社の前に位置しています。左右には厳島神社の建物が見え、山のふもとにはホテルがあります。この鳥居は水面に立っているように見えます。", "input_text": "写真の鳥居はどこにありますか？", "heron-bench": {"score": 10, "score_gpt": 9}}
+{"question_id": "63", "text": "この画像には、美しい湖の上に浮かぶ赤い鳥居が写っています。鳥居は日本の神社の入り口を示すもので、湖の水面に映り込んでいます。周囲にはいくつかの建物が見え、その中には神社の一部と思われる建物も含まれています。\n\n湖の近くには、いくつかのボートが停泊しており、その中には小型のボートがいくつか見受けられます。また、湖の対岸には山々が広がっており、自然の美しさが際立っています。全体的に、静かで穏やかな雰囲気が漂う風景です。", "answer": "この写真は、緑豊かな山々を背景に、その前に位置する厳島神社の象徴的な赤い鳥居を捉えています。鳥居の左右には、厳島神社の伝統的な建築様式を持つ建物が見え、山のふもとにはホテルが建っています。この構図は、自然と人間の作り出した美が調和している様子を示しています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "64", "text": "この写真の面白い点は、高層ビルの上に立つ大きなテレビ塔と、その近くに置かれた小さなレッドブルの缶の対比です。テレビ塔は非常に高く、その存在感が際立っていますが、その隣にあるレッドブ��の缶は、その大きさの違いがユーモラスに感じられます。この対比は、日常的なアイテムが巨大な構造物の近くに置かれることで、視覚的に面白い効果を生み出しています。また、レッドブルの缶は、エネルギー飲料としてのイメージを強調し、都市の活気やエネルギーを象徴しているとも考えられます。このように、異なるスケールの物体が共存することで、視覚的な興味を引き、観る人に印象的なシーンを提供しています。", "answer": "この写真の面白い点は、東京スカイツリーとRedBullの缶を模した広告オブジェが同一の視界に収まっていることです。スカイツリーの優雅な構造と、エネルギッシュなイメージのRedBull缶オブジェが対照的な印象を与え、都市のダイナミズムとクリエイティビティを象徴しています。また、高層ビルの背景がこのユニークなコントラストを一層際立たせています。", "input_text": "この写真の面白い点は何か説明してもらえますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "65", "text": "この写真は、東京スカイツリーとその隣にある建物の高さを対比しています。スカイツリーは非常に高く、その隣の建物と比べてその大きさが際立っています。", "answer": "この写真は、東京スカイツリーという現代的な建築物と、RedBullの缶を模した広告オブジェがあるビルの屋上という、商業的なアイテムを対比しています。また、これらの特徴的なオブジェクトと背後に見える高層ビル群との間の都市的な景観も対比しています。", "input_text": "この写真は何と何を対比していますか？", "heron-bench": {"score": 6, "score_gpt": 9}}
+{"question_id": "66", "text": "写真に写っているランドマークは東京スカイツリーです。", "answer": "写真に写っているランドマークの名称は東京スカイツリーです。", "input_text": "写真に写っているランドマークの名称はなんですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "67", "text": "左側に写っている飲みものは「レッドブル」の缶です。", "answer": "左側に写っている飲み物はRedBullです。これはビルの屋上に設置されたRedBullの缶を模した広告オブジェによるものです。", "input_text": "左側に写っている飲みものは何ですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "68", "text": "この画像には、東京スカイツリーの近くに立つ建物の前にある赤い缶が写っています。スカイツリーは非常に高く、その存在感が際立っています。赤い缶は、建物の前に置かれており、周囲の景色と対比を成しています。\n\nスカイツリーの背景には、青空が広がっており、その下に建物が立ち並んでいます。また、スカイツリーの近くには、いくつかの人が見え、彼らはスカイツリーを眺めたり、写真を撮ったりしているようです。全体的に、都市の活気とスカイツリーの壮大さが感じられるシーンです。", "answer": "この写真は、地上から見上げた青空の下の風景を捉えています。右側には、東京スカイツリーがそびえ立ち、左手にあるビルの屋上にはRedBullの缶を模した大きな広告オブジェが設置されています。このオブジェの後方には、いくつかの高層ビルが見え、都市の活気を感じさせます。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "69", "text": "この建物は16世紀に建設されました。", "answer": "金閣寺は室町時代に建設されました。", "input_text": "この建物はいつの時代に建設されましたか？", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "70", "text": "写真の池には、金色の建物が写っています。この建物は池の中央に位置しており、その美しい反射が水面に映っています。", "answer": "写真の池には金閣寺が反射して写っています。池の表面は静かで、金閣寺の鮮やかな金色の外観が鮮明に映し出されています。また、池の周囲には緑豊かな庭園が広がっており、その一部も水面に映り込んでいます。", "input_text": "写真の池には何が写っていますか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "71", "text": "この建物は「金閣寺」です。正式には「鹿苑寺」と呼ばれ、京都にある有名な寺院です。", "answer": "写真に写っている建物の名称は金閣寺です。", "input_text": "写真に写っている建物の名称はなん���すか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "72", "text": "この建物は金色です。", "answer": "この建物は金色です。金閣寺はその名の通り、外壁が金箔で覆われており、その輝きが池の水面にも美しく反映されています。", "input_text": "この建物は何色でしょうか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "73", "text": "この画像には、美しい池の中央に浮かぶ金色の建物が写っています。この建物は、池の水面に映り込み、非常に印象的な景観を作り出しています。池の周りには緑豊かな木々が立ち並び、自然の美しさが際立っています。\n\n建物は二階建てで、金色の外観が特徴的です。池の水面には、その建物の反射が映り、まるで鏡のように見えます。また、池の周囲にはいくつかの木が点在しており、自然との調和が感じられます。全体的に、静かで穏やかな雰囲気が漂う風景です。", "answer": "この写真は、金色に輝く金閣寺が中心に位置しており、その美しい姿が静かな池に映り込んでいます。池の周囲には、様々な種類の緑の樹木が豊かに生い茂り、金閣寺の周囲を囲むように広がっています。金閣寺とその反射、そして周囲の緑が、和の美しさを際立たせる調和のとれた風景を作り出しています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "74", "text": "ある日、二羽の白鳥が美しい緑の田んぼに降り立ちました。彼らは水辺で静かに佇み、周囲の自然を楽しんでいました。この田んぼは、穏やかな水が流れ、豊かな緑の草が生い茂る場所で、白鳥たちにとって理想的な休息の場所でした。\n\n白鳥たちは、お互いに寄り添いながら、穏やかな時間を過ごしていました。しかし、突然、空から大きな影が落ち、彼らの周りに何かが近づいてくるのを感じました。それは、田んぼの上を飛ぶ大きな鳥でした。白鳥たちは驚き、少し警戒しながらも、その鳥が何者であるのかを観察しました。\n\nその鳥は、白鳥たちに優しく近づき、彼らに話しかけました。その鳥は、この田んぼの守護者であり、白鳥たちが安全に過ごせるように見守っていると告げました。白鳥たちは、その鳥の優しさに感動し、安心してその場に留まることにしました。\n\nその後、白鳥た", "answer": "かつて、綺麗な水田に住む2羽の白鷺がいました。彼らは互いに背を向けて立っているものの、実は深い絆で結ばれた兄弟でした。ある日、彼らは互いの違いを乗り越え、力を合わせて水田を守る大冒険に出ることになります。", "input_text": "この白鳥二羽を主人公にしたおとぎ話を書いてください。", "heron-bench": {"score": 7, "score_gpt": 9}}
+{"question_id": "75", "text": "2羽の鳥は、互いに反対の方向を向いています。", "answer": "2羽の白鷺は互いに反対方向を向いています。", "input_text": "2羽の鳥はどちらを向いていますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "76", "text": "画像には2羽の鳥が写っています。", "answer": "画像には2羽の白鷺がいます。", "input_text": "鳥は何羽いますか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "77", "text": "白い鳥が2羽います。", "answer": "画像に基づいて、鳥の種類は白鷺です。画像には2羽の白鷺が写っており、それぞれが反対方向を向いて静かに立っています。白鷺はその純白の羽毛と優雅な姿で知られています。", "input_text": "鳥の種類はなんですか？", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "78", "text": "画像には、緑豊かな田んぼの中で2羽の白い鳥が水を飲んでいる様子が写っています。鳥たちはそれぞれ異なる位置に立っており、一方は画面の左側に、もう一方は右側にいます。彼らは穏やかな水辺でリラックスしているように見え、周囲の自然と調和しています。\n\n田んぼの水面は穏やかで、鳥たちの姿が映り込んでいます。背景には広がる緑の田んぼが見え、自然の美しさが感じられます。このシーンは、静かな田園風景を思わせる、穏やかな雰囲気を醸し出しています。", "answer": "この写真は、緑豊かな水田の中に2羽の白鷺が映っています。白鷺たちは互いに背を向け合っており、一羽は左を、もう一羽は右を向いて静かに立っています。彼らの周りは静寂に包まれており、自然の美しさと調和しています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 8, "score_gpt": 10}}
+{"question_id": "79", "text": "この写真は富士山が背景に見えることから、静岡県または山梨県で撮影されたと考えられます。", "answer": "この写真は静岡県で撮影されたと考えられます。宝永山が富士山の右側に見えることから、この地理的特徴は静岡県からの視点に一致します。また、東海道線の電車が写っていることも、この地域で撮影されたことを示唆しています。", "input_text": "この写真はどの都道府県で撮影されたと考えられますか？", "heron-bench": {"score": 8, "score_gpt": 10}}
+{"question_id": "80", "text": "写真に写っているのは、富士山です。", "answer": "写真に写っているのは富士山です。また、富士山の右側には宝永山が見えます。これらの山々は、静岡県から撮影されたことが示唆されています。", "input_text": "写真に写っているのは、何という山ですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "81", "text": "写真に写っている山の頂上付近は雪に覆われており、白く輝いています。", "answer": "富士山の頂上付近は雪に覆われており、その白さが澄み切った空に映えています。頂上は平らで、特徴的な円錐形のシルエットを描いています。宝永山が右側に見えることから、富士山の頂上は画像の中央よりやや左側に位置しています。", "input_text": "写真に写っている山の頂上付近はどのように見えますか？", "heron-bench": {"score": 10, "score_gpt": 9}}
+{"question_id": "82", "text": "この画像には、電車が駅に停車している様子が写っています。電車はオレンジとグレーの色合いで、背景には美しい山が見えます。電車の近くには、数人の人々が立っており、彼らは電車の到着を待っているようです。\n\n電車の左側には、一人の人物が立っており、その近くにはもう一人がいます。さらに、右側には別の人物がいて、彼らは電車のドアの近くに集まっています。電車のドアは開いており、乗客が乗り降りする準備をしている様子が伺えます。全体的に、駅の活気ある雰囲気が感じられます。", "answer": "この写真は、雪を頂に戴いた壮大な富士山を背景に、東海道線の特徴的なオレンジラインが入った電車が走行している様子を捉えています。電車は画像の中央やや下に位置し、その背後には宝永山が富士山の右側に小さく見えています。また、電柱とその間に張られた電線が画像の前景に複数存在し、日本の典型的な田園風景の一部として描かれています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 3, "score_gpt": 10}}
+{"question_id": "83", "text": "この写真は春に撮影されたと考えられます。桜の花が咲いていることから、春の訪れを示しています。", "answer": "この写真は春の季節に撮影されたものです。満開の桜が春の訪れを告げており、青空がその背景を美しく彩っています。桜の花がこの時期特有の風景を作り出しています。", "input_text": "この写真が撮影された季節はいつでしょうか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "84", "text": "この写真は桜の花が満開の時期に撮影されたもので、春の訪れを象徴しています。この時期、多くの学校では「入学式」や「卒業式」などの重要な行事が行われます。特に桜の花は、新しい始まりや別れを象徴するため、これらの行事と結びついています。また、桜の花が咲く時期は、多くの学生や保護者が学校の行事に参加するため、特別な意味を持つことが多いです。このように、桜の花が咲く時期は、学校生活の重要な節目を祝う機会でもあるのです。", "answer": "この写真が取られた時期に実施される学校行事は、入学式や卒業式です。春の季節は新しい生活が始まる時期であり、学生たちが新たなスタートを切る象徴的な行事が行われます。特に、桜の満開の下での式典は、日本の学校において美しい風物詩となっています。", "input_text": "この写真が取られた時期に実施される学校行事は何でしょうか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "85", "text": "写真には桜の花が写っています。桜の花は、春の象徴として知られ、美しいピンク色の花びらを持っています。", "answer": "写真に写っている花は桜です。��開の状態で、春の青空の下に美しく咲いています。これらの桜は、季節の象徴として、新しい始まりを感じさせます。", "input_text": "写真に写っている花は何ですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "86", "text": "写真に写っている花は満開で、多くの花が咲いている様子が見られます。", "answer": "写真に写っている桜は満開です。空に向かって広がる枝には、ほとんどの花が開いており、見事な花見景色を提供しています。この光景は春の訪れを象徴しています。", "input_text": "写真に写っている花は何分咲きでしょうか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "87", "text": "この画像には、白い花が咲き誇る木が写っています。木の枝には、たくさんの小さな白い花が密集しており、まるで雪が降り積もったかのように見えます。花は非常に美しく、青空を背景に鮮やかに映えています。\n\n木の全体が花で覆われており、特に枝の先端には多くの花が集まっています。花の色合いは淡いピンクから白にかけて変化しており、全体的に柔らかい印象を与えています。この木は、春の訪れを感じさせるような、穏やかで心地よい雰囲気を醸し出しています。", "answer": "この写真は、春の青空を背景に、満開の桜の木が中心に写っています。桜の花びらが密集して咲いており、その下では数人が花見を楽しんでいる様子が見受けられます。画像の前景には、桜の花びらが散り始めている様子も捉えられており、季節の移り変わりを感じさせます。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 10, "score_gpt": 9}}
+{"question_id": "88", "text": "茶畑が段々に分かれている理由は、主に地形の利用と効率的な栽培方法に関連しています。この画像では、茶畑が山の斜面に沿って段々に配置されていることがわかります。このような設計は、急な斜面での水はけを良くし、土壌の浸食を防ぐために重要です。また、段々にすることで、茶の木が均等に日光を受けることができ、成長を促進します。さらに、段々の茶畑は、作業効率を高めるためにも役立ちます。農家は、段々の畑を移動しながら作業を行いやすくなり、収穫や管理が容易になります。このように、段々に分かれた茶畑は、自然環境を最大限に活用し、持続可能な農業を実現するための工夫がなされているのです。", "answer": "茶畑が段々に分かれているのは、斜面を最大限に利用して栽培面積を増やし、水はけと日照条件を改善するためと考えられます。このような構造は、不均一な地形においても農作物の生育環境を最適化するのに役立ちます。また、段々畑は土壌の浸食を防ぎ、水資源の管理にも有効です。", "input_text": "茶畑はなぜ段々に分かれていると考えられるか論じてください。", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "89", "text": "この写真に写っているのは茶畑であり、主に緑茶がつくられる農作物です。", "answer": "この写真の農作物から作られる飲み物はお茶です。", "input_text": "この写真の農作物からつくられる飲み物は何でしょうか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "90", "text": "この写真は、緑豊かな丘陵地帯の風景を対象にして撮影されています。特に、茶畑が広がる美しい景色が特徴です。", "answer": "この写真は、日本の茶畑を対象に撮影されています。茶畑は複数の段々に分かれており、中央左手では3人の農家が作業をしています。また、写真の中央右手には木造の作業小屋が見えます。", "input_text": "この写真はどのような風景を対象にして、撮影されていますか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "91", "text": "作業小屋は写真の中央やや右寄りに位置しています。", "answer": "作業小屋は写真の中央右手に位置しています。木造であり、茶畑の段々の上部に見えます。小屋は農家が作業している場所から見て、右側にあります。", "input_text": "作業小屋は写真のどのあたりにうつっていますか？", "heron-bench": {"score": 10, "score_gpt": 9}}
+{"question_id": "92", "text": "この画像には、緑豊かな丘陵地帯が広がっており、その中には小さな小屋がいくつか点在しています。丘の斜面には、茶畑が広がっており、その美しい景観が印象的です。小屋は、丘の上���中腹に位置しており、自然に溶け込んでいます。\n\n丘の上には、数人の人々が見え、彼らはおそらく茶畑の手入れや観光を楽しんでいるのでしょう。人々はそれぞれ異なる位置に立っており、丘の斜面を歩いている様子が伺えます。全体的に、穏やかで自然豊かな風景が広がっており、リラックスした雰囲気が漂っています。", "answer": "この写真は、緑豊かな茶畑が特徴的な風景を捉えています。茶畑は美しく整備された段々に分かれており、中央左手では3人の農家が丁寧に作業を行っています。中央右手には、伝統的な木造の作業小屋があり、その周囲には自然な緑が広がっています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "93", "text": "いいえ、交差点に向かうことはできません。画像には、交通整理をしている作業員が立っており、その作業員は交通を制御しています。このため、交差点を通過する前に、作業員の指示に従う必要があります。作業員が交通を制御している場合、特に工事や道路の状況によっては、通行が制限されている可能性があります。安全を確保するためにも、作業員の指示に従い、交差点を通過する際には注意が必要です。", "answer": "いいえ、このまま交差点に向かってはいけません。信号機が赤色であり、手前の交通誘導員が「止まれ」の合図をしているため、停止する必要があります。また、片側交互通行の指示があるため、交通誘導員の指示に従う必要があります。", "input_text": "あなたは車を運転しており、交差点で右折しようと思っています。このまま交差点に向かってよいでしょうか？", "heron-bench": {"score": 9, "score_gpt": 10}}
+{"question_id": "94", "text": "はい、青信号になったら車を発進させても問題ありません。画像には、青信号が点灯している交通信号が映っており、その信号は車両の進行を許可しています。ただし、周囲の状況にも注意が必要です。特に、道路工事のためのコーンが設置されていることから、工事区域に近づく際には、他の作業員や車両の動きにも気を配る必要があります。また、工事区域では通常の交通ルールが異なる場合があるため、工事の指示に従うことが重要です。安全運転を心がけ、周囲の状況を確認しながら進むことが大切です。", "answer": "いいえ、青信号になったとしても、交通誘導員が「止まれ」の合図をしているため、車を発進させてはいけません。交通誘導員の指示は信号よりも優先されます。また、片側交互通行の状況下では、誘導員の指示に従う必要があります。", "input_text": "映像の状況において、青信号になったら車を発進させてもよいでしょうか？", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "95", "text": "信号機の色は赤色です。", "answer": "信号機の色は赤色です。", "input_text": "信号機の色は何色ですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "96", "text": "交通誘導員は、道路を横断しようとしている人々に対して、赤い旗を振って停止の合図を出しています。", "answer": "交通誘導員は「止まれ」の合図を出しています。信号機が赤色であり、手前の交通誘導員が手を上げて車両の進行を停止させるジェスチャーをしています。これは、ドライバーに対して現在進行方向に進むことを止めるよう指示していることを意味します。", "input_text": "交通誘導員はどのような指示を出していますか？", "heron-bench": {"score": 4, "score_gpt": 10}}
+{"question_id": "97", "text": "画像には、道路の真ん中に立っている作業員が写っています。彼は青い作業服を着ており、交通を制御するために赤い旗を持っています。周囲にはいくつかの交通信号が見え、道路の両側には車が停まっている様子が確認できます。\n\n作業員の近くには、数台の車が並んでおり、その中にはトラックも含まれています。また、道路の一部にはコーンが設置されており、作業区域を示しているようです。全体的に、交通の流れを管理するための作業が行われている様子が伺えます。", "answer": "この写真は、道路工事現場の様子を捉えており、前方には「止まれ」の合図をしている交通誘導員が一人立っています。信号機は赤色で、左側には片側交互通行を示す看板と右矢印がありま��。また、多数の青色カラーコーンが配置されており、その内側には白いトラックが停車しています。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}
+{"question_id": "98", "text": "奥に見える山の名前は「富士山」です。", "answer": "奥に見える山の名前は羊蹄山です。", "input_text": "奥に見える山の名前はなんですか？", "heron-bench": {"score": 1, "score_gpt": 10}}
+{"question_id": "99", "text": "この道を運転する際には、いくつかの重要な点に注意する必要があります。まず、雪が積もっているため、路面が滑りやすくなっています。特に、雪が深い場所や凍結している部分では、ブレーキをかける際やカーブを曲がる際に慎重になる必要があります。また、周囲の交通標識や信号も確認し、他の車両や歩行者との距離を保つことが重要です。さらに、視界が悪くなる可能性があるため、ヘッドライトを点灯し、他のドライバーに自分の存在を知らせることが大切です。最後に、急な天候の変化や雪の降り方にも注意を払い、必要に応じて運転を控える判断も必要です。安全運転を心がけることで、事故を未然に防ぐことができます。", "answer": "この道を運転する上で、まず雪に覆われた道路のため滑りやすい可能性があるので、慎重に運転する必要があります。次に、制限速度40kmの標識があるため、速度を適切に抑えることが重要です。また、青信号ではありますが、交差点や他の車両にも注意を払う必要があります。", "input_text": "この道を運転する上で気をつけるべきことは何でしょうか？", "heron-bench": {"score": 10, "score_gpt": 9}}
+{"question_id": "100", "text": "この場所の制限速度は40キロメートル毎時です。", "answer": "この場所の制限速度は40kmです。左側に制限速度40kmを示す標識が見えます。", "input_text": "この場所における制限速度はいくつですか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "101", "text": "12kmです。", "answer": "現在地からニセコまでは12kmです。", "input_text": "現在地からニセコまで何kmでしょうか？", "heron-bench": {"score": 10, "score_gpt": 10}}
+{"question_id": "102", "text": "この画像には、雪に覆われた道路が写っており、その先には大きな山がそびえています。道路は雪で覆われており、周囲には雪が積もった建物が見えます。道路の両側には、いくつかの交通標識が立っており、運転手に道案内をしています。\n\n道路の左側には、雪に覆われた建物がいくつか見え、その中にはバスも停まっているようです。また、道路の中央付近には信号機が設置されており、交通の流れを管理しています。全体的に、冬の静けさと雪の美しさが感じられる風景です。", "answer": "この写真は、冬の日に車窓から撮影された風景を捉えています。雪に覆われた道路が中央を通り、遠くには美しい羊蹄山が青空の下にそびえ立っています。画像の左側には制限速度40kmを示す標識があり、中央には青色の信号が見えます。また、ニセコまで12km、倶知安まで23kmであることを示す標識もあります。", "input_text": "この写真について詳しく説明してください。", "heron-bench": {"score": 6, "score_gpt": 10}}

qwen2_jp.jinja

@@ -0,0 +1,11 @@
+{% if messages[0]['role'] != 'system' %}
+    {{ '<|im_start|>system\n以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。<|im_end|>\n' }}
+{% endif %}
+
+{% for message in messages if message['content'] is not none %}
+    {{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n' }}
+{% endfor %}
+
+{% if add_generation_prompt %}
+    {{ '<|im_start|>assistant\n' }}
+{% endif %}

siglip_encoder.py

@@ -0,0 +1,288 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from accelerate.hooks import add_hook_to_module
+from einops import rearrange
+from s2wrapper import forward as multiscale_forward
+from transformers import AutoConfig, PretrainedConfig, PreTrainedModel, SiglipImageProcessor
+from transformers.image_processing_utils import BaseImageProcessor
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+from transformers.models.siglip import SiglipVisionModel
+
+
+class VisionTower(nn.Module):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__()
+
+        self.is_loaded = False
+
+        self.vision_tower_name = vision_tower
+        self.select_layer = getattr(args, "mm_vision_select_layer", -2)
+        self.select_feature = getattr(args, "mm_vision_select_feature", "patch")
+
+        self.cfg_only = None
+
+    def feature_select(self, image_forward_outs):
+        image_features = image_forward_outs.hidden_states[self.select_layer]
+        if self.select_feature == "patch":
+            image_features = image_features[:, 1:]
+        elif self.select_feature == "cls_patch":
+            image_features = image_features
+        else:
+            raise ValueError(f"Unexpected select feature: {self.select_feature}")
+        return image_features
+
+    def _maybe_resize_pos_embeds(
+        self,
+        model: PreTrainedModel,
+        image_processor: BaseImageProcessor,
+        resolution: int = -1,
+        interpolate_mode: str = "linear",
+    ):
+        if resolution in [model.config.image_size, -1]:
+            return
+        print(
+            f"Resizing vision model's position embeddings to support higher vision resolution: from {model.config.image_size} to {resolution} ..."
+        )
+        embeddings = model.vision_model.embeddings
+        patch_size = embeddings.patch_size
+        num_new_tokens = int((resolution // patch_size) ** 2)
+
+        old_embeddings = embeddings.position_embedding
+        match interpolate_mode:
+            case "linear":
+                ## Step 1: Calculate the corresponding patch ID (pid) in the current resolution (M patches) based on the target resolution (N patches). Formula: pid = pid / N * M
+                ## Step 2:  Obtain new embeddings by interpolating between the embeddings of the two nearest calculated patch IDs. Formula: new_embeds = (pid - floor(pid)) * embeds[ceil(pid)] + (ceil(pid) - pid) * embeds[floor(pid)]
+                import torch
+                import torch.nn as nn
+
+                if is_deepspeed_zero3_enabled():
+                    try:
+                        import deepspeed
+                    except ImportError:
+                        raise ImportError("DeepSpeed is not installed. Please install it with `pip install deepspeed`.")
+                    with deepspeed.zero.GatheredParameters([old_embeddings.weight], modifier_rank=None):
+                        old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
+                else:
+                    old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
+                new_embeddings = nn.Embedding(
+                    num_new_tokens,
+                    old_embedding_dim,
+                    dtype=old_embeddings.weight.dtype,
+                    device=old_embeddings.weight.device,
+                )
+                mapped_indices = (
+                    torch.arange(num_new_tokens).to(old_embeddings.weight.device)
+                    / (num_new_tokens - 1)
+                    * (old_num_tokens - 1)
+                )
+                floor_indices = torch.clamp(mapped_indices.floor().long(), min=0, max=old_num_tokens - 1)
+                ceil_indices = torch.clamp(mapped_indices.ceil().long(), min=0, max=old_num_tokens - 1)
+                if is_deepspeed_zero3_enabled():
+                    params = [old_embeddings.weight, new_embeddings.weight]
+                    with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                        interpolated_embeds = (mapped_indices - floor_indices)[:, None] * old_embeddings.weight.data[
+                            ceil_indices, :
+                        ] + (ceil_indices - mapped_indices)[:, None] * old_embeddings.weight.data[floor_indices, :]
+                else:
+                    interpolated_embeds = (mapped_indices - floor_indices)[:, None] * old_embeddings.weight.data[
+                        ceil_indices, :
+                    ] + (ceil_indices - mapped_indices)[:, None] * old_embeddings.weight.data[floor_indices, :]
+                new_embeddings.weight.data = interpolated_embeds
+            case _:
+                raise NotImplementedError
+
+        if hasattr(old_embeddings, "_hf_hook"):
+            hook = old_embeddings._hf_hook
+            add_hook_to_module(new_embeddings, hook)
+        new_embeddings.requires_grad_(old_embeddings.weight.requires_grad)
+        ## update vision encoder's configurations
+        model.config.image_size = resolution
+        if hasattr(image_processor, "crop_size"):
+            # CLIP vision tower
+            image_processor.crop_size = resolution
+        else:
+            # SIGLIP vision tower
+            assert hasattr(image_processor, "size")
+            image_processor.size = {"height": resolution, "width": resolution}
+        ## TODO define a '_reinitialize' method for VisionTower
+        embeddings.position_embedding = new_embeddings
+        embeddings.image_size = resolution
+        embeddings.num_patches = embeddings.num_positions = num_new_tokens
+        embeddings.position_ids = (
+            torch.arange(embeddings.num_positions).expand((1, -1)).to(old_embeddings.weight.device)
+        )
+
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_tower(
+                    image.to(device=self.device, dtype=self.dtype).unsqueeze(0),
+                    output_hidden_states=True,
+                )
+                image_feature = self.feature_select(image_forward_out).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            image_forward_outs = self.vision_tower(
+                images.to(device=self.device, dtype=self.dtype),
+                output_hidden_states=True,
+            )
+            image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device
+
+    @property
+    def config(self):
+        if self.is_loaded:
+            return self.vision_tower.config
+        else:
+            return self.cfg_only
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2
+
+
+class VisionTowerS2(VisionTower):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__(vision_tower, args, delay_load)
+
+        self.scales = list(map(int, args.s2_scales.split(",")))
+        self.scales.sort()
+        self.max_split_size = args.s2_max_split_size
+        self.resize_output_to_scale_idx = getattr(args, "s2_resize_output_to_scale_idx", 0)
+
+    def forward_feature(self, images):
+        image_forward_outs = self.vision_tower(
+            images.to(device=self.device, dtype=self.dtype), output_hidden_states=True
+        )
+        image_features = self.feature_select(image_forward_outs).to(images.dtype)
+        return image_features
+
+    def forward(self, images):
+        if type(images) is list:
+            image_feature = []
+            for image in images:
+                image_feature = multiscale_forward(
+                    self.forward_feature,
+                    image.unsqueeze(0),
+                    img_sizes=self.scales,
+                    max_split_size=self.max_split_size,
+                    resize_output_to_idx=self.resize_output_to_scale_idx,
+                )
+                image_features.append(image_feature)
+        else:
+            image_features = multiscale_forward(
+                self.forward_feature,
+                images,
+                img_sizes=self.scales,
+                max_split_size=self.max_split_size,
+                resize_output_to_idx=self.resize_output_to_scale_idx,
+            )
+
+        return image_features
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size * len(self.scales)
+
+
+class VisionTowerDynamicS2(VisionTower):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__(vision_tower, args, delay_load)
+
+        self.scales = list(map(int, args.s2_scales.split(",")))
+        self.scales.sort()
+        self.max_split_size = args.s2_max_split_size
+        self.resize_output_to_scale_idx = getattr(args, "s2_resize_output_to_scale_idx", 0)
+
+    def forward_feature(self, images):
+        image_forward_outs = self.vision_tower(
+            images.to(device=self.device, dtype=self.dtype), output_hidden_states=True
+        )
+        image_features = self.feature_select(image_forward_outs).to(images.dtype)
+        return image_features
+
+    def forward(self, images):
+        assert type(images) is not list
+        image_features = self.forward_feature(images)
+
+        return image_features
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size * len(self.scales)
+
+
+class SiglipVisionTower(VisionTower):
+    def __init__(self, model_name_or_path: str, config: PretrainedConfig) -> None:
+        super().__init__(model_name_or_path, config)
+        # TODO(ligengl): why pass config here leading to errors?
+        self.vision_tower = SiglipVisionModel.from_pretrained(
+            model_name_or_path,
+            attn_implementation=config._attn_implementation,
+            torch_dtype=eval(config.model_dtype),
+        )
+        self.image_processor = SiglipImageProcessor.from_pretrained(model_name_or_path)
+        self.is_loaded = True
+
+
+class SiglipVisionTowerS2(VisionTowerS2):
+    def __init__(self, model_name_or_path: str, config: PretrainedConfig) -> None:
+        super().__init__(model_name_or_path, config)
+        self.vision_tower = SiglipVisionModel.from_pretrained(
+            model_name_or_path,
+            attn_implementation=config._attn_implementation,
+            torch_dtype=eval(config.model_dtype),
+        )
+        self.image_processor = SiglipImageProcessor.from_pretrained(model_name_or_path)
+        # Make sure it crops/resizes the image to the largest scale in self.scales to maintain high-res information
+        self.image_processor.size["height"] = self.image_processor.size["width"] = self.scales[-1]
+        self.is_loaded = True
+
+
+class SiglipVisionTowerDynamicS2(VisionTowerDynamicS2):
+    def __init__(self, model_name_or_path: str, config: PretrainedConfig) -> None:
+        super().__init__(model_name_or_path, config)
+        self.vision_tower = SiglipVisionModel.from_pretrained(
+            model_name_or_path,
+            attn_implementation="flash_attention_2",
+            torch_dtype=eval(config.model_dtype),
+        )
+        self.image_processor = SiglipImageProcessor.from_pretrained(model_name_or_path)
+        # Make sure it crops/resizes the image to the largest scale in self.scales to maintain high-res information
+        self.image_processor.size["height"] = self.image_processor.size["width"] = self.scales[0]
+        self.is_loaded = True

tokenizer_utils.py

@@ -0,0 +1,182 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Dict, List, Optional, Sequence
+
+import torch
+import transformers
+
+from .constants import IGNORE_INDEX, SENTINEL_TOKEN
+from .conversation import SeparatorStyle, default_conversation
+from .mm_utils import tokenizer_image_token
+
+# __all__ = [
+#     "tokenize_conversation",
+#     "preprocess_conversation",
+#     "infer_stop_tokens",
+# ]
+
+DUMMY_CONVERSATION = [
+    {"from": "human", "value": "question"},
+    {"from": "gpt", "value": "answer"},
+] * 10
+
+
+def tokenize_conversation_legacy(
+    messages: Sequence[Dict[str, str]],
+    tokenizer: transformers.PreTrainedTokenizer,
+    add_generation_prompt: bool = False,
+    overrides: Optional[Dict[str, str]] = None,
+    no_system_prompt: bool = False,
+) -> torch.Tensor:
+    conv = default_conversation.copy()
+    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
+
+    if no_system_prompt:
+        conv.system = ""
+
+    # Skip the first message if it is not from human
+    if messages[0]["from"] != "human":
+        messages = messages[1:]
+
+    # Add a generation prompt if needed
+    if add_generation_prompt:
+        messages.append({"from": "gpt", "value": None})
+
+    conv.messages = []
+    for turn, message in enumerate(messages):
+        role = roles[message["from"]]
+        assert role == conv.roles[turn % 2]
+        if overrides is not None and message["from"] in overrides:
+            conv.append_message(role, overrides[message["from"]])
+        else:
+            conv.append_message(role, message["value"])
+
+    return tokenizer_image_token(conv.get_prompt(), tokenizer, return_tensors="pt")
+
+
+def tokenize_conversation(
+    messages: Sequence[Dict[str, str]],
+    tokenizer: transformers.PreTrainedTokenizer,
+    add_generation_prompt: bool = False,
+    overrides: Optional[Dict[str, str]] = None,
+    no_system_prompt: bool = False,
+) -> torch.Tensor:
+    # Normalize the conversation before tokenization
+    for message in messages:
+        message["value"] = message["value"].strip()
+
+    if default_conversation.sep_style != SeparatorStyle.AUTO:
+        return tokenize_conversation_legacy(
+            messages,
+            tokenizer,
+            add_generation_prompt=add_generation_prompt,
+            overrides=overrides,
+            no_system_prompt=no_system_prompt,
+        )
+
+    conversation = []
+    for m in messages:
+        message = {}
+        if m["from"] == "human":
+            message["role"] = "user"
+        elif m["from"] == "gpt":
+            message["role"] = "assistant"
+        else:
+            raise ValueError(f"Unexpected sender '{m['from']}' in conversation entry.")
+
+        message["content"] = m["value"]
+        if overrides is not None and m["from"] in overrides:
+            message["content"] = overrides[m["from"]]
+        conversation.append(message)
+
+    if no_system_prompt:
+        conversation = [{"role": "system", "content": ""}] + conversation
+
+    text = tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=add_generation_prompt,
+        tokenize=False,
+    )
+    return tokenizer_image_token(text, tokenizer, return_tensors="pt")
+
+
+def _maybe_add_sentinel_token(tokenizer: transformers.PreTrainedTokenizer) -> None:
+    if not hasattr(tokenizer, "sentinel_token"):
+        tokenizer.add_tokens([SENTINEL_TOKEN], special_tokens=True)
+        tokenizer.sentinel_token = SENTINEL_TOKEN
+        tokenizer.sentinel_token_id = tokenizer.convert_tokens_to_ids(SENTINEL_TOKEN)
+
+
+def preprocess_conversation(
+    conversation: Sequence[Dict[str, str]],
+    tokenizer: transformers.PreTrainedTokenizer,
+    no_system_prompt: bool = False,
+    retried: bool = False,
+) -> Dict[str, Any]:
+    inputs = tokenize_conversation(conversation, tokenizer, no_system_prompt=no_system_prompt)
+    labels = torch.ones_like(inputs) * IGNORE_INDEX
+
+    # Generate the template by replacing the assistant's response with a sentinel.
+    _maybe_add_sentinel_token(tokenizer)
+    template = tokenize_conversation(
+        conversation, tokenizer, overrides={"gpt": SENTINEL_TOKEN}, no_system_prompt=no_system_prompt
+    )
+
+    # Remove sentinel tokens from the template.
+    mask = torch.ones_like(template, dtype=torch.bool)
+    for k in range(template.size(0) - 1):
+        if template[k] == tokenizer.sentinel_token_id:
+            mask[k : k + 2] = False
+            # NOTE(zhijianl): This is to handle the corner case where there is an empty token before the sentinel token.
+            if k > 0 and retried:
+                mask[k - 1] = False
+    template = template[mask]
+
+    # Match the tokenized conversation with the template (with no assistant's response).
+    # Every token that is not matched will be included in the label for training.
+    p = 0
+    for k in range(inputs.size(0)):
+        if p < template.size(0) and inputs[k] == template[p]:
+            p += 1
+        else:
+            labels[k] = inputs[k]
+
+    # Mask all tokens in the label if the template is not fully matched.
+    if p < template.size(0):
+        if not retried:
+            return preprocess_conversation(
+                conversation,
+                tokenizer,
+                no_system_prompt=no_system_prompt,
+                retried=True,
+            )
+        print(f"Failed to process the conversation: '{conversation}'. All tokens will be masked in the label.")
+        labels[:] = IGNORE_INDEX
+
+    return {"input_ids": inputs, "labels": labels}
+
+
+def infer_stop_tokens(tokenizer: transformers.PreTrainedTokenizer) -> List[str]:
+    _maybe_add_sentinel_token(tokenizer)
+    template = tokenize_conversation(DUMMY_CONVERSATION, tokenizer, overrides={"gpt": SENTINEL_TOKEN})
+
+    stop_tokens = {tokenizer.eos_token}
+    for k in range(template.size(0) - 1):
+        if template[k] == tokenizer.sentinel_token_id:
+            stop_token = tokenizer.decode(template[k + 1])
+            stop_tokens.add(stop_token)
+    return list(stop_tokens)

trainer_state.json

@@ -0,0 +1,3983 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0,
+  "eval_steps": 500,
+  "global_step": 563,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0,
+      "grad_norm": 8.507063209836788,
+      "learning_rate": 1.1764705882352942e-06,
+      "loss": 1.4995,
+      "step": 1
+    },
+    {
+      "epoch": 0.0,
+      "grad_norm": 8.380837976348499,
+      "learning_rate": 2.3529411764705885e-06,
+      "loss": 1.4748,
+      "step": 2
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 16.153712765062593,
+      "learning_rate": 3.529411764705883e-06,
+      "loss": 1.5095,
+      "step": 3
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 6.32918606424435,
+      "learning_rate": 4.705882352941177e-06,
+      "loss": 1.5469,
+      "step": 4
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 7.9385610864888605,
+      "learning_rate": 5.882352941176471e-06,
+      "loss": 1.3485,
+      "step": 5
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 8.466808791062784,
+      "learning_rate": 7.058823529411766e-06,
+      "loss": 1.1535,
+      "step": 6
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 9.197832520423505,
+      "learning_rate": 8.23529411764706e-06,
+      "loss": 1.1178,
+      "step": 7
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 6.450381889902437,
+      "learning_rate": 9.411764705882354e-06,
+      "loss": 1.1143,
+      "step": 8
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 5.943842042765044,
+      "learning_rate": 1.0588235294117648e-05,
+      "loss": 1.0598,
+      "step": 9
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 5.210792511369386,
+      "learning_rate": 1.1764705882352942e-05,
+      "loss": 1.0608,
+      "step": 10
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 6.530214268858284,
+      "learning_rate": 1.2941176470588238e-05,
+      "loss": 0.9929,
+      "step": 11
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 4.087686549913091,
+      "learning_rate": 1.4117647058823532e-05,
+      "loss": 1.0114,
+      "step": 12
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 3.9477045662395813,
+      "learning_rate": 1.5294117647058822e-05,
+      "loss": 0.9954,
+      "step": 13
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 4.1326622498960095,
+      "learning_rate": 1.647058823529412e-05,
+      "loss": 0.9655,
+      "step": 14
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 2.8320217201459146,
+      "learning_rate": 1.7647058823529414e-05,
+      "loss": 0.9806,
+      "step": 15
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 3.2554966057925836,
+      "learning_rate": 1.8823529411764708e-05,
+      "loss": 0.9477,
+      "step": 16
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 3.2675871177766234,
+      "learning_rate": 2e-05,
+      "loss": 0.951,
+      "step": 17
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 3.934063085197528,
+      "learning_rate": 1.999983446750305e-05,
+      "loss": 0.9663,
+      "step": 18
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 3.241094261577659,
+      "learning_rate": 1.9999337875492412e-05,
+      "loss": 0.9121,
+      "step": 19
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 4.544380597079295,
+      "learning_rate": 1.9998510240408495e-05,
+      "loss": 0.935,
+      "step": 20
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 3.9294827890898,
+      "learning_rate": 1.999735158965141e-05,
+      "loss": 0.9038,
+      "step": 21
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 8.623886416931795,
+      "learning_rate": 1.9995861961580024e-05,
+      "loss": 0.9155,
+      "step": 22
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 5.132365358995958,
+      "learning_rate": 1.9994041405510705e-05,
+      "loss": 0.9172,
+      "step": 23
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 19.8213188609718,
+      "learning_rate": 1.9991889981715696e-05,
+      "loss": 0.8447,
+      "step": 24
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 2.2943297073592803,
+      "learning_rate": 1.998940776142111e-05,
+      "loss": 0.9511,
+      "step": 25
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 2.621703554351971,
+      "learning_rate": 1.9986594826804563e-05,
+      "loss": 0.8948,
+      "step": 26
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 10.369292417211861,
+      "learning_rate": 1.998345127099248e-05,
+      "loss": 0.9247,
+      "step": 27
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 2.1594855526771366,
+      "learning_rate": 1.9979977198056986e-05,
+      "loss": 0.896,
+      "step": 28
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 1.95675488611705,
+      "learning_rate": 1.997617272301248e-05,
+      "loss": 0.9416,
+      "step": 29
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 2.8649894461349144,
+      "learning_rate": 1.9972037971811802e-05,
+      "loss": 0.8633,
+      "step": 30
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 1.7651028451161137,
+      "learning_rate": 1.9967573081342103e-05,
+      "loss": 0.933,
+      "step": 31
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 2.8772648372928935,
+      "learning_rate": 1.9962778199420265e-05,
+      "loss": 0.8624,
+      "step": 32
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 1.619118612925012,
+      "learning_rate": 1.9957653484788054e-05,
+      "loss": 0.8723,
+      "step": 33
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 1.847616171745648,
+      "learning_rate": 1.995219910710682e-05,
+      "loss": 0.9186,
+      "step": 34
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 1.8003449391174806,
+      "learning_rate": 1.9946415246951928e-05,
+      "loss": 0.8574,
+      "step": 35
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 1.5572245128267164,
+      "learning_rate": 1.9940302095806733e-05,
+      "loss": 0.9054,
+      "step": 36
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 1.429846538560356,
+      "learning_rate": 1.9933859856056264e-05,
+      "loss": 0.8959,
+      "step": 37
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 2.3135286173820884,
+      "learning_rate": 1.992708874098054e-05,
+      "loss": 0.7889,
+      "step": 38
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 1.7675896646055616,
+      "learning_rate": 1.9919988974747473e-05,
+      "loss": 0.8598,
+      "step": 39
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 1.6088010620771458,
+      "learning_rate": 1.9912560792405465e-05,
+      "loss": 0.8749,
+      "step": 40
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 1.4194262632230676,
+      "learning_rate": 1.9904804439875635e-05,
+      "loss": 0.8915,
+      "step": 41
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 2.234751299184324,
+      "learning_rate": 1.9896720173943657e-05,
+      "loss": 0.8515,
+      "step": 42
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 1.6675796150462,
+      "learning_rate": 1.9888308262251286e-05,
+      "loss": 0.8868,
+      "step": 43
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 1.5198896510379474,
+      "learning_rate": 1.9879568983287468e-05,
+      "loss": 0.8712,
+      "step": 44
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 1.4817760559180928,
+      "learning_rate": 1.9870502626379127e-05,
+      "loss": 0.8518,
+      "step": 45
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 1.2950833873505025,
+      "learning_rate": 1.986110949168162e-05,
+      "loss": 0.8903,
+      "step": 46
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 1.8499963418893206,
+      "learning_rate": 1.9851389890168738e-05,
+      "loss": 0.7847,
+      "step": 47
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 1.7242837081062508,
+      "learning_rate": 1.9841344143622472e-05,
+      "loss": 0.847,
+      "step": 48
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 1.2806401368832885,
+      "learning_rate": 1.9830972584622327e-05,
+      "loss": 0.8828,
+      "step": 49
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 1.5629477066468145,
+      "learning_rate": 1.9820275556534306e-05,
+      "loss": 0.8251,
+      "step": 50
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 1.4745129474955476,
+      "learning_rate": 1.9809253413499565e-05,
+      "loss": 0.8397,
+      "step": 51
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 1.3870891516090702,
+      "learning_rate": 1.979790652042268e-05,
+      "loss": 0.8159,
+      "step": 52
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 1.4153257317082082,
+      "learning_rate": 1.9786235252959555e-05,
+      "loss": 0.8729,
+      "step": 53
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 1.5032889887849479,
+      "learning_rate": 1.9774239997505e-05,
+      "loss": 0.8311,
+      "step": 54
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 1.4793054274886537,
+      "learning_rate": 1.9761921151179937e-05,
+      "loss": 0.8304,
+      "step": 55
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 1.3623327446068922,
+      "learning_rate": 1.9749279121818235e-05,
+      "loss": 0.8246,
+      "step": 56
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 1.711010136615301,
+      "learning_rate": 1.9736314327953242e-05,
+      "loss": 0.8575,
+      "step": 57
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 2.1030806528939197,
+      "learning_rate": 1.9723027198803896e-05,
+      "loss": 0.884,
+      "step": 58
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 1.8114692078531625,
+      "learning_rate": 1.9709418174260523e-05,
+      "loss": 0.832,
+      "step": 59
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 4.199754342127851,
+      "learning_rate": 1.969548770487029e-05,
+      "loss": 0.8548,
+      "step": 60
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 1.8197457491567695,
+      "learning_rate": 1.9681236251822274e-05,
+      "loss": 0.7631,
+      "step": 61
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 3.9438599409073496,
+      "learning_rate": 1.9666664286932198e-05,
+      "loss": 0.8848,
+      "step": 62
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 1.2465928817428147,
+      "learning_rate": 1.9651772292626804e-05,
+      "loss": 0.8544,
+      "step": 63
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 1.5494633367087811,
+      "learning_rate": 1.96365607619279e-05,
+      "loss": 0.7678,
+      "step": 64
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 1.7179370824588993,
+      "learning_rate": 1.9621030198436007e-05,
+      "loss": 0.8439,
+      "step": 65
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 2.271674256525051,
+      "learning_rate": 1.9605181116313725e-05,
+      "loss": 0.8706,
+      "step": 66
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 1.8765382998169675,
+      "learning_rate": 1.9589014040268678e-05,
+      "loss": 0.8203,
+      "step": 67
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 1.7244243284677303,
+      "learning_rate": 1.957252950553616e-05,
+      "loss": 0.835,
+      "step": 68
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 1.4614578596506242,
+      "learning_rate": 1.955572805786141e-05,
+      "loss": 0.7874,
+      "step": 69
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 1.5281931827918855,
+      "learning_rate": 1.9538610253481542e-05,
+      "loss": 0.8225,
+      "step": 70
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 1.82847520738767,
+      "learning_rate": 1.952117665910714e-05,
+      "loss": 0.8493,
+      "step": 71
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 1.3951596613953592,
+      "learning_rate": 1.950342785190349e-05,
+      "loss": 0.8324,
+      "step": 72
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 2.0697823165315503,
+      "learning_rate": 1.9485364419471454e-05,
+      "loss": 0.7796,
+      "step": 73
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 1.2907212439977902,
+      "learning_rate": 1.9466986959828063e-05,
+      "loss": 0.8388,
+      "step": 74
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 1.8757593351488622,
+      "learning_rate": 1.9448296081386656e-05,
+      "loss": 0.8262,
+      "step": 75
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 1.0842852719284728,
+      "learning_rate": 1.9429292402936804e-05,
+      "loss": 0.8655,
+      "step": 76
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 1.849794221228235,
+      "learning_rate": 1.9409976553623767e-05,
+      "loss": 0.7859,
+      "step": 77
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 2.0412614125126365,
+      "learning_rate": 1.93903491729277e-05,
+      "loss": 0.8108,
+      "step": 78
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 1.2056345210077357,
+      "learning_rate": 1.9370410910642473e-05,
+      "loss": 0.8035,
+      "step": 79
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 1.374462281110202,
+      "learning_rate": 1.9350162426854152e-05,
+      "loss": 0.836,
+      "step": 80
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 1.3148936377412426,
+      "learning_rate": 1.932960439191915e-05,
+      "loss": 0.7965,
+      "step": 81
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 1.108816627540348,
+      "learning_rate": 1.9308737486442045e-05,
+      "loss": 0.8034,
+      "step": 82
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 1.0106641852379432,
+      "learning_rate": 1.9287562401253023e-05,
+      "loss": 0.8895,
+      "step": 83
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 1.0643031247749164,
+      "learning_rate": 1.926607983738504e-05,
+      "loss": 0.7938,
+      "step": 84
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 1.6202434008438316,
+      "learning_rate": 1.924429050605057e-05,
+      "loss": 0.8174,
+      "step": 85
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 1.1277481220625152,
+      "learning_rate": 1.9222195128618108e-05,
+      "loss": 0.7965,
+      "step": 86
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 1.2288767977862927,
+      "learning_rate": 1.9199794436588244e-05,
+      "loss": 0.8108,
+      "step": 87
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 1.8547807452034146,
+      "learning_rate": 1.9177089171569484e-05,
+      "loss": 0.7788,
+      "step": 88
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 1.02129245587083,
+      "learning_rate": 1.9154080085253665e-05,
+      "loss": 0.8123,
+      "step": 89
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 1.0460991831102135,
+      "learning_rate": 1.9130767939391092e-05,
+      "loss": 0.8053,
+      "step": 90
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 1.4621605349515425,
+      "learning_rate": 1.9107153505765308e-05,
+      "loss": 0.8817,
+      "step": 91
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 1.072411126394728,
+      "learning_rate": 1.908323756616754e-05,
+      "loss": 0.8013,
+      "step": 92
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 1.260321384657037,
+      "learning_rate": 1.9059020912370836e-05,
+      "loss": 0.8172,
+      "step": 93
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 1.3234441337963823,
+      "learning_rate": 1.9034504346103825e-05,
+      "loss": 0.7731,
+      "step": 94
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 1.2328175633896823,
+      "learning_rate": 1.900968867902419e-05,
+      "loss": 0.8484,
+      "step": 95
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 1.550206231336332,
+      "learning_rate": 1.898457473269181e-05,
+      "loss": 0.7686,
+      "step": 96
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 1.1944254976680975,
+      "learning_rate": 1.895916333854152e-05,
+      "loss": 0.7743,
+      "step": 97
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 1.3111106744110506,
+      "learning_rate": 1.8933455337855633e-05,
+      "loss": 0.824,
+      "step": 98
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 1.27000399731967,
+      "learning_rate": 1.8907451581736052e-05,
+      "loss": 0.7887,
+      "step": 99
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 1.0853884228742083,
+      "learning_rate": 1.8881152931076126e-05,
+      "loss": 0.8801,
+      "step": 100
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 1.1505416374261161,
+      "learning_rate": 1.8854560256532098e-05,
+      "loss": 0.7937,
+      "step": 101
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 1.136615577492445,
+      "learning_rate": 1.882767443849435e-05,
+      "loss": 0.817,
+      "step": 102
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 1.2158975241416552,
+      "learning_rate": 1.8800496367058187e-05,
+      "loss": 0.7749,
+      "step": 103
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 1.7879520459848115,
+      "learning_rate": 1.877302694199442e-05,
+      "loss": 0.7924,
+      "step": 104
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 1.3883544371253298,
+      "learning_rate": 1.8745267072719555e-05,
+      "loss": 0.7704,
+      "step": 105
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 1.7872291898285468,
+      "learning_rate": 1.8717217678265682e-05,
+      "loss": 0.7794,
+      "step": 106
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 1.3311382018956643,
+      "learning_rate": 1.8688879687250067e-05,
+      "loss": 0.8308,
+      "step": 107
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 1.2893581866568486,
+      "learning_rate": 1.866025403784439e-05,
+      "loss": 0.7793,
+      "step": 108
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 1.139241829453439,
+      "learning_rate": 1.863134167774369e-05,
+      "loss": 0.8892,
+      "step": 109
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 2.6133328625978893,
+      "learning_rate": 1.860214356413501e-05,
+      "loss": 0.8023,
+      "step": 110
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 1.5841826966676458,
+      "learning_rate": 1.8572660663665672e-05,
+      "loss": 0.8248,
+      "step": 111
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 1.4270048951849372,
+      "learning_rate": 1.8542893952411302e-05,
+      "loss": 0.7757,
+      "step": 112
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 1.7093625023294083,
+      "learning_rate": 1.8512844415843514e-05,
+      "loss": 0.7922,
+      "step": 113
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 1.3617190025737091,
+      "learning_rate": 1.8482513048797267e-05,
+      "loss": 0.787,
+      "step": 114
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 1.8071893998904143,
+      "learning_rate": 1.845190085543795e-05,
+      "loss": 0.8116,
+      "step": 115
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.3814019945231775,
+      "learning_rate": 1.842100884922812e-05,
+      "loss": 0.802,
+      "step": 116
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.4512890833634728,
+      "learning_rate": 1.838983805289396e-05,
+      "loss": 0.8536,
+      "step": 117
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.2549934379051162,
+      "learning_rate": 1.835838949839143e-05,
+      "loss": 0.7743,
+      "step": 118
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.208372599156025,
+      "learning_rate": 1.8326664226872063e-05,
+      "loss": 0.8024,
+      "step": 119
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.0972220457411088,
+      "learning_rate": 1.829466328864856e-05,
+      "loss": 0.8344,
+      "step": 120
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.2822175154411732,
+      "learning_rate": 1.826238774315995e-05,
+      "loss": 0.797,
+      "step": 121
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 1.2107730495640208,
+      "learning_rate": 1.8229838658936566e-05,
+      "loss": 0.7739,
+      "step": 122
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 0.9365728161852094,
+      "learning_rate": 1.819701711356464e-05,
+      "loss": 0.8338,
+      "step": 123
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 1.0864214560239218,
+      "learning_rate": 1.816392419365065e-05,
+      "loss": 0.7793,
+      "step": 124
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 0.9218244394223452,
+      "learning_rate": 1.8130560994785325e-05,
+      "loss": 0.7612,
+      "step": 125
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 2.080436030466144,
+      "learning_rate": 1.809692862150739e-05,
+      "loss": 0.8612,
+      "step": 126
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 1.4900488161819043,
+      "learning_rate": 1.8063028187266987e-05,
+      "loss": 0.7464,
+      "step": 127
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.8798162644545249,
+      "learning_rate": 1.8028860814388826e-05,
+      "loss": 0.8367,
+      "step": 128
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.8812260971377114,
+      "learning_rate": 1.7994427634035016e-05,
+      "loss": 0.755,
+      "step": 129
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.9676476461621292,
+      "learning_rate": 1.7959729786167616e-05,
+      "loss": 0.8508,
+      "step": 130
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.9256406288268074,
+      "learning_rate": 1.7924768419510906e-05,
+      "loss": 0.7292,
+      "step": 131
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.760232197066633,
+      "learning_rate": 1.7889544691513357e-05,
+      "loss": 0.797,
+      "step": 132
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 1.0639511668001849,
+      "learning_rate": 1.7854059768309292e-05,
+      "loss": 0.7521,
+      "step": 133
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 0.9073895992533008,
+      "learning_rate": 1.78183148246803e-05,
+      "loss": 0.8134,
+      "step": 134
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 0.905361446098707,
+      "learning_rate": 1.778231104401634e-05,
+      "loss": 0.7921,
+      "step": 135
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 1.0000546401715837,
+      "learning_rate": 1.7746049618276545e-05,
+      "loss": 0.7719,
+      "step": 136
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 1.2991918331253864,
+      "learning_rate": 1.7709531747949796e-05,
+      "loss": 0.7825,
+      "step": 137
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 0.9040019002054857,
+      "learning_rate": 1.767275864201494e-05,
+      "loss": 0.7751,
+      "step": 138
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 1.4065727827842003,
+      "learning_rate": 1.7635731517900782e-05,
+      "loss": 0.828,
+      "step": 139
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 1.6988580043162866,
+      "learning_rate": 1.759845160144579e-05,
+      "loss": 0.8597,
+      "step": 140
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 0.900317786852657,
+      "learning_rate": 1.756092012685749e-05,
+      "loss": 0.7999,
+      "step": 141
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 1.2873306206831179,
+      "learning_rate": 1.7523138336671628e-05,
+      "loss": 0.8174,
+      "step": 142
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 1.133316576026692,
+      "learning_rate": 1.7485107481711014e-05,
+      "loss": 0.7603,
+      "step": 143
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 1.953496025978647,
+      "learning_rate": 1.7446828821044123e-05,
+      "loss": 0.7743,
+      "step": 144
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 1.3312782549575295,
+      "learning_rate": 1.7408303621943417e-05,
+      "loss": 0.7661,
+      "step": 145
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 1.4629213396365295,
+      "learning_rate": 1.7369533159843368e-05,
+      "loss": 0.753,
+      "step": 146
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 0.9279779565723256,
+      "learning_rate": 1.7330518718298263e-05,
+      "loss": 0.8161,
+      "step": 147
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 1.1467927655177639,
+      "learning_rate": 1.729126158893969e-05,
+      "loss": 0.7701,
+      "step": 148
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 0.946811505107358,
+      "learning_rate": 1.7251763071433767e-05,
+      "loss": 0.8386,
+      "step": 149
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 1.33369128616317,
+      "learning_rate": 1.7212024473438145e-05,
+      "loss": 0.7646,
+      "step": 150
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 1.6320245391058008,
+      "learning_rate": 1.71720471105587e-05,
+      "loss": 0.7407,
+      "step": 151
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 1.5229292166958899,
+      "learning_rate": 1.7131832306305964e-05,
+      "loss": 0.8306,
+      "step": 152
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 0.949232772230085,
+      "learning_rate": 1.7091381392051333e-05,
+      "loss": 0.7453,
+      "step": 153
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 1.1465110573538513,
+      "learning_rate": 1.7050695706982977e-05,
+      "loss": 0.766,
+      "step": 154
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 1.5154779450711096,
+      "learning_rate": 1.7009776598061496e-05,
+      "loss": 0.7694,
+      "step": 155
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 1.0139789984795697,
+      "learning_rate": 1.6968625419975355e-05,
+      "loss": 0.8019,
+      "step": 156
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 1.1079537252923364,
+      "learning_rate": 1.6927243535095995e-05,
+      "loss": 0.7735,
+      "step": 157
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 1.0291789122473398,
+      "learning_rate": 1.6885632313432772e-05,
+      "loss": 0.8687,
+      "step": 158
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 1.0108245776654075,
+      "learning_rate": 1.684379313258757e-05,
+      "loss": 0.7863,
+      "step": 159
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 0.8450335525589678,
+      "learning_rate": 1.6801727377709195e-05,
+      "loss": 0.7804,
+      "step": 160
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 0.9063906180650181,
+      "learning_rate": 1.6759436441447544e-05,
+      "loss": 0.7411,
+      "step": 161
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 1.4659728097146998,
+      "learning_rate": 1.6716921723907472e-05,
+      "loss": 0.7311,
+      "step": 162
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 1.1450889946261562,
+      "learning_rate": 1.6674184632602447e-05,
+      "loss": 0.7997,
+      "step": 163
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 2.2474639045647633,
+      "learning_rate": 1.6631226582407954e-05,
+      "loss": 0.7823,
+      "step": 164
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 1.0473794040487285,
+      "learning_rate": 1.6588048995514657e-05,
+      "loss": 0.759,
+      "step": 165
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 1.131898412640055,
+      "learning_rate": 1.6544653301381318e-05,
+      "loss": 0.822,
+      "step": 166
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 1.3959580024538538,
+      "learning_rate": 1.6501040936687444e-05,
+      "loss": 0.8124,
+      "step": 167
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 1.1510662240838245,
+      "learning_rate": 1.645721334528577e-05,
+      "loss": 0.7697,
+      "step": 168
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 0.8682977945159439,
+      "learning_rate": 1.641317197815442e-05,
+      "loss": 0.7743,
+      "step": 169
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 1.4212192981288094,
+      "learning_rate": 1.6368918293348893e-05,
+      "loss": 0.7987,
+      "step": 170
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 1.005593257735101,
+      "learning_rate": 1.6324453755953772e-05,
+      "loss": 0.7801,
+      "step": 171
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 1.4643600406891248,
+      "learning_rate": 1.627977983803424e-05,
+      "loss": 0.7466,
+      "step": 172
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 1.3023912771038668,
+      "learning_rate": 1.6234898018587336e-05,
+      "loss": 0.8037,
+      "step": 173
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 1.311996017760399,
+      "learning_rate": 1.6189809783492984e-05,
+      "loss": 0.7564,
+      "step": 174
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 0.9928204683020392,
+      "learning_rate": 1.6144516625464813e-05,
+      "loss": 0.8662,
+      "step": 175
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 1.6857925213687284,
+      "learning_rate": 1.609902004400073e-05,
+      "loss": 0.7446,
+      "step": 176
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 1.2049608818051798,
+      "learning_rate": 1.605332154533328e-05,
+      "loss": 0.7564,
+      "step": 177
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 1.037483746060024,
+      "learning_rate": 1.600742264237979e-05,
+      "loss": 0.7439,
+      "step": 178
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 1.2577074894891214,
+      "learning_rate": 1.5961324854692254e-05,
+      "loss": 0.8322,
+      "step": 179
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 1.1897741362610719,
+      "learning_rate": 1.5915029708407056e-05,
+      "loss": 0.773,
+      "step": 180
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 1.2430173497955095,
+      "learning_rate": 1.5868538736194427e-05,
+      "loss": 0.7524,
+      "step": 181
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 1.3567406768823702,
+      "learning_rate": 1.582185347720771e-05,
+      "loss": 0.8127,
+      "step": 182
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 1.4881992614206114,
+      "learning_rate": 1.57749754770324e-05,
+      "loss": 0.8067,
+      "step": 183
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 1.0306297593822549,
+      "learning_rate": 1.5727906287634986e-05,
+      "loss": 0.7811,
+      "step": 184
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 2.0461723719467053,
+      "learning_rate": 1.568064746731156e-05,
+      "loss": 0.7193,
+      "step": 185
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 1.0088861798153181,
+      "learning_rate": 1.563320058063622e-05,
+      "loss": 0.803,
+      "step": 186
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 1.1671931811922227,
+      "learning_rate": 1.55855671984093e-05,
+      "loss": 0.7398,
+      "step": 187
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 0.9106562766396383,
+      "learning_rate": 1.553774889760533e-05,
+      "loss": 0.8089,
+      "step": 188
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 1.1983507765271422,
+      "learning_rate": 1.5489747261320866e-05,
+      "loss": 0.7691,
+      "step": 189
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 1.0575825543404433,
+      "learning_rate": 1.5441563878722042e-05,
+      "loss": 0.7866,
+      "step": 190
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 0.9806900643594609,
+      "learning_rate": 1.5393200344991993e-05,
+      "loss": 0.8081,
+      "step": 191
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 0.9313407419052266,
+      "learning_rate": 1.5344658261278013e-05,
+      "loss": 0.7764,
+      "step": 192
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 0.992472391192496,
+      "learning_rate": 1.5295939234638566e-05,
+      "loss": 0.768,
+      "step": 193
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 2.174356436614505,
+      "learning_rate": 1.524704487799008e-05,
+      "loss": 0.7436,
+      "step": 194
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 0.9838293314441882,
+      "learning_rate": 1.5197976810053546e-05,
+      "loss": 0.759,
+      "step": 195
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 1.1277555782814377,
+      "learning_rate": 1.5148736655300922e-05,
+      "loss": 0.8158,
+      "step": 196
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 1.027047047472426,
+      "learning_rate": 1.5099326043901361e-05,
+      "loss": 0.7511,
+      "step": 197
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 1.7265548361992147,
+      "learning_rate": 1.504974661166724e-05,
+      "loss": 0.7734,
+      "step": 198
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 17.66044356313852,
+      "learning_rate": 1.5000000000000002e-05,
+      "loss": 0.7617,
+      "step": 199
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 1.6424117340223483,
+      "learning_rate": 1.4950087855835816e-05,
+      "loss": 0.8207,
+      "step": 200
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 1.368367595003515,
+      "learning_rate": 1.490001183159105e-05,
+      "loss": 0.7698,
+      "step": 201
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 1.1195914174852997,
+      "learning_rate": 1.4849773585107574e-05,
+      "loss": 0.7486,
+      "step": 202
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 0.9960466064073642,
+      "learning_rate": 1.4799374779597866e-05,
+      "loss": 0.7687,
+      "step": 203
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 3.40677503554284,
+      "learning_rate": 1.4748817083589943e-05,
+      "loss": 0.7689,
+      "step": 204
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 1.1744057232173357,
+      "learning_rate": 1.4698102170872142e-05,
+      "loss": 0.7953,
+      "step": 205
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 0.9709376500781769,
+      "learning_rate": 1.4647231720437687e-05,
+      "loss": 0.7996,
+      "step": 206
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 1.0780628041187534,
+      "learning_rate": 1.4596207416429119e-05,
+      "loss": 0.7793,
+      "step": 207
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 1.198486325581988,
+      "learning_rate": 1.4545030948082518e-05,
+      "loss": 0.7584,
+      "step": 208
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 0.9476564594417786,
+      "learning_rate": 1.4493704009671614e-05,
+      "loss": 0.8241,
+      "step": 209
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 0.876018870277804,
+      "learning_rate": 1.4442228300451655e-05,
+      "loss": 0.7778,
+      "step": 210
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 0.8272624022636804,
+      "learning_rate": 1.439060552460318e-05,
+      "loss": 0.6966,
+      "step": 211
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 0.8534456906469495,
+      "learning_rate": 1.4338837391175582e-05,
+      "loss": 0.8315,
+      "step": 212
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 0.8151679715897892,
+      "learning_rate": 1.4286925614030542e-05,
+      "loss": 0.7912,
+      "step": 213
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 1.0240572196453213,
+      "learning_rate": 1.423487191178528e-05,
+      "loss": 0.7632,
+      "step": 214
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 0.9373128897118324,
+      "learning_rate": 1.4182678007755653e-05,
+      "loss": 0.7991,
+      "step": 215
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 0.9074042298617393,
+      "learning_rate": 1.4130345629899116e-05,
+      "loss": 0.7241,
+      "step": 216
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 0.9820070319451613,
+      "learning_rate": 1.4077876510757502e-05,
+      "loss": 0.8254,
+      "step": 217
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.5066789735427286,
+      "learning_rate": 1.4025272387399676e-05,
+      "loss": 0.6886,
+      "step": 218
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.047602052004599,
+      "learning_rate": 1.3972535001364014e-05,
+      "loss": 0.8107,
+      "step": 219
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.9028584981292163,
+      "learning_rate": 1.3919666098600753e-05,
+      "loss": 0.7859,
+      "step": 220
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.1236913663983221,
+      "learning_rate": 1.3866667429414188e-05,
+      "loss": 0.7656,
+      "step": 221
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.0891172388783499,
+      "learning_rate": 1.3813540748404734e-05,
+      "loss": 0.8203,
+      "step": 222
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 1.1667104560101977,
+      "learning_rate": 1.3760287814410822e-05,
+      "loss": 0.7695,
+      "step": 223
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 1.0970968340081568,
+      "learning_rate": 1.3706910390450679e-05,
+      "loss": 0.7794,
+      "step": 224
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 1.0353263645139252,
+      "learning_rate": 1.3653410243663953e-05,
+      "loss": 0.7557,
+      "step": 225
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 1.1246530594282715,
+      "learning_rate": 1.3599789145253226e-05,
+      "loss": 0.7557,
+      "step": 226
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 1.036301793654716,
+      "learning_rate": 1.3546048870425356e-05,
+      "loss": 0.7556,
+      "step": 227
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 1.092416722925642,
+      "learning_rate": 1.3492191198332725e-05,
+      "loss": 0.8048,
+      "step": 228
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 1.3287906600486554,
+      "learning_rate": 1.3438217912014318e-05,
+      "loss": 0.7952,
+      "step": 229
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 1.3462381401254768,
+      "learning_rate": 1.3384130798336705e-05,
+      "loss": 0.7581,
+      "step": 230
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 1.3294253157215177,
+      "learning_rate": 1.3329931647934884e-05,
+      "loss": 0.7723,
+      "step": 231
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 1.534142920030296,
+      "learning_rate": 1.3275622255152994e-05,
+      "loss": 0.7712,
+      "step": 232
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 1.122000700902407,
+      "learning_rate": 1.3221204417984907e-05,
+      "loss": 0.7998,
+      "step": 233
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.219606607341259,
+      "learning_rate": 1.3166679938014728e-05,
+      "loss": 0.7467,
+      "step": 234
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.2888174749050882,
+      "learning_rate": 1.311205062035711e-05,
+      "loss": 0.783,
+      "step": 235
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 0.9381076410902784,
+      "learning_rate": 1.3057318273597531e-05,
+      "loss": 0.7808,
+      "step": 236
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.0511440079542627,
+      "learning_rate": 1.300248470973239e-05,
+      "loss": 0.7559,
+      "step": 237
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.0349236600724503,
+      "learning_rate": 1.2947551744109044e-05,
+      "loss": 0.7867,
+      "step": 238
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 0.9962346927775738,
+      "learning_rate": 1.2892521195365679e-05,
+      "loss": 0.7784,
+      "step": 239
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 1.01211478570776,
+      "learning_rate": 1.2837394885371124e-05,
+      "loss": 0.7505,
+      "step": 240
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 1.1594840054734818,
+      "learning_rate": 1.2782174639164528e-05,
+      "loss": 0.8209,
+      "step": 241
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 1.5411694924722625,
+      "learning_rate": 1.2726862284894939e-05,
+      "loss": 0.7022,
+      "step": 242
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 11.599019209766738,
+      "learning_rate": 1.2671459653760781e-05,
+      "loss": 0.7805,
+      "step": 243
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 1.440097235149696,
+      "learning_rate": 1.2615968579949225e-05,
+      "loss": 0.7316,
+      "step": 244
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.3372366896231092,
+      "learning_rate": 1.2560390900575472e-05,
+      "loss": 0.7589,
+      "step": 245
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.105604126041751,
+      "learning_rate": 1.2504728455621935e-05,
+      "loss": 0.795,
+      "step": 246
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.5393643288243302,
+      "learning_rate": 1.244898308787731e-05,
+      "loss": 0.7405,
+      "step": 247
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 2.468274498361346,
+      "learning_rate": 1.2393156642875579e-05,
+      "loss": 0.761,
+      "step": 248
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.3238659895664049,
+      "learning_rate": 1.2337250968834912e-05,
+      "loss": 0.7449,
+      "step": 249
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.0300122533315417,
+      "learning_rate": 1.2281267916596474e-05,
+      "loss": 0.8029,
+      "step": 250
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 1.7441854429035688,
+      "learning_rate": 1.2225209339563144e-05,
+      "loss": 0.7578,
+      "step": 251
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 1.019686827811183,
+      "learning_rate": 1.2169077093638175e-05,
+      "loss": 0.7683,
+      "step": 252
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 1.1348511043608238,
+      "learning_rate": 1.2112873037163728e-05,
+      "loss": 0.7656,
+      "step": 253
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 1.226402214745372,
+      "learning_rate": 1.2056599030859367e-05,
+      "loss": 0.8597,
+      "step": 254
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 0.945689755126611,
+      "learning_rate": 1.2000256937760446e-05,
+      "loss": 0.7545,
+      "step": 255
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 1.3506734624035819,
+      "learning_rate": 1.1943848623156436e-05,
+      "loss": 0.762,
+      "step": 256
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 1.1218264228304475,
+      "learning_rate": 1.1887375954529167e-05,
+      "loss": 0.734,
+      "step": 257
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 1.1724721829239755,
+      "learning_rate": 1.1830840801491014e-05,
+      "loss": 0.7525,
+      "step": 258
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 0.8740511029135943,
+      "learning_rate": 1.1774245035722984e-05,
+      "loss": 0.8227,
+      "step": 259
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 1.1338144219674435,
+      "learning_rate": 1.1717590530912764e-05,
+      "loss": 0.7997,
+      "step": 260
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 2.4527177700373994,
+      "learning_rate": 1.1660879162692676e-05,
+      "loss": 0.6995,
+      "step": 261
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 1.1079167279809168,
+      "learning_rate": 1.1604112808577603e-05,
+      "loss": 0.8415,
+      "step": 262
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 1.3963442463864142,
+      "learning_rate": 1.1547293347902813e-05,
+      "loss": 0.7237,
+      "step": 263
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 1.060519941420844,
+      "learning_rate": 1.1490422661761744e-05,
+      "loss": 0.7969,
+      "step": 264
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 2.4247537601070173,
+      "learning_rate": 1.1433502632943735e-05,
+      "loss": 0.7365,
+      "step": 265
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 1.0395011318393783,
+      "learning_rate": 1.1376535145871685e-05,
+      "loss": 0.7323,
+      "step": 266
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 1.2872576522201582,
+      "learning_rate": 1.1319522086539666e-05,
+      "loss": 0.7514,
+      "step": 267
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 0.8579921224078123,
+      "learning_rate": 1.1262465342450494e-05,
+      "loss": 0.7888,
+      "step": 268
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.1043236121717428,
+      "learning_rate": 1.1205366802553231e-05,
+      "loss": 0.7512,
+      "step": 269
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.314140212229249,
+      "learning_rate": 1.1148228357180657e-05,
+      "loss": 0.7945,
+      "step": 270
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 0.8645287535152224,
+      "learning_rate": 1.1091051897986677e-05,
+      "loss": 0.7964,
+      "step": 271
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.0606036174497515,
+      "learning_rate": 1.10338393178837e-05,
+      "loss": 0.7485,
+      "step": 272
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.3434898828929964,
+      "learning_rate": 1.0976592510979982e-05,
+      "loss": 0.7659,
+      "step": 273
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 1.0226310780483951,
+      "learning_rate": 1.0919313372516894e-05,
+      "loss": 0.7436,
+      "step": 274
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 1.0194327081170713,
+      "learning_rate": 1.0862003798806195e-05,
+      "loss": 0.7326,
+      "step": 275
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 1.1420415262511832,
+      "learning_rate": 1.0804665687167262e-05,
+      "loss": 0.7716,
+      "step": 276
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 1.0405745082215876,
+      "learning_rate": 1.0747300935864245e-05,
+      "loss": 0.7562,
+      "step": 277
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 0.9411416560859605,
+      "learning_rate": 1.0689911444043249e-05,
+      "loss": 0.7892,
+      "step": 278
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.0638350695922085,
+      "learning_rate": 1.0632499111669453e-05,
+      "loss": 0.7924,
+      "step": 279
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 0.8896644123246137,
+      "learning_rate": 1.0575065839464207e-05,
+      "loss": 0.7944,
+      "step": 280
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.3851887010258916,
+      "learning_rate": 1.0517613528842096e-05,
+      "loss": 0.7523,
+      "step": 281
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 0.8564790731817052,
+      "learning_rate": 1.0460144081848006e-05,
+      "loss": 0.8007,
+      "step": 282
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.472385478361725,
+      "learning_rate": 1.0402659401094154e-05,
+      "loss": 0.742,
+      "step": 283
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.2343696078506177,
+      "learning_rate": 1.0345161389697083e-05,
+      "loss": 0.735,
+      "step": 284
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 1.0664256469995077,
+      "learning_rate": 1.0287651951214675e-05,
+      "loss": 0.7707,
+      "step": 285
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 2.7853431478072315,
+      "learning_rate": 1.0230132989583122e-05,
+      "loss": 0.7754,
+      "step": 286
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 1.4312677584129017,
+      "learning_rate": 1.0172606409053887e-05,
+      "loss": 0.7561,
+      "step": 287
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 1.434894152042441,
+      "learning_rate": 1.0115074114130676e-05,
+      "loss": 0.8314,
+      "step": 288
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 1.40363743791131,
+      "learning_rate": 1.0057538009506378e-05,
+      "loss": 0.7358,
+      "step": 289
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.4216319281561687,
+      "learning_rate": 1e-05,
+      "loss": 0.7821,
+      "step": 290
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.1843530926281,
+      "learning_rate": 9.942461990493627e-06,
+      "loss": 0.7269,
+      "step": 291
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.198960472374662,
+      "learning_rate": 9.884925885869326e-06,
+      "loss": 0.774,
+      "step": 292
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.3700679139092526,
+      "learning_rate": 9.827393590946116e-06,
+      "loss": 0.7289,
+      "step": 293
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.3134346846483278,
+      "learning_rate": 9.769867010416883e-06,
+      "loss": 0.735,
+      "step": 294
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.766483532838929,
+      "learning_rate": 9.71234804878533e-06,
+      "loss": 0.7761,
+      "step": 295
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 1.1321036509904978,
+      "learning_rate": 9.654838610302922e-06,
+      "loss": 0.7386,
+      "step": 296
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 1.1929839213242892,
+      "learning_rate": 9.597340598905851e-06,
+      "loss": 0.8203,
+      "step": 297
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 1.1981861529295657,
+      "learning_rate": 9.539855918151999e-06,
+      "loss": 0.7094,
+      "step": 298
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 1.0586477146206266,
+      "learning_rate": 9.482386471157905e-06,
+      "loss": 0.8288,
+      "step": 299
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 2.2947958743267374,
+      "learning_rate": 9.424934160535795e-06,
+      "loss": 0.6719,
+      "step": 300
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 1.1695655494977126,
+      "learning_rate": 9.367500888330546e-06,
+      "loss": 0.8039,
+      "step": 301
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.1394813339124892,
+      "learning_rate": 9.310088555956751e-06,
+      "loss": 0.7279,
+      "step": 302
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.2737270910437388,
+      "learning_rate": 9.252699064135759e-06,
+      "loss": 0.7462,
+      "step": 303
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.1045855514533351,
+      "learning_rate": 9.195334312832742e-06,
+      "loss": 0.7522,
+      "step": 304
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.1375565719120944,
+      "learning_rate": 9.137996201193807e-06,
+      "loss": 0.7923,
+      "step": 305
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.0738044490098744,
+      "learning_rate": 9.080686627483111e-06,
+      "loss": 0.7591,
+      "step": 306
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 1.2993614485597482,
+      "learning_rate": 9.02340748902002e-06,
+      "loss": 0.6994,
+      "step": 307
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 1.5117034251500119,
+      "learning_rate": 8.966160682116301e-06,
+      "loss": 0.8261,
+      "step": 308
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 1.3780522053648592,
+      "learning_rate": 8.908948102013326e-06,
+      "loss": 0.7309,
+      "step": 309
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 1.4972338805756236,
+      "learning_rate": 8.851771642819346e-06,
+      "loss": 0.727,
+      "step": 310
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 1.4680656432100794,
+      "learning_rate": 8.79463319744677e-06,
+      "loss": 0.7285,
+      "step": 311
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 2.313846141986661,
+      "learning_rate": 8.73753465754951e-06,
+      "loss": 0.7945,
+      "step": 312
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 1.2869463268384915,
+      "learning_rate": 8.680477913460339e-06,
+      "loss": 0.7959,
+      "step": 313
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 2.106280518292013,
+      "learning_rate": 8.62346485412832e-06,
+      "loss": 0.74,
+      "step": 314
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 1.3531581194667899,
+      "learning_rate": 8.566497367056268e-06,
+      "loss": 0.6986,
+      "step": 315
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 1.4197697588908693,
+      "learning_rate": 8.509577338238255e-06,
+      "loss": 0.781,
+      "step": 316
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 1.4349145256946136,
+      "learning_rate": 8.452706652097187e-06,
+      "loss": 0.7567,
+      "step": 317
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 1.677424093901339,
+      "learning_rate": 8.395887191422397e-06,
+      "loss": 0.7277,
+      "step": 318
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 3.0582740634732097,
+      "learning_rate": 8.339120837307324e-06,
+      "loss": 0.722,
+      "step": 319
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 3.7393227504364246,
+      "learning_rate": 8.28240946908724e-06,
+      "loss": 0.8167,
+      "step": 320
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 1.581227401509297,
+      "learning_rate": 8.225754964277018e-06,
+      "loss": 0.7352,
+      "step": 321
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 1.470503432709286,
+      "learning_rate": 8.169159198508987e-06,
+      "loss": 0.7977,
+      "step": 322
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 1.4594103330779356,
+      "learning_rate": 8.112624045470834e-06,
+      "loss": 0.7568,
+      "step": 323
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 1.5307349721492423,
+      "learning_rate": 8.056151376843568e-06,
+      "loss": 0.6967,
+      "step": 324
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 1.5603332209161358,
+      "learning_rate": 7.999743062239557e-06,
+      "loss": 0.7951,
+      "step": 325
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 2.5996132473445206,
+      "learning_rate": 7.943400969140635e-06,
+      "loss": 0.7474,
+      "step": 326
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 3.334122489866669,
+      "learning_rate": 7.887126962836274e-06,
+      "loss": 0.759,
+      "step": 327
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 1.9981089012739424,
+      "learning_rate": 7.830922906361827e-06,
+      "loss": 0.7641,
+      "step": 328
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 1.3813790955360155,
+      "learning_rate": 7.774790660436857e-06,
+      "loss": 0.7285,
+      "step": 329
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 1.2395963850748397,
+      "learning_rate": 7.718732083403531e-06,
+      "loss": 0.796,
+      "step": 330
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 1.8038963005397093,
+      "learning_rate": 7.662749031165093e-06,
+      "loss": 0.7354,
+      "step": 331
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 1.55966722915584,
+      "learning_rate": 7.606843357124426e-06,
+      "loss": 0.7729,
+      "step": 332
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 1.5927922672934738,
+      "learning_rate": 7.551016912122692e-06,
+      "loss": 0.6933,
+      "step": 333
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 1.0699588514141283,
+      "learning_rate": 7.4952715443780665e-06,
+      "loss": 0.8096,
+      "step": 334
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 1.6246603724667785,
+      "learning_rate": 7.4396090994245295e-06,
+      "loss": 0.7943,
+      "step": 335
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 1.1888479591553487,
+      "learning_rate": 7.384031420050779e-06,
+      "loss": 0.7953,
+      "step": 336
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 1.552386758593148,
+      "learning_rate": 7.328540346239222e-06,
+      "loss": 0.7619,
+      "step": 337
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 1.580381187107915,
+      "learning_rate": 7.273137715105063e-06,
+      "loss": 0.7629,
+      "step": 338
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 1.2304680598494862,
+      "learning_rate": 7.217825360835475e-06,
+      "loss": 0.7551,
+      "step": 339
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 7.949474583996868,
+      "learning_rate": 7.162605114628879e-06,
+      "loss": 0.7413,
+      "step": 340
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 1.7323775275331923,
+      "learning_rate": 7.107478804634324e-06,
+      "loss": 0.7075,
+      "step": 341
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 1.2608510194692342,
+      "learning_rate": 7.052448255890958e-06,
+      "loss": 0.7804,
+      "step": 342
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 1.347750741978453,
+      "learning_rate": 6.997515290267611e-06,
+      "loss": 0.7368,
+      "step": 343
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 1.6326688971820857,
+      "learning_rate": 6.942681726402474e-06,
+      "loss": 0.7809,
+      "step": 344
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 2.095587605698155,
+      "learning_rate": 6.887949379642893e-06,
+      "loss": 0.805,
+      "step": 345
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 2.3280279836708337,
+      "learning_rate": 6.833320061985278e-06,
+      "loss": 0.7814,
+      "step": 346
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 1.4129001593840296,
+      "learning_rate": 6.778795582015096e-06,
+      "loss": 0.7175,
+      "step": 347
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 1.8428763282151688,
+      "learning_rate": 6.724377744847013e-06,
+      "loss": 0.8029,
+      "step": 348
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 23.919369190129455,
+      "learning_rate": 6.67006835206512e-06,
+      "loss": 0.6863,
+      "step": 349
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 2.0252490459198977,
+      "learning_rate": 6.615869201663296e-06,
+      "loss": 0.7234,
+      "step": 350
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 1.4586085919502172,
+      "learning_rate": 6.561782087985682e-06,
+      "loss": 0.7848,
+      "step": 351
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 1.3615988764007254,
+      "learning_rate": 6.507808801667275e-06,
+      "loss": 0.7288,
+      "step": 352
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 1.2707870687019411,
+      "learning_rate": 6.453951129574644e-06,
+      "loss": 0.7994,
+      "step": 353
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 1.2789794436108992,
+      "learning_rate": 6.400210854746776e-06,
+      "loss": 0.7892,
+      "step": 354
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 173.76063353258326,
+      "learning_rate": 6.34658975633605e-06,
+      "loss": 0.7521,
+      "step": 355
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 1.451279782269762,
+      "learning_rate": 6.293089609549325e-06,
+      "loss": 0.6994,
+      "step": 356
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 1.353528520599467,
+      "learning_rate": 6.239712185589181e-06,
+      "loss": 0.8124,
+      "step": 357
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 1.3427971146726998,
+      "learning_rate": 6.186459251595267e-06,
+      "loss": 0.6997,
+      "step": 358
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 1.7436977861938026,
+      "learning_rate": 6.133332570585813e-06,
+      "loss": 0.7466,
+      "step": 359
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 1.1504667545933847,
+      "learning_rate": 6.080333901399252e-06,
+      "loss": 0.759,
+      "step": 360
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 1.198911606483234,
+      "learning_rate": 6.02746499863599e-06,
+      "loss": 0.7806,
+      "step": 361
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 1.6129378703199322,
+      "learning_rate": 5.9747276126003265e-06,
+      "loss": 0.7609,
+      "step": 362
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 1.1562370828669843,
+      "learning_rate": 5.922123489242499e-06,
+      "loss": 0.8708,
+      "step": 363
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 1.4562484055998943,
+      "learning_rate": 5.869654370100887e-06,
+      "loss": 0.7431,
+      "step": 364
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 1.0680361234204387,
+      "learning_rate": 5.8173219922443516e-06,
+      "loss": 0.8032,
+      "step": 365
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 0.9284927777737381,
+      "learning_rate": 5.765128088214723e-06,
+      "loss": 0.7889,
+      "step": 366
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 1.1936424602098312,
+      "learning_rate": 5.713074385969457e-06,
+      "loss": 0.8115,
+      "step": 367
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 0.9730264080342912,
+      "learning_rate": 5.66116260882442e-06,
+      "loss": 0.7615,
+      "step": 368
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.0168331506692085,
+      "learning_rate": 5.609394475396821e-06,
+      "loss": 0.7566,
+      "step": 369
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.0193464829148098,
+      "learning_rate": 5.557771699548347e-06,
+      "loss": 0.7821,
+      "step": 370
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.512465787439959,
+      "learning_rate": 5.5062959903283855e-06,
+      "loss": 0.801,
+      "step": 371
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 2.8070018653706734,
+      "learning_rate": 5.454969051917482e-06,
+      "loss": 0.8019,
+      "step": 372
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.29102477160499,
+      "learning_rate": 5.403792583570884e-06,
+      "loss": 0.717,
+      "step": 373
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.0476339722470762,
+      "learning_rate": 5.352768279562315e-06,
+      "loss": 0.7767,
+      "step": 374
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 1.0218859582674915,
+      "learning_rate": 5.3018978291278635e-06,
+      "loss": 0.7388,
+      "step": 375
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 0.8879863478698906,
+      "learning_rate": 5.251182916410061e-06,
+      "loss": 0.7781,
+      "step": 376
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 1.118103740784447,
+      "learning_rate": 5.200625220402139e-06,
+      "loss": 0.7485,
+      "step": 377
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 0.9151353492890569,
+      "learning_rate": 5.150226414892426e-06,
+      "loss": 0.735,
+      "step": 378
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 1.3193616869908815,
+      "learning_rate": 5.099988168408952e-06,
+      "loss": 0.8344,
+      "step": 379
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 0.8632307336970306,
+      "learning_rate": 5.049912144164186e-06,
+      "loss": 0.751,
+      "step": 380
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.9901488846066039,
+      "learning_rate": 5.000000000000003e-06,
+      "loss": 0.7699,
+      "step": 381
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 1.101506870811026,
+      "learning_rate": 4.950253388332763e-06,
+      "loss": 0.7011,
+      "step": 382
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 1.0746169854564547,
+      "learning_rate": 4.900673956098644e-06,
+      "loss": 0.7929,
+      "step": 383
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.8732023653035873,
+      "learning_rate": 4.8512633446990805e-06,
+      "loss": 0.8045,
+      "step": 384
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.6916355479779134,
+      "learning_rate": 4.802023189946454e-06,
+      "loss": 0.7143,
+      "step": 385
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 0.779756473268933,
+      "learning_rate": 4.75295512200992e-06,
+      "loss": 0.7683,
+      "step": 386
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 0.909706557431742,
+      "learning_rate": 4.7040607653614335e-06,
+      "loss": 0.7467,
+      "step": 387
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 0.7722540383691884,
+      "learning_rate": 4.655341738721989e-06,
+      "loss": 0.7815,
+      "step": 388
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 0.9684329548399441,
+      "learning_rate": 4.606799655008009e-06,
+      "loss": 0.7497,
+      "step": 389
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 1.0483349519846141,
+      "learning_rate": 4.558436121277957e-06,
+      "loss": 0.7319,
+      "step": 390
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 0.8099534417037388,
+      "learning_rate": 4.510252738679136e-06,
+      "loss": 0.7647,
+      "step": 391
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 1.1081247822447862,
+      "learning_rate": 4.462251102394669e-06,
+      "loss": 0.7461,
+      "step": 392
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 0.9893574445652334,
+      "learning_rate": 4.414432801590703e-06,
+      "loss": 0.7783,
+      "step": 393
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 1.8652426181110364,
+      "learning_rate": 4.3667994193637794e-06,
+      "loss": 0.7563,
+      "step": 394
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 1.4008789517131635,
+      "learning_rate": 4.319352532688444e-06,
+      "loss": 0.7786,
+      "step": 395
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 0.9945515863927269,
+      "learning_rate": 4.272093712365016e-06,
+      "loss": 0.7605,
+      "step": 396
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 0.7065491597356962,
+      "learning_rate": 4.225024522967602e-06,
+      "loss": 0.8059,
+      "step": 397
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 0.8085670100730566,
+      "learning_rate": 4.178146522792296e-06,
+      "loss": 0.7675,
+      "step": 398
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 0.7493202614166902,
+      "learning_rate": 4.1314612638055766e-06,
+      "loss": 0.6753,
+      "step": 399
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 0.9780288557512398,
+      "learning_rate": 4.084970291592948e-06,
+      "loss": 0.8119,
+      "step": 400
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 0.9642423794575019,
+      "learning_rate": 4.038675145307747e-06,
+      "loss": 0.7209,
+      "step": 401
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 2.805757303861435,
+      "learning_rate": 3.99257735762021e-06,
+      "loss": 0.7739,
+      "step": 402
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.7352143355461409,
+      "learning_rate": 3.946678454666719e-06,
+      "loss": 0.7686,
+      "step": 403
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.9660000138982177,
+      "learning_rate": 3.900979955999271e-06,
+      "loss": 0.7576,
+      "step": 404
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 1.0101800258798939,
+      "learning_rate": 3.855483374535191e-06,
+      "loss": 0.7244,
+      "step": 405
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.6558303005788094,
+      "learning_rate": 3.810190216507017e-06,
+      "loss": 0.814,
+      "step": 406
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.855723897165516,
+      "learning_rate": 3.7651019814126656e-06,
+      "loss": 0.7209,
+      "step": 407
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.8337824890986832,
+      "learning_rate": 3.7202201619657575e-06,
+      "loss": 0.7474,
+      "step": 408
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.8106968198738221,
+      "learning_rate": 3.6755462440462288e-06,
+      "loss": 0.7723,
+      "step": 409
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.7782544106560358,
+      "learning_rate": 3.6310817066511106e-06,
+      "loss": 0.7685,
+      "step": 410
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.9143992276666147,
+      "learning_rate": 3.58682802184558e-06,
+      "loss": 0.7259,
+      "step": 411
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.7313441515558586,
+      "learning_rate": 3.542786654714234e-06,
+      "loss": 0.7734,
+      "step": 412
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.9886120170991388,
+      "learning_rate": 3.4989590633125583e-06,
+      "loss": 0.7262,
+      "step": 413
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 0.9615177391847125,
+      "learning_rate": 3.455346698618689e-06,
+      "loss": 0.7533,
+      "step": 414
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 0.7015094641187799,
+      "learning_rate": 3.4119510044853434e-06,
+      "loss": 0.793,
+      "step": 415
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 0.7446915794929314,
+      "learning_rate": 3.3687734175920505e-06,
+      "loss": 0.7465,
+      "step": 416
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 0.7526032154187763,
+      "learning_rate": 3.325815367397557e-06,
+      "loss": 0.7649,
+      "step": 417
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 1.3172737231291454,
+      "learning_rate": 3.2830782760925285e-06,
+      "loss": 0.7545,
+      "step": 418
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 1.1499109712233648,
+      "learning_rate": 3.2405635585524566e-06,
+      "loss": 0.7566,
+      "step": 419
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 0.9632828455463713,
+      "learning_rate": 3.1982726222908046e-06,
+      "loss": 0.748,
+      "step": 420
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 1.0604062010414916,
+      "learning_rate": 3.1562068674124347e-06,
+      "loss": 0.7856,
+      "step": 421
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 0.8284502574122435,
+      "learning_rate": 3.114367686567228e-06,
+      "loss": 0.6855,
+      "step": 422
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 1.355412483487692,
+      "learning_rate": 3.0727564649040066e-06,
+      "loss": 0.7446,
+      "step": 423
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 0.7929060735164787,
+      "learning_rate": 3.0313745800246508e-06,
+      "loss": 0.7914,
+      "step": 424
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 0.9112317992163422,
+      "learning_rate": 2.9902234019385056e-06,
+      "loss": 0.7512,
+      "step": 425
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 1.5432777353176643,
+      "learning_rate": 2.949304293017029e-06,
+      "loss": 0.7482,
+      "step": 426
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 1.010125898550166,
+      "learning_rate": 2.908618607948669e-06,
+      "loss": 0.7362,
+      "step": 427
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 0.9406215385352928,
+      "learning_rate": 2.8681676936940397e-06,
+      "loss": 0.7671,
+      "step": 428
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 0.9620579229038209,
+      "learning_rate": 2.827952889441302e-06,
+      "loss": 0.7827,
+      "step": 429
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 0.875563723116615,
+      "learning_rate": 2.7879755265618558e-06,
+      "loss": 0.6912,
+      "step": 430
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 0.8120003143081626,
+      "learning_rate": 2.748236928566238e-06,
+      "loss": 0.7559,
+      "step": 431
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 0.820381236475311,
+      "learning_rate": 2.7087384110603156e-06,
+      "loss": 0.755,
+      "step": 432
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 1.0748157041278616,
+      "learning_rate": 2.669481281701739e-06,
+      "loss": 0.781,
+      "step": 433
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 1.0867626411644806,
+      "learning_rate": 2.6304668401566334e-06,
+      "loss": 0.7433,
+      "step": 434
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 1.188744881776241,
+      "learning_rate": 2.5916963780565897e-06,
+      "loss": 0.7639,
+      "step": 435
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 1.0238209184489064,
+      "learning_rate": 2.55317117895588e-06,
+      "loss": 0.7549,
+      "step": 436
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 0.9935203010274218,
+      "learning_rate": 2.514892518288988e-06,
+      "loss": 0.7363,
+      "step": 437
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 0.8010410348003562,
+      "learning_rate": 2.4768616633283747e-06,
+      "loss": 0.7762,
+      "step": 438
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 0.8397648419805517,
+      "learning_rate": 2.43907987314251e-06,
+      "loss": 0.6987,
+      "step": 439
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 0.9413314360541944,
+      "learning_rate": 2.401548398554213e-06,
+      "loss": 0.7732,
+      "step": 440
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 0.9583425145861847,
+      "learning_rate": 2.364268482099218e-06,
+      "loss": 0.7114,
+      "step": 441
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 0.8277369009396708,
+      "learning_rate": 2.327241357985063e-06,
+      "loss": 0.8339,
+      "step": 442
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 1.2734782043124575,
+      "learning_rate": 2.290468252050204e-06,
+      "loss": 0.7192,
+      "step": 443
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 0.7443093962471002,
+      "learning_rate": 2.2539503817234553e-06,
+      "loss": 0.7181,
+      "step": 444
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 2.5435419934304613,
+      "learning_rate": 2.2176889559836655e-06,
+      "loss": 0.7381,
+      "step": 445
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 0.7478928464914032,
+      "learning_rate": 2.1816851753197023e-06,
+      "loss": 0.7657,
+      "step": 446
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 0.9130630498620563,
+      "learning_rate": 2.145940231690713e-06,
+      "loss": 0.8151,
+      "step": 447
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 1.208520757596276,
+      "learning_rate": 2.1104553084866476e-06,
+      "loss": 0.7135,
+      "step": 448
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 0.9346797974039266,
+      "learning_rate": 2.075231580489098e-06,
+      "loss": 0.7467,
+      "step": 449
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 0.9041571142725298,
+      "learning_rate": 2.0402702138323895e-06,
+      "loss": 0.7875,
+      "step": 450
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 1.136292512279568,
+      "learning_rate": 2.0055723659649907e-06,
+      "loss": 0.7891,
+      "step": 451
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 1.4622534646655785,
+      "learning_rate": 1.971139185611176e-06,
+      "loss": 0.7325,
+      "step": 452
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 0.9288453975224165,
+      "learning_rate": 1.936971812733012e-06,
+      "loss": 0.7239,
+      "step": 453
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 0.9248099202923302,
+      "learning_rate": 1.9030713784926115e-06,
+      "loss": 0.7539,
+      "step": 454
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 1.109980406035709,
+      "learning_rate": 1.8694390052146737e-06,
+      "loss": 0.7161,
+      "step": 455
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 1.1303363535074193,
+      "learning_rate": 1.8360758063493511e-06,
+      "loss": 0.7596,
+      "step": 456
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 0.9683620967798374,
+      "learning_rate": 1.8029828864353583e-06,
+      "loss": 0.7729,
+      "step": 457
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 0.9130376128054942,
+      "learning_rate": 1.7701613410634367e-06,
+      "loss": 0.8044,
+      "step": 458
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 1.0497968691825215,
+      "learning_rate": 1.7376122568400533e-06,
+      "loss": 0.7659,
+      "step": 459
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 0.9712735742009617,
+      "learning_rate": 1.705336711351443e-06,
+      "loss": 0.7362,
+      "step": 460
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 1.2095566979767287,
+      "learning_rate": 1.6733357731279375e-06,
+      "loss": 0.7853,
+      "step": 461
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 0.9487731058979595,
+      "learning_rate": 1.641610501608575e-06,
+      "loss": 0.7338,
+      "step": 462
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 0.7578020160286263,
+      "learning_rate": 1.6101619471060415e-06,
+      "loss": 0.6903,
+      "step": 463
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 0.7690765997202936,
+      "learning_rate": 1.5789911507718824e-06,
+      "loss": 0.7635,
+      "step": 464
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 1.002240773181561,
+      "learning_rate": 1.5480991445620541e-06,
+      "loss": 0.7275,
+      "step": 465
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 0.6958769525618868,
+      "learning_rate": 1.5174869512027368e-06,
+      "loss": 0.8401,
+      "step": 466
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 0.8904677751822176,
+      "learning_rate": 1.4871555841564889e-06,
+      "loss": 0.7106,
+      "step": 467
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 0.9090397305558453,
+      "learning_rate": 1.457106047588701e-06,
+      "loss": 0.7718,
+      "step": 468
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 0.8803863054747346,
+      "learning_rate": 1.4273393363343324e-06,
+      "loss": 0.7106,
+      "step": 469
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 0.8690462193016681,
+      "learning_rate": 1.3978564358649926e-06,
+      "loss": 0.7498,
+      "step": 470
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 0.8346189883438835,
+      "learning_rate": 1.368658322256311e-06,
+      "loss": 0.694,
+      "step": 471
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 0.8360259755034392,
+      "learning_rate": 1.339745962155613e-06,
+      "loss": 0.7266,
+      "step": 472
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 1.3601609064643252,
+      "learning_rate": 1.311120312749935e-06,
+      "loss": 0.7349,
+      "step": 473
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 1.4273303270709006,
+      "learning_rate": 1.2827823217343182e-06,
+      "loss": 0.7139,
+      "step": 474
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 0.851993104311948,
+      "learning_rate": 1.2547329272804477e-06,
+      "loss": 0.7842,
+      "step": 475
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.9109512771042799,
+      "learning_rate": 1.2269730580055806e-06,
+      "loss": 0.7818,
+      "step": 476
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.839921234928505,
+      "learning_rate": 1.1995036329418153e-06,
+      "loss": 0.7696,
+      "step": 477
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.814661256852612,
+      "learning_rate": 1.172325561505653e-06,
+      "loss": 0.7705,
+      "step": 478
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.6837807595128806,
+      "learning_rate": 1.1454397434679022e-06,
+      "loss": 0.7494,
+      "step": 479
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.9985810756968562,
+      "learning_rate": 1.1188470689238794e-06,
+      "loss": 0.7093,
+      "step": 480
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.9835361439566525,
+      "learning_rate": 1.0925484182639467e-06,
+      "loss": 0.7514,
+      "step": 481
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 1.4847672648252264,
+      "learning_rate": 1.066544662144371e-06,
+      "loss": 0.7197,
+      "step": 482
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 0.8865210358285674,
+      "learning_rate": 1.040836661458482e-06,
+      "loss": 0.7832,
+      "step": 483
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 0.7826635968526742,
+      "learning_rate": 1.0154252673081943e-06,
+      "loss": 0.7706,
+      "step": 484
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 1.1823763400819176,
+      "learning_rate": 9.903113209758098e-07,
+      "loss": 0.7805,
+      "step": 485
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 0.899288350168961,
+      "learning_rate": 9.65495653896179e-07,
+      "loss": 0.7434,
+      "step": 486
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 0.790636862611698,
+      "learning_rate": 9.409790876291658e-07,
+      "loss": 0.7653,
+      "step": 487
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 1.5778550765050696,
+      "learning_rate": 9.167624338324599e-07,
+      "loss": 0.6827,
+      "step": 488
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 0.762966136049774,
+      "learning_rate": 8.928464942346949e-07,
+      "loss": 0.7544,
+      "step": 489
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 0.7152745679330409,
+      "learning_rate": 8.692320606089088e-07,
+      "loss": 0.7117,
+      "step": 490
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 0.938749881629017,
+      "learning_rate": 8.459199147463371e-07,
+      "loss": 0.7272,
+      "step": 491
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 0.7887243971433457,
+      "learning_rate": 8.229108284305176e-07,
+      "loss": 0.7599,
+      "step": 492
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 1.1187390973902265,
+      "learning_rate": 8.002055634117578e-07,
+      "loss": 0.7813,
+      "step": 493
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 1.0611976127431488,
+      "learning_rate": 7.778048713818975e-07,
+      "loss": 0.7291,
+      "step": 494
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 0.775779064137104,
+      "learning_rate": 7.557094939494325e-07,
+      "loss": 0.7374,
+      "step": 495
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 1.1230838323655812,
+      "learning_rate": 7.33920162614965e-07,
+      "loss": 0.8131,
+      "step": 496
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 0.8286549343306245,
+      "learning_rate": 7.124375987469767e-07,
+      "loss": 0.6944,
+      "step": 497
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 0.8510070990308295,
+      "learning_rate": 6.912625135579587e-07,
+      "loss": 0.7423,
+      "step": 498
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 0.7443707131929154,
+      "learning_rate": 6.703956080808515e-07,
+      "loss": 0.726,
+      "step": 499
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 1.2002747293068283,
+      "learning_rate": 6.498375731458529e-07,
+      "loss": 0.7308,
+      "step": 500
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 1.8988796327176072,
+      "learning_rate": 6.295890893575296e-07,
+      "loss": 0.805,
+      "step": 501
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 0.8806297015318687,
+      "learning_rate": 6.096508270723023e-07,
+      "loss": 0.7082,
+      "step": 502
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 0.8651385101087814,
+      "learning_rate": 5.900234463762367e-07,
+      "loss": 0.7248,
+      "step": 503
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.7449934824918327,
+      "learning_rate": 5.707075970631981e-07,
+      "loss": 0.7295,
+      "step": 504
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.9977941207308976,
+      "learning_rate": 5.517039186133433e-07,
+      "loss": 0.8009,
+      "step": 505
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.930119838929777,
+      "learning_rate": 5.330130401719413e-07,
+      "loss": 0.7334,
+      "step": 506
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.8268138801139723,
+      "learning_rate": 5.146355805285452e-07,
+      "loss": 0.7481,
+      "step": 507
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.8381948856631107,
+      "learning_rate": 4.965721480965136e-07,
+      "loss": 0.7552,
+      "step": 508
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.717461765264355,
+      "learning_rate": 4.788233408928588e-07,
+      "loss": 0.7909,
+      "step": 509
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 1.1030961588444228,
+      "learning_rate": 4.6138974651845914e-07,
+      "loss": 0.7126,
+      "step": 510
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 1.5079778253927174,
+      "learning_rate": 4.4427194213859216e-07,
+      "loss": 0.7532,
+      "step": 511
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 0.9963116969948781,
+      "learning_rate": 4.27470494463843e-07,
+      "loss": 0.6903,
+      "step": 512
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 0.7736026624934934,
+      "learning_rate": 4.109859597313237e-07,
+      "loss": 0.772,
+      "step": 513
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 0.7562011399269606,
+      "learning_rate": 3.9481888368627764e-07,
+      "loss": 0.7279,
+      "step": 514
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 1.2119074253843334,
+      "learning_rate": 3.7896980156399533e-07,
+      "loss": 0.7628,
+      "step": 515
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 0.806901432869074,
+      "learning_rate": 3.6343923807210523e-07,
+      "loss": 0.7608,
+      "step": 516
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 0.7981792877376421,
+      "learning_rate": 3.482277073731988e-07,
+      "loss": 0.7552,
+      "step": 517
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 0.899137343142627,
+      "learning_rate": 3.33335713067805e-07,
+      "loss": 0.7758,
+      "step": 518
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 0.966293413969323,
+      "learning_rate": 3.187637481777284e-07,
+      "loss": 0.7083,
+      "step": 519
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 0.7500728280625207,
+      "learning_rate": 3.045122951297119e-07,
+      "loss": 0.7648,
+      "step": 520
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 0.8343628114364402,
+      "learning_rate": 2.905818257394799e-07,
+      "loss": 0.6902,
+      "step": 521
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 0.9641684433229529,
+      "learning_rate": 2.7697280119610816e-07,
+      "loss": 0.7569,
+      "step": 522
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 1.0501897485858462,
+      "learning_rate": 2.636856720467573e-07,
+      "loss": 0.7323,
+      "step": 523
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 0.8440715396044582,
+      "learning_rate": 2.507208781817638e-07,
+      "loss": 0.8247,
+      "step": 524
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 0.9266257349789722,
+      "learning_rate": 2.3807884882006583e-07,
+      "loss": 0.7063,
+      "step": 525
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 0.6988784112931856,
+      "learning_rate": 2.257600024950013e-07,
+      "loss": 0.7425,
+      "step": 526
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 0.8446132979847808,
+      "learning_rate": 2.1376474704044693e-07,
+      "loss": 0.7789,
+      "step": 527
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 0.8591922340666319,
+      "learning_rate": 2.0209347957732328e-07,
+      "loss": 0.7288,
+      "step": 528
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 1.0926686710744615,
+      "learning_rate": 1.9074658650043764e-07,
+      "loss": 0.7114,
+      "step": 529
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 0.7624493253931196,
+      "learning_rate": 1.7972444346569752e-07,
+      "loss": 0.7459,
+      "step": 530
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 0.7923149172686518,
+      "learning_rate": 1.6902741537767607e-07,
+      "loss": 0.7612,
+      "step": 531
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 0.9137607028579776,
+      "learning_rate": 1.5865585637752644e-07,
+      "loss": 0.7652,
+      "step": 532
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 1.709281268111564,
+      "learning_rate": 1.4861010983126202e-07,
+      "loss": 0.7674,
+      "step": 533
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 0.82294274146925,
+      "learning_rate": 1.3889050831838336e-07,
+      "loss": 0.7802,
+      "step": 534
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 0.9209804385310738,
+      "learning_rate": 1.2949737362087156e-07,
+      "loss": 0.7084,
+      "step": 535
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 0.7554554097180233,
+      "learning_rate": 1.2043101671253553e-07,
+      "loss": 0.7576,
+      "step": 536
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 2.5327776012868326,
+      "learning_rate": 1.1169173774871478e-07,
+      "loss": 0.6961,
+      "step": 537
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.29647899469185,
+      "learning_rate": 1.0327982605634279e-07,
+      "loss": 0.7273,
+      "step": 538
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 0.8544062605377822,
+      "learning_rate": 9.519556012436815e-08,
+      "loss": 0.7458,
+      "step": 539
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.025833225658084,
+      "learning_rate": 8.743920759453739e-08,
+      "loss": 0.7298,
+      "step": 540
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.0753264133055775,
+      "learning_rate": 8.00110252525299e-08,
+      "loss": 0.7683,
+      "step": 541
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.1943324017414,
+      "learning_rate": 7.291125901946027e-08,
+      "loss": 0.7811,
+      "step": 542
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 0.7694222713894205,
+      "learning_rate": 6.61401439437348e-08,
+      "loss": 0.7681,
+      "step": 543
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 1.2370289673252275,
+      "learning_rate": 5.969790419326993e-08,
+      "loss": 0.674,
+      "step": 544
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 0.8002799193353513,
+      "learning_rate": 5.3584753048073756e-08,
+      "loss": 0.8067,
+      "step": 545
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 1.5301982389376738,
+      "learning_rate": 4.7800892893178354e-08,
+      "loss": 0.7349,
+      "step": 546
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 1.5057671436094202,
+      "learning_rate": 4.234651521194844e-08,
+      "loss": 0.7347,
+      "step": 547
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 0.8253550569542444,
+      "learning_rate": 3.7221800579735346e-08,
+      "loss": 0.7457,
+      "step": 548
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 0.9341885237844507,
+      "learning_rate": 3.242691865790071e-08,
+      "loss": 0.7243,
+      "step": 549
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 0.9554329055220271,
+      "learning_rate": 2.796202818819871e-08,
+      "loss": 0.8415,
+      "step": 550
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 0.966722990989304,
+      "learning_rate": 2.3827276987524738e-08,
+      "loss": 0.7286,
+      "step": 551
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 0.7901189093094835,
+      "learning_rate": 2.002280194301487e-08,
+      "loss": 0.722,
+      "step": 552
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 0.9000774799749912,
+      "learning_rate": 1.654872900752169e-08,
+      "loss": 0.7212,
+      "step": 553
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 1.0953002583055267,
+      "learning_rate": 1.340517319543877e-08,
+      "loss": 0.7818,
+      "step": 554
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 0.8753114117581223,
+      "learning_rate": 1.0592238578892577e-08,
+      "loss": 0.7321,
+      "step": 555
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 1.6427409735808294,
+      "learning_rate": 8.110018284304132e-09,
+      "loss": 0.7249,
+      "step": 556
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 0.7993301297504886,
+      "learning_rate": 5.958594489295921e-09,
+      "loss": 0.7669,
+      "step": 557
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 1.0213659754539555,
+      "learning_rate": 4.138038419978507e-09,
+      "loss": 0.7661,
+      "step": 558
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 0.8677144101553231,
+      "learning_rate": 2.648410348592423e-09,
+      "loss": 0.7804,
+      "step": 559
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 0.8598685631934427,
+      "learning_rate": 1.4897595915053242e-09,
+      "loss": 0.7757,
+      "step": 560
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.749563526002078,
+      "learning_rate": 6.621245075910665e-10,
+      "loss": 0.7367,
+      "step": 561
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 1.0165817159651056,
+      "learning_rate": 1.6553249694850616e-10,
+      "loss": 0.7488,
+      "step": 562
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 1.012753603929472,
+      "learning_rate": 0.0,
+      "loss": 0.7666,
+      "step": 563
+    },
+    {
+      "epoch": 1.0,
+      "step": 563,
+      "total_flos": 0.0,
+      "train_loss": 0.7891389744531622,
+      "train_runtime": 46961.4176,
+      "train_samples_per_second": 23.135,
+      "train_steps_per_second": 0.012
+    }
+  ],
+  "logging_steps": 1.0,
+  "max_steps": 563,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 100,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": false,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.0,
+  "train_batch_size": 4,
+  "trial_name": null,
+  "trial_params": null
+}

utils.py

@@ -0,0 +1,212 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+# This file is modified from https://github.com/haotian-liu/LLaVA/
+import os
+import os.path as osp
+
+from huggingface_hub import repo_exists, snapshot_download
+from huggingface_hub.utils import HFValidationError, validate_repo_id
+from transformers import AutoConfig, AutoTokenizer, PretrainedConfig
+
+from .configuration_vila import VILAConfig
+from .constants import MEDIA_TOKENS
+from .tokenizer_utils import infer_stop_tokens
+
+
+def load_tokenizer_then_handle_media_tokens_and_chat_template(
+    model_name_or_path, config: VILAConfig, model_max_length=None
+):
+    # TODO(ligeng): a lot of copy-paste code, refactor to make a single function
+    tokenizer = AutoTokenizer.from_pretrained(
+        osp.join(model_name_or_path, "llm"), padding_side="right", use_fast=True, legacy=False
+    )
+    if model_max_length is not None:
+        tokenizer.model_max_length = model_max_length
+
+    # Load chat template if specified.
+    if getattr(config, "chat_template", None) is not None:
+        print(f"Using chat template: {config.chat_template}")
+        fpath = os.path.join(os.path.dirname(__file__), "chat_templates", f"{config.chat_template}.jinja")
+        if not os.path.exists(fpath):
+            fpath = os.path.join(model_name_or_path, f"{config.chat_template}.jinja")
+        with open(fpath) as fd:
+            chat_template = fd.read()
+        tokenizer.chat_template = chat_template.replace("    ", "").replace("\n", "")
+
+    # Set stop tokens for the tokenizer
+    tokenizer.stop_tokens = infer_stop_tokens(tokenizer)
+    tokenizer.stop_token_ids = tokenizer.convert_tokens_to_ids(tokenizer.stop_tokens)
+
+    # Add media tokens to the tokenizer
+    tokenizer.media_tokens = MEDIA_TOKENS
+    tokenizer.media_token_ids = {}
+    for name, token in MEDIA_TOKENS.items():
+        tokenizer.add_tokens([token], special_tokens=True)
+        tokenizer.media_token_ids[name] = tokenizer.convert_tokens_to_ids(token)
+
+    return tokenizer
+
+
+def get_model_config(config):
+    default_keys = ["llm_cfg", "vision_tower_cfg", "mm_projector_cfg"]
+
+    if hasattr(config, "_name_or_path") and len(config._name_or_path) >= 2:
+        root_path = config._name_or_path
+    else:
+        root_path = config.resume_path
+
+    # download from huggingface
+    if root_path is not None and not osp.exists(root_path):
+        try:
+            valid_hf_repo = repo_exists(root_path)
+        except HFValidationError as e:
+            valid_hf_repo = False
+        if valid_hf_repo:
+            root_path = snapshot_download(root_path)
+
+    return_list = []
+    for key in default_keys:
+        cfg = getattr(config, key, None)
+        if isinstance(cfg, dict):
+            try:
+                return_list.append(os.path.join(root_path, key[:-4]))
+            except:
+                raise ValueError(f"Cannot find resume path in config for {key}!")
+        elif isinstance(cfg, PretrainedConfig):
+            return_list.append(os.path.join(root_path, key[:-4]))
+        elif isinstance(cfg, str):
+            return_list.append(cfg)
+
+    return return_list
+
+
+def get_model_config_fp8(config):
+    default_keys = ["llm_cfg", "vision_tower_cfg", "mm_projector_cfg"]
+
+    if hasattr(config, "_name_or_path") and len(config._name_or_path) >= 2:
+        root_path = config._name_or_path
+    else:
+        root_path = config.resume_path
+
+    # download from huggingface
+    if root_path is not None and not osp.exists(root_path):
+        try:
+            valid_hf_repo = repo_exists(root_path)
+        except HFValidationError as e:
+            valid_hf_repo = False
+        if valid_hf_repo:
+            root_path = snapshot_download(root_path)
+
+    return_list = []
+    for key in default_keys:
+        cfg = getattr(config, key, None)
+        if isinstance(cfg, dict):
+            try:
+                return_list.append(os.path.join(root_path, key[:-4]))
+            except:
+                raise ValueError(f"Cannot find resume path in config for {key}!")
+        elif isinstance(cfg, PretrainedConfig):
+            return_list.append(os.path.join(root_path, key[:-4]))
+        elif isinstance(cfg, str):
+            return_list.append(cfg)
+
+    # fp8_llm
+    key = "fp8_llm_cfg"
+    directory_path = os.path.join(root_path, key[:-4])
+    assert os.path.isdir(directory_path) and os.listdir(
+        directory_path
+    ), "You need to first convert the model weights to FP8 explicitly."
+    return_list.append(directory_path)
+
+    return return_list
+
+
+def get_model_config_fp8(config):
+    default_keys = ["llm_cfg", "vision_tower_cfg", "mm_projector_cfg"]
+
+    if hasattr(config, "_name_or_path") and len(config._name_or_path) >= 2:
+        root_path = config._name_or_path
+    else:
+        root_path = config.resume_path
+
+    # download from huggingface
+    if root_path is not None and not osp.exists(root_path):
+        try:
+            valid_hf_repo = repo_exists(root_path)
+        except HFValidationError as e:
+            valid_hf_repo = False
+        if valid_hf_repo:
+            root_path = snapshot_download(root_path)
+
+    return_list = []
+    for key in default_keys:
+        cfg = getattr(config, key, None)
+        if isinstance(cfg, dict):
+            try:
+                return_list.append(os.path.join(root_path, key[:-4]))
+            except:
+                raise ValueError(f"Cannot find resume path in config for {key}!")
+        elif isinstance(cfg, PretrainedConfig):
+            return_list.append(os.path.join(root_path, key[:-4]))
+        elif isinstance(cfg, str):
+            return_list.append(cfg)
+
+    # fp8_llm
+    key = "fp8_llm_cfg"
+    directory_path = os.path.join(root_path, key[:-4])
+    assert os.path.isdir(directory_path) and os.listdir(
+        directory_path
+    ), "You need to first convert the model weights to FP8 explicitly."
+    return_list.append(directory_path)
+
+    return return_list
+
+
+def is_mm_model(model_path):
+    """
+    Check if the model at the given path is a visual language model.
+
+    Args:
+        model_path (str): The path to the model.
+
+    Returns:
+        bool: True if the model is an MM model, False otherwise.
+    """
+    config = AutoConfig.from_pretrained(model_path)
+    architectures = config.architectures
+    for architecture in architectures:
+        if "llava" in architecture.lower():
+            return True
+    return False
+
+
+def auto_upgrade(config):
+    cfg = AutoConfig.from_pretrained(config)
+    if "llava" in config and "llava" not in cfg.model_type:
+        assert cfg.model_type == "llama"
+        print("You are using newer LLaVA code base, while the checkpoint of v0 is from older code base.")
+        print("You must upgrade the checkpoint to the new code base (this can be done automatically).")
+        confirm = input("Please confirm that you want to upgrade the checkpoint. [Y/N]")
+        if confirm.lower() in ["y", "yes"]:
+            print("Upgrading checkpoint...")
+            assert len(cfg.architectures) == 1
+            setattr(cfg.__class__, "model_type", "llava")
+            cfg.architectures[0] = "LlavaLlamaForCausalLM"
+            cfg.save_pretrained(config)
+            print("Checkpoint upgraded.")
+        else:
+            print("Checkpoint upgrade aborted.")
+            exit(1)