PR Changes

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

Author: Abhishek-TAMU

tests/data/test_data_preprocessing.py

@@ -20,9 +20,11 @@
 
 # Third Party
 from datasets import Dataset, IterableDataset
-from transformers import AutoTokenizer, DataCollatorForSeq2Seq
+from PIL import Image
+from transformers import AutoProcessor, AutoTokenizer, DataCollatorForSeq2Seq
 from trl import DataCollatorForCompletionOnlyLM
 import datasets
+import numpy as np
 import pyarrow
 import pytest
 import yaml
@@ -62,6 +64,7 @@
 # Local
 from tuning.config import configs
 from tuning.config.acceleration_configs import AttentionAndDistributedPackingConfig
+from tuning.data.collators import VisionDataCollator
 from tuning.data.data_config import DataPreProcessorConfig, DataSetConfig
 from tuning.data.data_preprocessing_utils import get_data_collator
 from tuning.data.data_processors import DataPreProcessor, get_datapreprocessor
@@ -71,6 +74,8 @@
     process_dataargs,
 )
 
+LLAMA_VISION_MODEL_NAME = "tests/artifacts/tiny-llama-vision-model"
+
 
 @pytest.mark.parametrize(
     "datafile, column_names",
@@ -1831,3 +1836,51 @@ def test_get_processed_dataset(datafile, datasetconfigname):
             "train_dataset",
         )
         assert len(os.listdir(train_dataset_dir)) == num_dataset_shards
+
+
+def test_vision_data_collator():
+    """Test the VisionDataCollator with dummy Image data."""
+
+    processor = AutoProcessor.from_pretrained(LLAMA_VISION_MODEL_NAME)
+    collator = VisionDataCollator(processor)
+    processor_kwargs = {}
+    processor_kwargs["return_tensors"] = "pt"
+    processor_kwargs["padding"] = True
+    image_size = (32, 32)
+
+    def generate_pil_image(size=image_size):
+        """Generate a dummy image array of the specified size and return PIL Image."""
+        image_array = np.random.randint(0, 256, size=(*size, 3), dtype=np.uint8)
+        return Image.fromarray(image_array)
+
+    image1 = generate_pil_image()
+    image2 = generate_pil_image()
+
+    features = [
+        {
+            "processor_kwargs": processor_kwargs,
+            "fields_name": {
+                "dataset_text_field": "text",
+                "dataset_image_field": "image",
+            },
+            "text": "Describe the image.",
+            "image": [image1],
+        },
+        {
+            "processor_kwargs": processor_kwargs,
+            "fields_name": {
+                "dataset_text_field": "text",
+                "dataset_image_field": "image",
+            },
+            "text": "What is in the image?",
+            "image": [image2],
+        },
+    ]
+
+    # Call the collator which returns a batch dictionary containing "input_ids" and "labels"
+    batch = collator(features)
+
+    assert "input_ids" in batch
+    assert "labels" in batch
+    assert "attention_mask" in batch
+    assert batch["input_ids"].shape == batch["labels"].shape

tests/test_sft_trainer.py

@@ -27,10 +27,7 @@
 
 # Third Party
 from datasets.exceptions import DatasetGenerationError, DatasetNotFoundError
-from PIL import Image
-from transformers import AutoProcessor
 from transformers.trainer_callback import TrainerCallback
-import numpy as np
 import pytest
 import torch
 import transformers
@@ -93,7 +90,6 @@
     DataHandlerType,
     add_tokenizer_eos_token,
 )
-from tuning.utils.collators import VisionDataCollator
 from tuning.utils.import_utils import is_fms_accelerate_available
 
 MODEL_ARGS = configs.ModelArguments(
@@ -127,7 +123,6 @@
 )
 
 PEFT_LORA_ARGS = peft_config.LoraConfig(r=8, lora_alpha=32, lora_dropout=0.05)
-LLAMA_VISION_MODEL_NAME = "tests/artifacts/tiny-llama-vision-model"
 
 
 @pytest.mark.parametrize(
@@ -1974,51 +1969,3 @@ def test_handler(element, tokenizer, **kwargs):
             },
         )
         _validate_training(tempdir)
-
-
-def test_vision_data_collator():
-    """Test the VisionDataCollator with dummy Image data."""
-
-    processor = AutoProcessor.from_pretrained(LLAMA_VISION_MODEL_NAME)
-    collator = VisionDataCollator(processor)
-    processor_kwargs = {}
-    processor_kwargs["return_tensors"] = "pt"
-    processor_kwargs["padding"] = True
-    image_size = (32, 32)
-
-    def generate_pil_image(size=image_size):
-        """Generate a dummy image array of the specified size and return PIL Image."""
-        image_array = np.random.randint(0, 256, size=(*size, 3), dtype=np.uint8)
-        return Image.fromarray(image_array)
-
-    image1 = generate_pil_image()
-    image2 = generate_pil_image()
-
-    features = [
-        {
-            "processor_kwargs": processor_kwargs,
-            "fields_name": {
-                "dataset_text_field": "text",
-                "dataset_image_field": "image",
-            },
-            "text": "Describe the image.",
-            "image": [image1],
-        },
-        {
-            "processor_kwargs": processor_kwargs,
-            "fields_name": {
-                "dataset_text_field": "text",
-                "dataset_image_field": "image",
-            },
-            "text": "What is in the image?",
-            "image": [image2],
-        },
-    ]
-
-    # Call the collator which returns a batch dictionary containing "input_ids" and "labels"
-    batch = collator(features)
-
-    assert "input_ids" in batch
-    assert "labels" in batch
-    assert "attention_mask" in batch
-    assert batch["input_ids"].shape == batch["labels"].shape

tests/utils/test_embedding_resize.py

@@ -235,7 +235,14 @@ def test_resize_llama_vision_model():
     assert "<unk>" in tokenizer.get_vocab()
     assert resize_result["num_new_tokens"] == 1
 
-    # 2 new tokens were added: <unk> and <image>
+    # Resizing adds 2 tokens (<unk> and <image>) because the tokenizer vocab size (128257)
+    # is 1 greater than the output embedding size (128256),
+    # i.e., len(tokenizer) == model.get_output_embeddings().weight.shape[0] + 1.
+
+    # When special_tokens_dict only contains <unk>, the embedding size calculation
+    # increases the embedding size from 128256 to 128258 (adding 2 tokens in total).
+    # Consequently, the model's input embeddings are resized with an increase of 2 tokens as well.
+
     assert (
         resized_output_embeddings.weight.shape[0]
         == current_output_embeddings.weight.shape[0] + 2

tuning/utils/collators.py tuning/data/collators.pytuning/utils/collators.py renamed to tuning/data/collators.py

@@ -26,7 +26,7 @@
 
 # Local
 from tuning.config import configs
-from tuning.utils.collators import VisionDataCollator
+from tuning.data.collators import VisionDataCollator
 
 logger = logging.getLogger(__name__)
 

tuning/data/data_preprocessing_utils.py

@@ -252,15 +252,15 @@ def _get_vision_dataset_handlers(data_args, processor_kwargs):
     )
 
     # Second data handler configuration
-    fn_kwargs2 = {
+    handler_fn_kwargs2 = {
         "fields_name": {
             "dataset_text_field": data_args.dataset_text_field,
             "dataset_image_field": data_args.dataset_image_field,
         },
         "processor_kwargs": processor_kwargs,
     }
     kwargs2 = {
-        "fn_kwargs": fn_kwargs2,
+        "fn_kwargs": handler_fn_kwargs2,
     }
     handlers.append(
         DataHandlerConfig("prepare_multimodal_data_processor", arguments=kwargs2)