PR Change 2

Signed-off-by: Abhishek <maurya.abhishek@ibm.com>

Author: Abhishek-TAMU

tuning/data/data_preprocessing_utils.py

@@ -69,6 +69,11 @@ def get_data_collator(
     """
 
     if processor:
+        if is_padding_free or packing:
+            raise ValueError(
+                "Vision model tuning does not support packing or padding_free tuning."
+                "Please set packing=False and is_padding_free=False."
+            )
         return VisionDataCollator(processor)
 
     if packing:

tuning/data/setup_dataprocessor.py

@@ -345,8 +345,7 @@ def _process_raw_data_args(
         handlers, dataset_text_field = _get_pretokenized_dataset_handlers(
             data_args, (is_eval_dataset_present and not is_evaldata_tokenized)
         )
-    # TODO: Better way to handle vision this elif condition
-    elif data_args.dataset_text_field and data_args.dataset_image_field:
+    elif processor and data_args.dataset_text_field and data_args.dataset_image_field:
 
         handlers, dataset_text_field = _get_vision_dataset_handlers(
             data_args, processor_kwargs

tuning/sft_trainer.py

@@ -230,8 +230,7 @@ def train(
     processor = None
     try:
         # try to load vision model
-        model_loader = AutoModelForVision2Seq.from_pretrained
-        model = model_loader(
+        model = AutoModelForVision2Seq.from_pretrained(
             model_args.model_name_or_path,
             cache_dir=train_args.cache_dir,
             torch_dtype=get_torch_dtype(model_args.torch_dtype),

tuning/utils/collators.py

@@ -18,18 +18,11 @@
 
 class VisionDataCollator:
     """
-    A data collator specialized for multi-modal (text + image) inputs.
+    A data collator specialized for vision model (text + image) inputs.
     It uses a processor (e.g., LlavaProcessor or MllamaProcessor) to
     combine text and images into model-ready tensors.
 
-    For padding-free tuning, configure the processor's arguments
-    in `processor_kwargs`, for example:
-        processor_kwargs = {
-            "padding": False,
-            "max_length": 1024,
-            ...
-        }
-
+    Padding-free tuning is not supported.
     Args:
         processor: A processor (like `LlavaProcessor`, `MllamaProcessor`, etc.).
     """

tuning/utils/tokenizer_data_utils.py

@@ -107,24 +107,10 @@ def tokenizer_and_embedding_resize(
     embedding_size = int(multiple_of * math.ceil(len(tokenizer) / multiple_of))
     num_new_tokens = num_new_tokens + embedding_size - len(tokenizer)
 
+    # For Mllama models, we need to resize the input and output embeddings
+    # separately, as the model has a different input and output embeddings.
     if isinstance(model, MllamaForConditionalGeneration):
-        # Get new input embedding size
-        current_input_embeddings = model.get_input_embeddings()
-        current_output_embeddings = model.get_output_embeddings()
-        input_embedding_size = current_input_embeddings.weight.shape[0] + (
-            embedding_size - current_output_embeddings.weight.shape[0]
-        )
-
-        # Save current input embedding
-        resized_input_embeddings = model._get_resized_embeddings(
-            current_input_embeddings,
-            new_num_tokens=input_embedding_size,
-            mean_resizing=True,
-        )
-        resized_input_embeddings = copy.deepcopy(resized_input_embeddings)
-        resized_input_embeddings.requires_grad_(
-            current_input_embeddings.weight.requires_grad
-        )
+        resized_input_embeddings = get_resized_input_embeddings(model, embedding_size)
 
         # Resize input and output embeddings
         model.resize_token_embeddings(embedding_size)
@@ -153,3 +139,32 @@ def tokenizer_and_embedding_resize(
         output_embeddings[-num_new_tokens:] = output_embeddings_avg
 
     return {"num_new_tokens": num_new_tokens, "new_embedding_size": embedding_size}
+
+
+def get_resized_input_embeddings(model, embedding_size):
+    """Get resized input embeddings for Mllama models.
+    Args:
+        model: Mllama models.
+        embedding_size: Size of the new embeddings.
+    Returns:
+        resized_input_embeddings: Resized input embeddings.
+    """
+    # Get current input and output embeddings
+    # and their respective vocab sizes
+    current_input_embeddings = model.get_input_embeddings()
+    current_output_embeddings = model.get_output_embeddings()
+    input_embedding_size = current_input_embeddings.weight.shape[0] + (
+        embedding_size - current_output_embeddings.weight.shape[0]
+    )
+
+    # Save current input embedding
+    resized_input_embeddings = model._get_resized_embeddings(
+        current_input_embeddings,
+        new_num_tokens=input_embedding_size,
+        mean_resizing=True,
+    )
+    resized_input_embeddings = copy.deepcopy(resized_input_embeddings)
+    resized_input_embeddings.requires_grad_(
+        current_input_embeddings.weight.requires_grad
+    )
+    return resized_input_embeddings