[PixtralLarge] Update Pixtral conversion script to support large format!

src/transformers/models/llava/configuration_llava.py

@@ -50,6 +50,8 @@ class LlavaConfig(PretrainedConfig):
             The index of the layer to select the vision feature.
         image_seq_length (`int`, *optional*, defaults to 576):
             Sequence length of one image embedding.
+        multimodal_projector_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in the multimodal projector.
 
     Example:
 
@@ -85,6 +87,7 @@ def __init__(
         vision_feature_select_strategy="default",
         vision_feature_layer=-2,
         image_seq_length=576,
+        multimodal_projector_bias=True,
         **kwargs,
     ):
         self.ignore_index = ignore_index
@@ -127,6 +130,7 @@ def __init__(
             text_config = CONFIG_MAPPING["llama"]()
 
         self.text_config = text_config
+        self.multimodal_projector_bias = multimodal_projector_bias
 
         super().__init__(**kwargs)
 

src/transformers/models/llava/modeling_llava.py

@@ -86,10 +86,13 @@ class LlavaCausalLMOutputWithPast(ModelOutput):
 class LlavaMultiModalProjector(nn.Module):
     def __init__(self, config: LlavaConfig):
         super().__init__()
-
-        self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
         self.act = ACT2FN[config.projector_hidden_act]
-        self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_2 = nn.Linear(
+            config.text_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
 
     def forward(self, image_features):
         hidden_states = self.linear_1(image_features)

src/transformers/models/llava_next/configuration_llava_next.py

@@ -55,6 +55,8 @@ class LlavaNextConfig(PretrainedConfig):
             Whether the model's input and output word embeddings should be tied.
         image_seq_length (`int`, *optional*, defaults to 576):
             Sequence length of one image embedding.
+        multimodal_projector_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in the multimodal projector.
 
     Example:
 
@@ -92,12 +94,14 @@ def __init__(
         image_grid_pinpoints=None,
         tie_word_embeddings=False,
         image_seq_length=576,
+        multimodal_projector_bias=True,
         **kwargs,
     ):
         self.ignore_index = ignore_index
         self.image_token_index = image_token_index
         self.projector_hidden_act = projector_hidden_act
         self.image_seq_length = image_seq_length
+        self.multimodal_projector_bias = multimodal_projector_bias
 
         if vision_feature_select_strategy not in ["default", "full"]:
             raise ValueError(

src/transformers/models/llava_next/modeling_llava_next.py

@@ -194,10 +194,13 @@ class LlavaNextCausalLMOutputWithPast(ModelOutput):
 class LlavaNextMultiModalProjector(nn.Module):
     def __init__(self, config: LlavaNextConfig):
         super().__init__()
-
-        self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
         self.act = ACT2FN[config.projector_hidden_act]
-        self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_2 = nn.Linear(
+            config.text_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
 
     def forward(self, image_features):
         hidden_states = self.linear_1(image_features)

src/transformers/models/llava_next_video/configuration_llava_next_video.py

@@ -44,6 +44,8 @@ class LlavaNextVideoConfig(PretrainedConfig):
             The image token index to encode the image prompt.
         projector_hidden_act (`str`, *optional*, defaults to `"gelu"`):
             The activation function used by the multimodal projector.
+        multimodal_projector_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in the multimodal projector.
         vision_feature_select_strategy (`str`, *optional*, defaults to `"default"`):
             The feature selection strategy used to select the vision feature from the vision backbone.
             Can be one of `"default"` or `"full"`. If `"default"`, the CLS token is removed from the vision features.
@@ -95,6 +97,7 @@ def __init__(
         ignore_index=-100,
         image_token_index=32001,
         projector_hidden_act="gelu",
+        multimodal_projector_bias=True,
         vision_feature_select_strategy="default",
         vision_feature_layer=-2,
         image_grid_pinpoints=None,
@@ -114,6 +117,7 @@ def __init__(
         self.ignore_index = ignore_index
         self.image_token_index = image_token_index
         self.projector_hidden_act = projector_hidden_act
+        self.multimodal_projector_bias = multimodal_projector_bias
 
         if vision_feature_select_strategy not in ["default", "full"]:
             raise ValueError(

src/transformers/models/llava_next_video/modeling_llava_next_video.py

@@ -179,10 +179,13 @@ def _init_weights(self, module):
 class LlavaNextVideoMultiModalProjector(nn.Module):
     def __init__(self, config: LlavaNextVideoConfig):
         super().__init__()
-
-        self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
         self.act = ACT2FN[config.projector_hidden_act]
-        self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_2 = nn.Linear(
+            config.text_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
 
     def forward(self, image_features):
         hidden_states = self.linear_1(image_features)

src/transformers/models/llava_next_video/modular_llava_next_video.py

@@ -58,6 +58,8 @@ class LlavaNextVideoConfig(PretrainedConfig):
             The image token index to encode the image prompt.
         projector_hidden_act (`str`, *optional*, defaults to `"gelu"`):
             The activation function used by the multimodal projector.
+        multimodal_projector_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in the multimodal projector.
         vision_feature_select_strategy (`str`, *optional*, defaults to `"default"`):
             The feature selection strategy used to select the vision feature from the vision backbone.
             Can be one of `"default"` or `"full"`. If `"default"`, the CLS token is removed from the vision features.
@@ -109,6 +111,7 @@ def __init__(
         ignore_index=-100,
         image_token_index=32001,
         projector_hidden_act="gelu",
+        multimodal_projector_bias=True,
         vision_feature_select_strategy="default",
         vision_feature_layer=-2,
         image_grid_pinpoints=None,
@@ -128,6 +131,7 @@ def __init__(
         self.ignore_index = ignore_index
         self.image_token_index = image_token_index
         self.projector_hidden_act = projector_hidden_act
+        self.multimodal_projector_bias = multimodal_projector_bias
 
         if vision_feature_select_strategy not in ["default", "full"]:
             raise ValueError(

src/transformers/models/llava_onevision/configuration_llava_onevision.py

@@ -58,6 +58,8 @@ class LlavaOnevisionConfig(PretrainedConfig):
             of the form `(height, width)`.
         tie_word_embeddings (`bool`, *optional*, defaults to `False`):
             Whether the model's input and output word embeddings should be tied.
+        multimodal_projector_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in the multimodal projector.
 
     Example:
 
@@ -95,11 +97,13 @@ def __init__(
         vision_aspect_ratio="anyres_max_9",
         image_grid_pinpoints=None,
         tie_word_embeddings=False,
+        multimodal_projector_bias=True,
         **kwargs,
     ):
         self.image_token_index = image_token_index
         self.video_token_index = video_token_index
         self.projector_hidden_act = projector_hidden_act
+        self.multimodal_projector_bias = multimodal_projector_bias
 
         if vision_feature_select_strategy not in ["default", "full"]:
             raise ValueError(

src/transformers/models/llava_onevision/modeling_llava_onevision.py

@@ -201,10 +201,13 @@ class LlavaOnevisionCausalLMOutputWithPast(ModelOutput):
 class LlavaOnevisionMultiModalProjector(nn.Module):
     def __init__(self, config: LlavaOnevisionConfig):
         super().__init__()
-
-        self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
         self.act = ACT2FN[config.projector_hidden_act]
-        self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_2 = nn.Linear(
+            config.text_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+        )
 
     def forward(self, image_features):
         hidden_states = self.linear_1(image_features)

src/transformers/models/pixtral/convert_pixtral_weights_to_hf.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import json
+import os
 
 import regex as re
 import torch
@@ -27,7 +29,6 @@
     PixtralImageProcessor,
     PixtralProcessor,
     PixtralVisionConfig,
-    PreTrainedTokenizerFast,
 )
 from transformers.convert_slow_tokenizer import bytes_to_unicode
 
@@ -156,29 +157,18 @@ def converted(self) -> Tokenizer:
         return tokenizer
 
 
-def convert_mistral_tokenizer():
-    model_name = "mistralai/Pixtral-12B-2409"
+def convert_mistral_tokenizer(model_file):
+    from transformers import LlamaTokenizer
 
-    tokenizer = MistralTokenizer.from_model(model_name)
-
-    vocab = tokenizer.instruct_tokenizer.tokenizer._tekken_token2id_nospecial
-    all_special = [
-        token.value if hasattr(token, "value") else token
-        for token in tokenizer.instruct_tokenizer.tokenizer._all_special_tokens
-    ]
-    specials_tokens = {token: all_special.index(token) for token in all_special}
-    specials_tokens.update(vocab)
-    vocab = specials_tokens
-
-    tokenizer = PreTrainedTokenizerFast(
-        tokenizer_object=MistralConverter(vocab=vocab, additional_special_tokens=all_special).converted(),
-        bos_token="<s>",
-        unk_token="<unk>",
-        eos_token="</s>",
-    )
-    tokenizer.model_input_names = ["input_ids", "attention_mask"]
-
-    return tokenizer
+    mistral_tokenizer = MistralTokenizer.from_file(model_file)
+    vocab = mistral_tokenizer.instruct_tokenizer.tokenizer.vocab()
+    control_token_ids = mistral_tokenizer.instruct_tokenizer.tokenizer._control_tokens
+    all_special = [vocab[id] for id in control_token_ids]
+    hf_tokenizer = LlamaTokenizer(model_file)
+    # Do I need to exclude tokens that are already special?
+    hf_tokenizer.add_special_tokens({"additional_special_tokens": all_special})
+    hf_tokenizer.model_input_names = ["input_ids", "attention_mask"]
+    return hf_tokenizer
 
 
 def permute_for_rope(value, n_heads, config):
@@ -208,7 +198,6 @@ def convert_dictionnary(original_state_dict, vision_config, text_config):
                 num_attention_heads = _config.num_attention_heads
             if "k_proj" in new_key:
                 num_attention_heads = _config.num_key_value_heads
-            # convert the text model (basically mistral model)
 
         if "q_proj" in new_key or "k_proj" in new_key:
             value = permute_for_rope(value, num_attention_heads, _config)
@@ -217,68 +206,111 @@ def convert_dictionnary(original_state_dict, vision_config, text_config):
     return new_dict
 
 
-def convert_mistral_model(input_dir, output_dir):
-    text_config = MistralConfig(
-        attention_dropout=0.0,
-        bos_token_id=1,
-        eos_token_id=2,
-        head_dim=128,
-        hidden_act="silu",
-        hidden_size=5120,
-        initializer_range=0.02,
-        intermediate_size=14336,
-        max_position_embeddings=1024000,
-        model_type="mistral",
-        num_attention_heads=32,
-        num_hidden_layers=40,
-        num_key_value_heads=8,
-        rms_norm_eps=1e-05,
-        rope_theta=1000000000.0,
-        sliding_window=None,
-        tie_word_embeddings=False,
-        vocab_size=131072,
-    )
+MISTRAL_CONFIG_MAPPING = {
+    "dim": "hidden_size",
+    "hidden_dim": "intermediate_size",
+    "n_kv_heads": "num_key_value_heads",
+    "n_heads": "num_attention_heads",
+    "n_layers": "num_hidden_layers",
+}
+
 
-    vision_config = PixtralVisionConfig()
+def convert_mistral_model(input_dir, output_dir):
+    vision_config = {}
+    if os.path.isfile(f"{input_dir}/params.json"):
+        with open(f"{input_dir}/params.json") as f:
+            param_json = json.load(f)
+        vision_config = param_json.pop("vision_encoder")
+        for k, v in MISTRAL_CONFIG_MAPPING.items():
+            value = param_json.pop(k)
+            param_json[v] = value
+        if "hidden_act" not in vision_config:
+            vision_config["hidden_act"] = "silu"
+        text_config = MistralConfig(
+            **param_json,
+            hidden_act="silu",
+            sliding_window=None,
+            tie_word_embeddings=False,
+            is_composition=True,
+            rms_norm_eps=1e-5,
+        )
+    else:
+        text_config = MistralConfig(
+            attention_dropout=0.0,
+            bos_token_id=1,
+            eos_token_id=2,
+            head_dim=128,
+            hidden_act="silu",
+            hidden_size=5120,
+            initializer_range=0.02,
+            intermediate_size=14336,
+            max_position_embeddings=1024000,
+            model_type="mistral",
+            num_attention_heads=32,
+            num_hidden_layers=40,
+            num_key_value_heads=8,
+            rms_norm_eps=1e-05,
+            rope_theta=1000000000.0,
+            sliding_window=None,
+            tie_word_embeddings=False,
+            vocab_size=131072,
+        )
+    adapter_bias = vision_config.pop("adapter_bias", True)
+    vision_config = PixtralVisionConfig(**vision_config)
     config = LlavaConfig(
         vision_config,
         text_config,
         vision_feature_layer=-1,
         image_token_index=10,
         vision_feature_select_strategy="full",
         image_seq_length=1,
+        multimodal_projector_bias=adapter_bias,
     )
     config.architectures = ["LlavaForConditionalGeneration"]
     config.save_pretrained(output_dir)
-
-    original_state_dict = safe_load_file(f"{input_dir}/consolidated.safetensors")
-    new_dict = convert_dictionnary(original_state_dict, vision_config, text_config)
-
+    full_original_state_dict = {}
+    safetensors_files = sorted([file for file in os.listdir(input_dir) if file.endswith(".safetensors")])
+    if len(safetensors_files) == 1:
+        full_original_state_dict = safe_load_file(f"{input_dir}/consolidated.safetensors")
+    else:
+        for file in safetensors_files:
+            loaded_dict = safe_load_file(f"{input_dir}/{file}")
+            full_original_state_dict.update(loaded_dict)
+
+    new_dict = convert_dictionnary(full_original_state_dict, vision_config, text_config)
     with torch.device("meta"):
         model = LlavaForConditionalGeneration(config)
     model.load_state_dict(new_dict, strict=True, assign=True)
-
     model.save_pretrained(output_dir)
 
-    tokenizer = convert_mistral_tokenizer()
-    image_processor = PixtralImageProcessor()
-    processor = PixtralProcessor(tokenizer=tokenizer, image_processor=image_processor, image_token="[IMG]")
-    processor.save_pretrained(output_dir)
-
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--input_dir",
         help="Location of LLaMA weights, which contains tokenizer.model and model folders",
+        required=True,
     )
     parser.add_argument(
         "--output_dir",
         help="Location to write HF model and tokenizer",
+        required=True,
+    )
+    parser.add_argument("--tokenizer_file", help="Location of the specific tokenizer model file to use.", required=True)
+    parser.add_argument(
+        "--chat_template_file",
+        help="Optional file containing a raw chat template. Will be set as the processor's chat template.",
+        required=False,
     )
 
     args = parser.parse_args()
     convert_mistral_model(args.input_dir, args.output_dir)
+    tokenizer = convert_mistral_tokenizer(args.tokenizer_file)
+    image_processor = PixtralImageProcessor()
+    processor = PixtralProcessor(tokenizer=tokenizer, image_processor=image_processor, image_token="[IMG]")
+    if args.chat_template_file:
+        processor.chat_template = open(args.chat_template_file).read()
+    processor.save_pretrained(args.output_dir)
 
 
 if __name__ == "__main__":