Add files using upload-large-folder tool

catty.py

@@ -0,0 +1,221 @@
+import os
+from typing import List, Tuple
+
+from PIL import Image
+
+from .dynamic_high_resolution import factorize_number
+
+
+def construct_mapping_dict(max_splits: int = 12) -> dict:
+    """Construct a mapping dictionary for the given max_splits.
+
+    Args:
+        max_splits (int, optional): The maximum number of splits.
+            Defaults to 12.
+
+    Returns:
+        dict: A mapping dictionary for the given max_splits.
+    """
+    mapping_dict = {}
+    for i in range(1, max_splits + 1):
+        factor_list = factorize_number(i)
+        for factor in factor_list:
+            ratio = factor[0] / factor[1]
+            if ratio not in mapping_dict:
+                mapping_dict[ratio] = [factor]
+            else:
+                mapping_dict[ratio].append(factor)
+    return mapping_dict
+
+
+def save_image_list(image_list: List[Image.Image], save_folder: str) -> None:
+    """Save a list of images to a folder.
+
+    Args:
+        image_list (List[Image.Image]): A list of images.
+        save_folder (str): The folder to save the images to.
+    """
+    os.makedirs(save_folder, exist_ok=True)
+    for i, image in enumerate(image_list):
+        image.save(os.path.join(save_folder, f'{i}.png'))
+
+
+def resize_to_best_size(image: Image.Image, best_slices: tuple,
+                        width_slices: int, height_slices: int,
+                        sub_image_size: int) -> Image.Image:
+    """Resize an image to the best size for the given number of slices.
+
+    Args:
+        image (Image.Image): The image to resize.
+        best_slices (tuple): The best number of slices for the image.
+        width_slices (int): The number of horizontal slices.
+        height_slices (int): The number of vertical slices.
+        sub_image_size (int): The size of the sub-images.
+
+    Returns:
+        Image.Image: The resized image.
+    """
+    width, height = image.size
+    best_width_slices, best_height_slices = best_slices
+    if width_slices < height_slices:
+        new_image_width = best_width_slices * sub_image_size
+        new_image_height = int(height / width * new_image_width)
+    else:
+        new_image_height = best_height_slices * sub_image_size
+        new_image_width = int(width / height * new_image_height)
+    new_image = image.resize((new_image_width, new_image_height), resample=2)
+    return new_image
+
+
+def compute_strides(height: int, width: int, sub_image_size: int,
+                    slices: Tuple[int, int]) -> Tuple[int, int]:
+    """Compute the strides for the given image size and slices.
+
+    Args:
+        height (int): The height of the image.
+        width (int): The width of the image.
+        sub_image_size (int): The size of the sub-images.
+        slices (Tuple[int, int]): The number of horizontal and vertical slices.
+
+    Returns:
+        Tuple[int, int]: The strides for the given image size and slices.
+    """
+    slice_width, slice_height = slices
+    if slice_width > 1:
+        stride_x = (width - sub_image_size) // (slice_width - 1)
+    else:
+        stride_x = 0
+    if slice_height > 1:
+        stride_y = (height - sub_image_size) // (slice_height - 1)
+    else:
+        stride_y = 0
+    return stride_x, stride_y
+
+
+def sliding_window_crop(image: Image.Image, window_size: int,
+                        slices: Tuple[int, int]) -> List[Image.Image]:
+    """Crop an image into sub-images using a sliding window.
+
+    Args:
+        image (Image.Image): The image to crop.
+        window_size (int): The size of the sub-images.
+        slices (Tuple[int, int]): The number of horizontal and vertical slices.
+
+    Returns:
+        List[Image]: A list of cropped images.
+    """
+    width, height = image.size
+    stride_x, stride_y = compute_strides(height, width, window_size, slices)
+    sub_images = []
+    if stride_x == 0:
+        stride_x = window_size
+
+    if stride_y == 0:
+        stride_y = window_size
+    for y in range(0, height - window_size + 1, stride_y):
+        for x in range(0, width - window_size + 1, stride_x):
+            sub_image = image.crop((x, y, x + window_size, y + window_size))
+            sub_images.append(sub_image)
+    return sub_images
+
+
+def find_best_slices(width_slices: int,
+                     height_slices: int,
+                     aspect_ratio: float,
+                     max_splits: int = 12) -> list:
+    """Find the best slices for the given image size and aspect ratio.
+
+    Args:
+        width_slices (int): The number of horizontal slices.
+        height_slices (int): The number of vertical slices.
+        aspect_ratio (float): The aspect ratio of the image.
+        max_splits (int, optional): The maximum number of splits.
+            Defaults to 12.
+
+    Returns:
+        list: the best slices for the given image.
+    """
+    mapping_dict = construct_mapping_dict(max_splits)
+    if aspect_ratio < 1:
+        mapping_dict = {
+            k: v
+            for k, v in mapping_dict.items() if k <= aspect_ratio
+        }
+    elif aspect_ratio > 1:
+        mapping_dict = {
+            k: v
+            for k, v in mapping_dict.items() if k >= aspect_ratio
+        }
+    # find the value which key is the closest to the ratio
+    best_ratio = min(mapping_dict.keys(), key=lambda x: abs(x - aspect_ratio))
+    # best_image_sizes is a list of image sizes
+    best_image_sizes = mapping_dict[best_ratio]
+    # find the image_size whose area is closest to the current image size
+    best_slices = min(
+        best_image_sizes,
+        key=lambda x: abs(x[0] * x[1] - width_slices * height_slices))
+    return best_slices
+
+
+def split_image_with_catty(pil_image: Image.Image,
+                           image_size: int = 336,
+                           max_crop_slices: int = 8,
+                           save_folder: str = None,
+                           add_thumbnail: bool = True,
+                           do_resize: bool = False,
+                           **kwargs) -> List[Image.Image]:
+    """Split an image into sub-images using Catty.
+
+    Args:
+        pil_image (Image.Image): The image to split.
+        image_size (int, optional): The size of the image.
+            Defaults to 336.
+        max_crop_slices (int, optional): The maximum number of slices.
+            Defaults to 8.
+        save_folder (str, optional): The folder to save the sub-images.
+            Defaults to None.
+        add_thumbnail (bool, optional): Whether to add a thumbnail.
+            Defaults to False.
+        do_resize (bool, optional): Whether to resize the image to fit the
+            maximum number of slices. Defaults to False.
+
+    Returns:
+        List[Image.Image]: A list of cropped images.
+    """
+    width, height = pil_image.size
+    ratio = width / height
+    if ratio > max_crop_slices or ratio < 1 / max_crop_slices:
+        if do_resize:
+            print(
+                f'Resizing image to fit maximum number of slices ({max_crop_slices})'  # noqa
+            )  # noqa
+            if width > height:
+                new_width = max_crop_slices * height
+                new_height = height
+            else:
+                new_width = width
+                new_height = max_crop_slices * width
+            pil_image = pil_image.resize((new_width, new_height), resample=2)
+            width, height = pil_image.size
+            ratio = width / height
+        else:
+            print(
+                f'Image aspect ratio ({ratio:.2f}) is out of range: ({1/max_crop_slices:.2f}, {max_crop_slices:.2f})'  # noqa
+            )
+            return None, None
+    width_slices = width / image_size
+    height_slices = height / image_size
+    best_slices = find_best_slices(width_slices, height_slices, ratio,
+                                   max_crop_slices)
+    pil_image = resize_to_best_size(pil_image, best_slices, width_slices,
+                                    height_slices, image_size)
+    width, height = pil_image.size
+    sub_images = sliding_window_crop(pil_image, image_size, best_slices)
+    if add_thumbnail:
+        thumbnail_image = pil_image.resize((image_size, image_size),
+                                           resample=2)
+        sub_images.append(thumbnail_image)
+    # save split images to folder for debugging
+    if save_folder is not None:
+        save_image_list(sub_images, save_folder)
+    return sub_images

config.json

@@ -0,0 +1,188 @@
+{
+  "_commit_hash": null,
+  "_name_or_path": "/mnt/cephfs/bensenliu/wfs/weights/mm/mmq-llava-20240826e1-sft-points-hf",
+  "architectures": [
+    "POINTSChatModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_points_chat.POINTSChatConfig",
+    "AutoModelForCausalLM": "modeling_points_chat.POINTSChatModel"
+  },
+  "llm_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "auto_map": {
+      "AutoConfig": "configuration_llama.CustomLlamaConfig",
+      "AutoModelForCausalLM": "modeling_llama.CustomLlamaForCausalLM"
+    },
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 0,
+    "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": [
+      2,
+      3
+    ],
+    "exponential_decay_length_penalty": null,
+    "ffn_hidden_size": 11008,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "swiglu",
+    "hidden_size": 4096,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layernorm_epsilon": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 16384,
+    "min_length": 0,
+    "mlp_fc1_bias": false,
+    "mlp_fc2_bias": false,
+    "model_type": "custom_llama",
+    "no_repeat_ngram_size": 0,
+    "norm_type": "rms_norm",
+    "num_attention_heads": 32,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 48,
+    "num_kv_heads": 4,
+    "num_layers": 48,
+    "num_return_sequences": 1,
+    "out_proj_bias": false,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "qkv_proj_bias": false,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rotary_compress": 1.0,
+    "rotary_emb_base": 5000000.0,
+    "rotary_pct": 1.0,
+    "sep_token_id": null,
+    "share_kv_num_layers": 1,
+    "sliding_window_size": -1,
+    "sliding_window_type": 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": null,
+    "torchscript": false,
+    "transformers_version": "4.44.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "use_gqa": true,
+    "vocab_size": 64000
+  },
+  "torch_dtype": "float32",
+  "transformers_version": null,
+  "vision_config": {
+    "_name_or_path": "/mnt/cephfs/bensenliu/exp_runs/weights/mm/CLIP-L-336/clip-vit-large-patch14-336",
+    "add_cross_attention": false,
+    "architectures": [
+      "CLIPVisionModel"
+    ],
+    "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,
+    "dropout": 0.0,
+    "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": "quick_gelu",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 336,
+    "initializer_factor": 1.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "clip_vision_model",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "projection_dim": 768,
+    "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": "float32",
+    "torchscript": false,
+    "transformers_version": "4.44.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  }
+}

configuration_llama.py

@@ -0,0 +1,103 @@
+# Modify the original configuration_llama.py to
+# be compatiable with our training framework.
+
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class CustomLlamaConfig(PretrainedConfig):
+    """
+    Args:
+        vocab_size (`int`, *optional*, defaults to 50432):
+            Vocabulary size of the WeLMV3 model. Defines the number of 
+                different tokens that can be represented by the
+            `inputs_ids` passed when calling [`WeLMV3Model`].
+        hidden_size (`int`, *optional*, defaults to 6144):
+            Dimension of the encoder layers and the pooler layer.
+        num_hidden_layers (`int`, *optional*, defaults to 44):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 64):
+            Number of attention heads for each attention layer in the 
+            Transformer encoder.
+        num_kv_heads (`int`, *optional*, defaults to 4):
+            Number of GQA groups.
+        intermediate_size (`int`, *optional*, defaults to 24576):
+            Dimension of the "intermediate" (i.e., feed-forward) layer in the 
+            Transformer encoder.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the 
+            encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` are supported.
+        rotary_pct (`float`, *optional*, defaults to 0.25):
+            percentage of hidden dimensions to allocate to rotary embeddings
+        rotary_emb_base (`int`, *optional*, defaults to 10000)
+            base for computing rotary embeddings frequency
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used 
+            with. Typically set this to something large just in case 
+            (e.g., 512 or 1024 or 2048).
+        initializer_range (`float`, *optional*, defaults to 1e-5):
+            The standard deviation of the truncated_normal_initializer for 
+            initializing all weight matrices.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values 
+            attentions (not used by all models). Only relevant if 
+            `config.is_decoder=True`.
+    """
+    model_type = "custom_llama"
+
+    def __init__(
+        self,
+        vocab_size=102400,
+        hidden_size=2560,
+        num_layers=32,
+        num_attention_heads=20,
+        num_kv_heads=4,
+        ffn_hidden_size=2560 * 4,
+        hidden_act="swiglu",
+        rotary_pct=1.0,
+        rotary_emb_base=10000,
+        rotary_compress=1.0,
+        max_position_embeddings=4096,
+        initializer_range=0.02,
+        layernorm_epsilon=1e-5,
+        use_cache=True,
+        bos_token_id=0,
+        eos_token_id=2,
+        rms_norm=None,
+        norm_type='layer_norm',
+        qkv_proj_bias=True,
+        out_proj_bias=True,
+        mlp_fc1_bias=True,
+        mlp_fc2_bias=True,
+        **kwargs,
+    ):
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_kv_heads = num_kv_heads
+        self.ffn_hidden_size = ffn_hidden_size
+        self.hidden_act = hidden_act
+        self.rotary_pct = rotary_pct
+        self.rotary_emb_base = rotary_emb_base
+        self.rotary_compress = rotary_compress
+        self.initializer_range = initializer_range
+        self.layernorm_epsilon = layernorm_epsilon
+        self.use_cache = use_cache
+        if rms_norm is not None:
+            self.norm_type = 'rms_norm' if rms_norm else 'layer_norm'
+        else:
+            self.norm_type = norm_type
+        self.qkv_proj_bias = qkv_proj_bias
+        self.out_proj_bias = out_proj_bias
+        self.mlp_fc1_bias = mlp_fc1_bias
+        self.mlp_fc2_bias = mlp_fc2_bias
+        self.num_hidden_layers = num_layers

configuration_points_chat.py

@@ -0,0 +1,39 @@
+import copy
+from typing import Any, Dict, Union
+
+from transformers import CLIPVisionConfig, PretrainedConfig
+
+from .configuration_llama import CustomLlamaConfig
+
+
+class POINTSChatConfig(PretrainedConfig):
+    model_type = "points_chat"
+    is_composition = True
+    """Configuration class for `POINTS`.
+    
+    Args:
+        vision_config (Union[dict, CLIPVisionConfig]): 
+            Configuration of the vision model.
+        llm_config (Union[dict, LlamaConfig]): 
+            Configuration of the language model.
+    """
+
+    def __init__(self,
+                 vision_config: Union[dict, CLIPVisionConfig],
+                 llm_config: Union[dict, CustomLlamaConfig],
+                 **kwargs) -> None:
+        super().__init__(**kwargs)
+        if isinstance(vision_config, dict):
+            self.vision_config = CLIPVisionConfig(**vision_config)
+        else:
+            self.vision_config = vision_config
+        if isinstance(llm_config, dict):
+            self.llm_config = CustomLlamaConfig(**llm_config)
+        else:
+            self.llm_config = llm_config
+
+    def to_dict(self) -> Dict[str, Any]:
+        output = copy.deepcopy(self.__dict__)
+        output["vision_config"] = self.vision_config.to_dict()
+        output["llm_config"] = self.llm_config.to_dict()
+        return output

dynamic_high_resolution.py

@@ -0,0 +1,112 @@
+from typing import List
+
+from PIL import Image
+
+
+def factorize_number(num: int) -> list:
+    """Factorize a number into its prime factors.
+
+    Args:
+        num (int): The number to factorize.
+
+    Returns:
+        list: A list of prime factors of the number.
+    """
+    factors = []
+    for i in range(1, int(num) + 1):
+        if num % i == 0:
+            factors.append([i, num // i])
+    return factors
+
+
+def construct_mapping_dict(max_splits: int = 8, image_size: int = 336) -> dict:
+    """Construct a mapping dictionary for image size reduction.
+
+    Args:
+        max_splits (int, optional): The maximum number of splits for each
+            dimension. Defaults to 8.
+        image_size (int, optional): The original image size.
+            Defaults to 336.
+
+    Returns:
+        dict: A dictionary containing the mapping of image sizes to
+            the corresponding factors.
+    """
+    mapping_dict = {}
+    for i in range(1, max_splits + 1):
+        factor_list = factorize_number(i)
+        for factor in factor_list:
+            ratio = factor[0] / factor[1]
+            if ratio not in mapping_dict:
+                mapping_dict[ratio] = [[
+                    factor[0] * image_size, factor[1] * image_size
+                ]]
+            else:
+                mapping_dict[ratio].append(
+                    [factor[0] * image_size, factor[1] * image_size])
+    return mapping_dict
+
+
+def find_best_image_size(cur_image_size: list,
+                         max_splits: int = 8,
+                         image_size: int = 336) -> list:
+    """Find the best image size for a given image size.
+
+    Args:
+        cur_image_size (list): The current image size.
+        max_splits (int, optional): The maximum number of splits for each
+            dimension. Defaults to 8.
+        image_size (int, optional): The original image size.
+            Defaults to 336.
+
+    Returns:
+        list: The best image size for the given image size.
+    """
+
+    mapping_dict = construct_mapping_dict(max_splits, image_size)
+    ratio = cur_image_size[0] / cur_image_size[1]
+    # find the value which key is the closest to the ratio
+    best_ratio = min(mapping_dict.keys(), key=lambda x: abs(x - ratio))
+    # best_image_sizes is a list of image sizes
+    best_image_sizes = mapping_dict[best_ratio]
+    # find the image_size whose area is closest to the current image size
+    best_image_size = min(
+        best_image_sizes,
+        key=lambda x: abs(x[0] * x[1] - cur_image_size[0] * cur_image_size[1]))
+    return best_image_size
+
+
+def split_image(pil_image: Image.Image,
+                image_size: int = 336,
+                max_splits: int = 8) -> List[Image.Image]:
+    """Split an image into sub-image.
+
+    Similar to that used in InternVL2。
+
+    Args:
+        pil_image (Image.Image): The input image.
+        image_size (int, optional): The size of the image.
+            Defaults to 336.
+        max_splits (int, optional): The maximum number of splits for each
+            dimension. Defaults to 8.
+
+    Returns:
+        List[Image.Image]: A list of cropped images.
+    """
+    whole_sub_image = pil_image.resize((image_size, image_size), resample=2)
+    best_size = find_best_image_size(pil_image.size,
+                                     max_splits=max_splits,
+                                     image_size=image_size)
+    pil_image = pil_image.resize(best_size, resample=2)
+    num_sub_images = ((best_size[0] // image_size),
+                      (best_size[1] // image_size))
+    # crop pil_image to sub_images
+    sub_images = []
+    for i in range(num_sub_images[1]):
+        for j in range(num_sub_images[0]):
+            sub_image = pil_image.crop(
+                (j * image_size, i * image_size, (j + 1) * image_size,
+                 (i + 1) * image_size))
+            sub_images.append(sub_image)
+    sub_images.append(whole_sub_image)
+    return sub_images

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.44.2"
+}

model-00001-of-00008.safetensors

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

model-00002-of-00008.safetensors

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

model-00003-of-00008.safetensors

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

model-00004-of-00008.safetensors

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

model-00005-of-00008.safetensors

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

model-00006-of-00008.safetensors

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

model-00007-of-00008.safetensors

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

model-00008-of-00008.safetensors

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

modeling_llama.py

@@ -0,0 +1,978 @@
+from functools import partial
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from packaging import version
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_llama import CustomLlamaConfig
+
+try:
+    from apex.megatron_layer_norm import MixedFusedLayerNorm as LayerNorm
+except ImportError:
+    from torch.nn import LayerNorm
+
+USE_FLASH_ATTN = False
+try:
+    import flash_attn
+    if version.parse(flash_attn.__version__) >= version.parse("2.1.0"):
+        USE_FLASH_ATTN = True
+        from flash_attn import flash_attn_func, flash_attn_varlen_func
+        from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
+except ImportError:
+    pass
+
+logger = logging.get_logger(__name__)
+
+
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = (attention_mask).sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0,
+                       dtype=torch.torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+def get_norm(config: CustomLlamaConfig):
+    norm_type = config.norm_type
+    if norm_type == 'rms_norm':
+        return partial(RMSNorm, eps=config.layernorm_epsilon)
+    elif norm_type == 'layer_norm':
+        return partial(LayerNorm, eps=config.layernorm_epsilon)
+    else:
+        raise ValueError(f'Unsupported norm type: {norm_type}')
+
+
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+    input_ids_shape: torch.Size,
+    dtype: torch.dtype,
+    device: torch.device,
+    past_key_values_length: int = 0,
+):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len),
+                      torch.finfo(dtype).min, device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat(
+            [
+                torch.zeros(
+                    tgt_len, past_key_values_length, dtype=dtype, device=device
+                ),
+                mask,
+            ],
+            dim=-1,
+        )
+    return mask[None, None, :, :].expand(
+        bsz, 1, tgt_len, tgt_len + past_key_values_length
+    )
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(
+        bsz, 1, tgt_len, src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(
+        inverted_mask.to(torch.bool), torch.finfo(dtype).min
+    )
+
+
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(self, dim, base=10000, compress=1.0):
+        super().__init__()
+        self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.seq_len_cached = 0
+        self.cos_cached = None
+        self.sin_cached = None
+        self.compress = compress
+
+    def forward(self, x, seq_len):
+        if seq_len > self.seq_len_cached:
+            self.seq_len_cached = seq_len
+            self.inv_freq = self.inv_freq.to(x.device)
+            t = (
+                torch.arange(seq_len, device=self.inv_freq.device,
+                             dtype=self.inv_freq.dtype)
+                * self.compress
+            )
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            self.cos_cached = emb.cos()[None, None, :, :]
+            self.sin_cached = emb.sin()[None, None, :, :]
+        return self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]
+
+
+# rotary pos emb helpers:
+
+
+def rotate_half(x):
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+@torch.jit.script
+def apply_rotary_pos_emb(q, k, cos, sin, offset: int = 0):
+    cos, sin = (
+        cos[..., offset: q.shape[-2] + offset, :],
+        sin[..., offset: q.shape[-2] + offset, :],
+    )
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
+def apply_rotary_pos_emb_torch(
+    q, k, cos, sin, offset: int = 0
+):  # jitting fails with bf16
+    cos, sin = (
+        cos[..., offset: q.shape[-2] + offset, :],
+        sin[..., offset: q.shape[-2] + offset, :],
+    )
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
+class CustomLlamaAttention(nn.Module):
+    def __init__(self, config: CustomLlamaConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_kv_heads
+        self.hidden_size = config.hidden_size
+        self.head_size = self.hidden_size // self.num_attention_heads
+        self.rotary_ndims = int(self.head_size * config.rotary_pct)
+        self.max_positions = config.max_position_embeddings
+        self.rotary_emb = RotaryEmbedding(
+            self.rotary_ndims,
+            base=config.rotary_emb_base,
+            compress=config.rotary_compress,
+        )
+        self.norm_factor = torch.sqrt(
+            torch.tensor(self.head_size, dtype=torch.float32)
+        ).to(torch.get_default_dtype())
+
+        if self.use_gqa:
+            self.query_dense = nn.Linear(
+                config.hidden_size,
+                config.hidden_size,
+                bias=getattr(config, "qkv_proj_bias", True)
+            )
+            self.key_value_dense = nn.Linear(
+                config.hidden_size,
+                self.head_size * 2 * config.num_kv_heads,
+                bias=getattr(config, "qkv_proj_bias", True),
+            )
+        else:
+            self.query_key_value = nn.Linear(
+                config.hidden_size,
+                3 * config.hidden_size,
+                bias=getattr(config, "qkv_proj_bias", True),
+            )
+
+        self.dense = nn.Linear(
+            config.hidden_size,
+            config.hidden_size,
+            bias=getattr(config, "out_proj_bias", True),
+        )
+        self.apply_rotary_fn = (
+            apply_rotary_pos_emb_torch
+            if config.torch_dtype == torch.bfloat16
+            else apply_rotary_pos_emb
+        )
+
+    @property
+    def use_gqa(self):
+        return self.num_kv_heads < self.num_attention_heads
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        head_mask=None,
+        layer_past=None,
+        use_cache=False,
+        output_attentions=False,
+    ):
+        has_layer_past = layer_past is not None
+
+        if self.use_gqa:
+            # Compute Q
+            # [batch, seq_len, hidden_size] --> [batch_size, seq_len, (num_heads * head_size)]
+            q = self.query_dense(hidden_states)
+
+            # [batch_size, seq_len, (num_heads * head_size)]
+            #     --> [batch, seq_len, num_attention_heads, head_size]
+            new_q_shape = q.size()[:-1] + \
+                (self.num_attention_heads, self.head_size)
+            q = q.view(*new_q_shape)
+
+            # Compute KV
+            # [batch, seq_len, hidden_size] --> [batch_size, seq_len, (num_attention_groups * 2 * head_size)]
+            kv = self.key_value_dense(hidden_states)
+
+            # [batch, seq_len, (num_attention_groups * 2 * head_size)]
+            #   --> [batch, seq_len, num_attention_groups, 2 * head_size]
+            new_kv_shape = kv.size()[:-1] + (
+                self.num_kv_heads,
+                2 * self.head_size,
+            )
+            kv = kv.view(*new_kv_shape)
+
+            # [batch, num_attention_heads, seq_len, head_size]
+            query = q.permute(0, 2, 1, 3)
+            # [batch, num_attention_groups, seq_len, head_size]
+            key = kv[..., : self.head_size].permute(0, 2, 1, 3)
+            value = kv[..., self.head_size:].permute(0, 2, 1, 3)
+        else:
+            # Compute QKV
+            # Attention heads [batch, seq_len, hidden_size]
+            #   --> [batch, seq_len, (np * 3 * head_size)]
+            qkv = self.query_key_value(hidden_states)
+
+            # [batch, seq_len, (num_heads * 3 * head_size)]
+            #   --> [batch, seq_len, num_heads, 3 * head_size]
+            new_qkv_shape = qkv.size()[:-1] + (
+                self.num_attention_heads,
+                3 * self.head_size,
+            )
+            qkv = qkv.view(*new_qkv_shape)
+
+            # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
+            query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
+            key = qkv[..., self.head_size: 2 *
+                      self.head_size].permute(0, 2, 1, 3)
+            value = qkv[..., 2 * self.head_size:].permute(0, 2, 1, 3)
+
+        # Compute rotary embeddings on rotary_ndims
+        query_rot = query[..., : self.rotary_ndims]
+        query_pass = query[..., self.rotary_ndims:]
+        key_rot = key[..., : self.rotary_ndims]
+        key_pass = key[..., self.rotary_ndims:]
+
+        # Compute token offset for rotary embeddings (when decoding)
+        seq_len = key.shape[-2]
+        offset = 0
+        if has_layer_past:
+            offset = layer_past[0].shape[-2]
+            seq_len += offset
+        cos, sin = self.rotary_emb(value, seq_len=seq_len)
+        query, key = self.apply_rotary_fn(
+            query_rot, key_rot, cos, sin, offset=offset)
+        query = torch.cat((query, query_pass), dim=-1)
+        key = torch.cat((key, key_pass), dim=-1)
+
+        # Cache QKV values
+        if has_layer_past:
+            past_key = layer_past[0]
+            past_value = layer_past[1]
+            key = torch.cat((past_key, key), dim=-2)
+            value = torch.cat((past_value, value), dim=-2)
+        present = (key, value) if use_cache else None
+
+        if USE_FLASH_ATTN:
+            # Compute attention
+            attn_output, attn_weights = self._flash_attn(
+                query, key, value, attention_mask, head_mask
+            )
+
+            # from [batch_size, ]
+            attn_output = attn_output.reshape(
+                attn_output.size(0), attn_output.size(1), self.hidden_size).contiguous()
+        else:
+            # Compute attention
+            attn_output, attn_weights = self._attn(
+                query, key, value, attention_mask, head_mask
+            )
+
+            # Reshape outputs
+            attn_output = self._merge_heads(
+                attn_output, self.num_attention_heads, self.head_size
+            )
+        attn_output = self.dense(attn_output)
+
+        outputs = (attn_output, present)
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+    @classmethod
+    def _split_heads(cls, tensor, num_attention_heads, attn_head_size):
+        """
+        Splits hidden dim into attn_head_size and num_attention_heads
+        """
+        # tensor: [bs, seq_len, hidden_size]
+        new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
+        # -> [bs, seq_len, num_attention_heads, attn_head_size]
+        tensor = tensor.view(new_shape)
+        # -> [bs, num_attention_heads, seq_len, attn_head_size]
+        tensor = tensor.permute(0, 2, 1, 3)
+        return tensor
+
+    @classmethod
+    def _merge_heads(cls, tensor, num_attention_heads, attn_head_size):
+        """
+        Merges attn_head_size dim and num_attn_heads dim into hidden dim
+        """
+        # tensor [bs, num_attention_heads, seq_len, attn_head_size]
+        tensor = tensor.permute(0, 2, 1, 3).contiguous()
+        # -> [bs, seq_len, num_attention_heads, attn_head_size]
+        tensor = tensor.view(
+            tensor.size(0), tensor.size(
+                1), num_attention_heads * attn_head_size
+        )
+        # -> [bs, seq_len, hidden_size]
+        return tensor
+
+    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
+        # q: [bs, num_attention_heads, seq_len, attn_head_size]
+        # k,v: [bs, num_attention_groups, seq_len, attn_head_size]
+        # compute causal mask from causal mask buffer
+        batch_size, num_attention_heads, query_length, attn_head_size = query.size()
+        _, num_attention_groups, key_length, _ = key.size()
+
+        group_size = num_attention_heads // num_attention_groups
+
+        if not self.use_gqa:
+            assert group_size == 1
+
+        # repeat key and value, so we can use normal MHA algorithm
+        key = (
+            key.view(batch_size, num_attention_groups,
+                     1, key_length, attn_head_size)
+            .repeat(1, 1, group_size, 1, 1)
+            .view(batch_size, num_attention_heads, key_length, attn_head_size)
+        )
+        value = (
+            value.view(batch_size, num_attention_groups,
+                       1, key_length, attn_head_size)
+            .repeat(1, 1, group_size, 1, 1)
+            .view(batch_size, num_attention_heads, key_length, attn_head_size)
+        )
+
+        query = query.view(
+            batch_size * num_attention_heads, query_length, attn_head_size
+        )
+        key = key.view(batch_size * num_attention_heads,
+                       key_length, attn_head_size)
+        attn_scores = torch.zeros(
+            batch_size * num_attention_heads,
+            query_length,
+            key_length,
+            dtype=query.dtype,
+            device=key.device,
+        )
+        attn_scores = torch.baddbmm(
+            attn_scores,
+            query,
+            key.transpose(1, 2),
+            beta=1.0,
+            alpha=(
+                torch.tensor(
+                    1.0, dtype=self.norm_factor.dtype, device=self.norm_factor.device
+                )
+                / self.norm_factor
+            ),
+        )
+        attn_scores = attn_scores.view(
+            batch_size, num_attention_heads, query_length, key_length
+        )
+
+        mask_value = torch.finfo(attn_scores.dtype).min
+        # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
+        # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
+        mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(
+            attn_scores.device
+        )
+
+        if attention_mask is not None:
+            # Apply the attention mask
+            attn_scores = attn_scores + attention_mask
+
+        attn_weights = nn.functional.softmax(attn_scores, dim=-1)
+        attn_weights = attn_weights.to(value.dtype)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attn_weights = attn_weights * head_mask
+
+        attn_output = torch.matmul(attn_weights, value)
+        return attn_output, attn_weights
+
+    def _flash_attn(self, query, key, value, attention_mask=None, head_mask=None):
+        assert head_mask is None, "head_mask is not supported in _flash_attn"
+        # q: [bs, num_attention_heads, seq_len, attn_head_size]
+        # k,v: [bs, num_attention_groups, seq_len, attn_head_size]
+
+        # flash_attn need the layout to be [batch_size, sequence_length, num_heads, head_dim]
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        query_length = query.size(1)
+        causal = query_length != 1
+
+        if attention_mask is not None:
+            batch_size = query.size(0)
+            query, key, value, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query, key, value, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query,
+                key,
+                value,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=0,
+                causal=causal,
+            )
+
+            attn_output = pad_input(
+                attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query, key, value, 0, causal=causal
+            )
+
+        return attn_output, None
+
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(
+            attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+        num_attention_heads = query_layer.shape[2]
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len,
+                              num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len,
+                                num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len,
+                                    num_attention_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
+                query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+def swiglu(x):
+    x1, x2 = x.chunk(2, dim=(x.ndim - 1))
+    return x1 * torch.nn.functional.silu(x2)
+
+
+def get_activation(act_name: str):
+    if act_name == "gelu":
+        return ACT2FN["gelu_new"]
+    elif act_name == "swiglu":
+        return swiglu
+    else:
+        return ACT2FN[act_name]
+
+
+class CustomLlamaMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        h_to_4h_out_channels = (
+            config.ffn_hidden_size * 2
+            if config.hidden_act == "swiglu"
+            else config.ffn_hidden_size
+        )
+        self.dense_h_to_4h = nn.Linear(
+            config.hidden_size,
+            h_to_4h_out_channels,
+            bias=getattr(config, "mlp_fc1_bias", True)
+        )
+        self.dense_4h_to_h = nn.Linear(
+            config.ffn_hidden_size,
+            config.hidden_size,
+            bias=getattr(config, "mlp_fc2_bias", True)
+        )
+        self.act = get_activation(config.hidden_act)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense_h_to_4h(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.dense_4h_to_h(hidden_states)
+        return hidden_states
+
+
+class CustomLlamaLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        norm_func = get_norm(config)
+        self.input_layernorm = norm_func(config.hidden_size)
+        self.post_attention_layernorm = norm_func(config.hidden_size)
+        self.attention = CustomLlamaAttention(config)
+        self.mlp = CustomLlamaMLP(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        use_cache=False,
+        layer_past=None,
+        output_attentions=False,
+    ):
+        attn_in = self.input_layernorm(hidden_states)
+        attention_layer_outputs = self.attention(
+            attn_in,
+            attention_mask=attention_mask,
+            layer_past=layer_past,
+            head_mask=head_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+        attn_output = attention_layer_outputs[
+            0
+        ]  # output_attn: attn_output, present, (attn_weights)
+        outputs = attention_layer_outputs[1:]
+        # pseudocode:
+        # x = x + attn(ln1(x))
+        # x = x + mlp(ln2(x))
+        attn_output = attn_output + hidden_states
+        mlp_input = self.post_attention_layernorm(attn_output)
+        mlp_output = self.mlp(mlp_input)
+        hidden_states = mlp_output + attn_output
+
+        if use_cache:
+            outputs = (
+                hidden_states,
+            ) + outputs  # hidden_states, present, (attn_weights)
+        else:
+            # hidden_states, (attn_weights)
+            outputs = (hidden_states,) + outputs[1:]
+
+        return outputs
+
+
+class CustomLlamaPreTrainedModel(PreTrainedModel):
+    config_class = CustomLlamaConfig
+    base_model_prefix = "lm"
+    _no_split_modules = ["CustomLlamaLayer"]
+
+
+class CustomLlamaModel(CustomLlamaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [CustomLlamaLayer(config) for _ in range(config.num_layers)]
+        )
+
+        norm_func = get_norm(config)
+        self.final_layer_norm = norm_func(config.hidden_size)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_in
+
+    def set_input_embeddings(self, value):
+        self.embed_in = value
+
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+    def _prepare_decoder_attention_mask(
+        self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+    ):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+            ).to(inputs_embeds.device)
+            combined_attention_mask = (
+                expanded_attn_mask
+                if combined_attention_mask is None
+                else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        r"""
+        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        """
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError(
+                "You have to specify either input_ids or inputs_embeds")
+
+        batch_size, seq_length = input_shape
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+        else:
+            past_key_values = tuple([None] * self.config.num_layers)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_in(input_ids)
+        # Attention mask.
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past),
+                dtype=torch.bool,
+                device=inputs_embeds.device,
+            )
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_layers x num_heads]
+        # and head_mask is converted to shape [num_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_layers)
+
+        if USE_FLASH_ATTN:
+            attention_mask = attention_mask if (
+                attention_mask is not None and 0 in attention_mask) else None
+        else:
+            attention_mask = self._prepare_decoder_attention_mask(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_key_values_length,
+            )
+
+        hidden_states = inputs_embeds
+        presents = () if use_cache else None
+        all_attentions = () if output_attentions else None
+        all_hidden_states = () if output_hidden_states else None
+        for i, (layer, layer_past) in enumerate(zip(self.layers, past_key_values)):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+            outputs = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                head_mask=head_mask[i],
+                layer_past=layer_past,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+            )
+            hidden_states = outputs[0]
+            if use_cache is True:
+                presents = presents + (outputs[1],)
+            if output_attentions:
+                all_attentions = all_attentions + \
+                    (outputs[2 if use_cache else 1],)
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        # Add last hidden state
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, presents, all_hidden_states, all_attentions]
+                if v is not None
+            )
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=presents,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class CustomLlamaForCausalLM(CustomLlamaPreTrainedModel):
+    _tied_weights_keys = ["embed_out.weight"]
+    _keys_to_ignore_on_load_unexpected = [
+        r"lm.layers.\d+.attention.rotary_emb.inv_freq"
+    ]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.lm = CustomLlamaModel(config)
+        self.embed_out = nn.Linear(
+            config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.embed_out
+
+    def set_output_embeddings(self, new_embeddings):
+        self.embed_out = new_embeddings
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
+            `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. The two additional tensors are
+            only required when the model is used as a decoder in a Sequence to Sequence model.
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+            `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
+            ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+
+        ```"""
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.lm(
+            input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        lm_logits = self.embed_out(hidden_states)
+
+        lm_loss = None
+        if labels is not None:
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shift_logits = lm_logits[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous()
+            loss_fct = CrossEntropyLoss()
+            lm_loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1)
+            )
+
+        if not return_dict:
+            output = (lm_logits,) + outputs[1:]
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=lm_loss,
+            logits=lm_logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **model_kwargs
+    ):
+        input_shape = input_ids.shape
+
+        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(input_shape)
+
+        # cut decoder_input_ids if past is used
+        if past_key_values and past_key_values[0] is not None:
+            input_ids = input_ids[:, -1:]
+
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values
+            }
+        )
+
+        return model_inputs
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(
+                    past_state.index_select(0, beam_idx)
+                    for past_state in layer_past[:2]
+                )
+                + layer_past[2:],
+            )
+        return reordered_past

modeling_points_chat.py

@@ -0,0 +1,236 @@
+from typing import List, Optional, Tuple
+
+import torch
+from PIL import Image
+from torch import nn
+from transformers import (
+    CLIPVisionModel,
+    GenerationMixin,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+from .catty import split_image_with_catty
+from .configuration_points_chat import POINTSChatConfig
+from .dynamic_high_resolution import split_image
+from .modeling_llama import CustomLlamaForCausalLM
+
+
+class POINTSChatModel(PreTrainedModel, GenerationMixin):
+    config_class = POINTSChatConfig
+    _no_split_modules = ["CLIPVisionModel", "LLamaDecoderLayer"]
+    """Chat model for POINTS.
+    
+    Official implementation of the paper "POINTS: Improving Your Vision-language Model with Affordable Strategies"  # noqa: E501
+    paper: https://huggingface.co/papers/2409.04828
+
+    Args:
+        config (PretrainedConfig): The model config.
+    """
+
+    def __init__(self, config: POINTSChatConfig) -> None:
+        super().__init__(config)
+        self.general_vit = CLIPVisionModel(config.vision_config)
+        self.ocr_vit = CLIPVisionModel(config.vision_config)
+        self.llm = CustomLlamaForCausalLM(config.llm_config)
+        self.vision_projector = nn.Sequential(
+            nn.Linear(config.vision_config.hidden_size *
+                      4, config.llm_config.hidden_size),
+            nn.GELU(),
+            nn.Linear(config.llm_config.hidden_size,
+                      config.llm_config.hidden_size)
+
+        )
+
+    def apply_chat_template(self, prompt: str, image_num: int) -> str:
+        """Apply the Yi-1.5-Chat template to the prompt.
+
+        Args:
+            prompt (str): The prompt to apply the template to.
+            image_num (int): The number of the image in the prompt.
+        Returns:
+            str: The prompt with the template applied.
+        """
+        image_tokens = ('<|endoftext|>' * 144) * image_num
+        prompt = f'<|im_start|>user\n{image_tokens}{prompt}<|im_end|>\n<|im_start|>assistant\n'  # noqa: E501
+        return prompt
+
+    def pixel_shuffle(self, feature_map: torch.Tensor,
+                      scale_factor: float = 0.5) -> torch.Tensor:
+        """Implementation of pixel shuffle.
+
+        Merge several patches into a single patch by concatenating
+        them across the channel dimension. Therefore, we can reduce 
+        the image sequence length. In POINTS, we merge 2x2 adjacent
+        patches into a single patch.
+
+        Args:
+            feature_map (torch.Tensor): The feature map to be pixel 
+                shuffled.
+            scale_factor (float, optional): The scale factor for the    
+        """
+
+        # taken from https://huggingface.co/OpenGVLab/InternVL-Chat-V1-5/blob/main/modeling_internvl_chat.py#L187 # noqa
+        n, w, h, c = feature_map.size()
+        # N, W, H, C --> N, W, H * scale, C // scale
+        feature_map = feature_map.view(
+            n, w, int(h * scale_factor), int(c / scale_factor))
+        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
+        feature_map = feature_map.permute(0, 2, 1, 3).contiguous()
+        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
+        feature_map = feature_map.view(
+            n,
+            int(h * scale_factor),
+            int(w * scale_factor),
+            int(c / (scale_factor * scale_factor)),
+        )
+        feature_map = feature_map.permute(0, 2, 1, 3).contiguous()
+        return feature_map
+
+    def extract_image_features(self, images: torch.Tensor,
+                               vision_encoder: str = 'general_vit') -> torch.Tensor:  # noqa: E501
+        """Extract the image features from the vision encoder.
+
+        Args:
+            images (torch.Tensor): The images to extract the features from.
+            vision_encoder (str, optional): The vision encoder to use.
+                Defaults to 'general_vit'.
+
+        Returns:
+            torch.Tensor: The extracted image features.
+        """
+        if vision_encoder == 'general_vit':
+            image_features = self.general_vit(
+                images, output_hidden_states=True
+            )
+        else:
+            image_features = self.ocr_vit(
+                images, output_hidden_states=True
+            )
+        image_features = image_features.hidden_states[-2]
+        image_features = image_features[:, 1:]
+        image_features = image_features.reshape(-1, 24, 24, 1024)
+        image_features = self.pixel_shuffle(image_features, 0.5)
+        image_features = image_features.view(-1, 144, 4096)
+        image_features = self.vision_projector(image_features)
+        return image_features
+
+    def get_pos_mapping(self, pos: List[list]) -> Tuple[dict, int]:
+        """Get the position mapping for the images.
+
+        Args:
+            pos (List[list]): The position of the images in the prompt.
+
+        Returns:
+            Tuple[dict, int]: The position mapping and the 
+            total number of images.
+        """
+        mapping = {}
+        total_images = 0
+        for i, (start, end) in enumerate(pos):
+            num_image = int((end - start) / 144)
+            mapping[i] = num_image
+            total_images += num_image
+        return mapping, total_images
+
+    @torch.no_grad()
+    def chat(self, pixel_values: Image, prompt: str,
+             tokenizer: PreTrainedTokenizer,
+             image_processor, catty: bool = True,
+             generation_config: dict = None,
+             max_splits: int = 8) -> str:
+        """Generate a response to the input prompt.
+
+        Args:
+            pixel_values (Image): The input image.
+            prompt (str): The input prompt.
+            tokenizer (PreTrainedTokenizer): The tokenizer to use.
+            image_processor: The image processor to use.
+            catty (bool, optional): Whether to use catty. Defaults to True.
+            generation_config (dict, optional): The generation config. 
+                Defaults to None.
+            max_splits (int, optional): The maximum number of splits. 
+                Defaults to 8.
+        Returns:
+            str: The generated response.
+        """
+        if catty:
+            cropped_images = split_image_with_catty(pixel_values,
+                                                    do_resize=True,
+                                                    max_crop_slices=max_splits)
+        else:
+            cropped_images = split_image(pixel_values, max_splits=max_splits)
+        prompt = self.apply_chat_template(prompt, len(cropped_images))
+        cropped_images = image_processor.preprocess(
+            cropped_images, return_tensors='pt')['pixel_values']
+        cropped_images = cropped_images.to(self.device)
+        # extract features with general_vit
+        general_vit_features = self.extract_image_features(
+            cropped_images, vision_encoder='general_vit')
+        # extract features with ocr_vit
+        ocr_vit_features = self.extract_image_features(
+            cropped_images, vision_encoder='ocr_vit')
+        image_features = 0.5 * general_vit_features + 0.5 * ocr_vit_features
+        model_inputs = tokenizer(prompt, return_tensors='pt')
+        input_ids = model_inputs['input_ids'].to(self.device)
+        attention_mask = model_inputs['attention_mask'].to(self.device)
+        # stop token
+        eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+        # image token
+        image_token_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
+        generation_config.update(
+            {
+                'eos_token_id': eos_token_id,
+            }
+        )
+        outputs = self.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            image_features=[image_features],
+            image_token_id=image_token_id,
+            **generation_config
+        )
+        response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+        return response
+
+    def generate(self,
+                 input_ids: torch.LongTensor,
+                 attention_mask: torch.LongTensor,
+                 image_features: List[torch.Tensor],
+                 image_token_id: int,
+                 generation_config: Optional[dict] = None,
+                 output_hidden_states: Optional[bool] = None,
+                 return_dict: Optional[bool] = None,
+                 **generate_kwargs) -> torch.LongTensor:
+        input_embeddings = self.llm.lm.embed_in(input_ids)
+        batch_size = input_ids.shape[0]
+        assert len(image_features) == batch_size
+        for i in range(batch_size):
+            special_pos = input_ids[i] == image_token_id
+            pos = (special_pos[:-1] != special_pos[1:]).nonzero() + 1
+            if pos.shape[0] % 2 != 0:
+                # when the sequence is <image><caption>
+                # we need to add a dummy token
+                pos = torch.cat([torch.tensor([[0]]).to(pos.device), pos])
+            pos = pos.reshape(-1, 2).tolist()
+            pos_mapping, total_images = self.get_pos_mapping(pos)
+            assert total_images == len(image_features[i])
+            img_offset = 0
+            for j, (start, end) in enumerate(pos):
+                num_images = pos_mapping[j]
+                input_embeddings[i, start:end] = torch.cat(
+                    [image_features[i][img_offset+k]
+                        for k in range(num_images)],
+                    dim=0
+                )
+                img_offset += num_images
+        outputs = self.llm.generate(
+            inputs_embeds=input_embeddings,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            use_cache=True,
+            **generate_kwargs
+        )
+        return outputs

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 336,
+    "width": 336
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 336
+  }
+}

special_tokens_map.json

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

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,53 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% endif %}{% if system_message is defined %}{{ system_message }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|im_start|>user\\n' + content + '<|im_end|>\\n<|im_start|>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '<|im_end|>' + '\\n' }}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "legacy": true,
+  "model_max_length": 4096,
+  "pad_token": "<unk>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "split_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}