udpates + adding epochs

Author: Goekdeniz-Guelmez

mlx_vlm/lora.py

@@ -5,7 +5,6 @@
 import mlx.optimizers as optim
 from datasets import load_dataset
 
-from .prompt_utils import apply_chat_template
 from .trainer import Dataset, TrainingArgs, Colors, train, print_trainable_parameters
 from .trainer.utils import (
     apply_lora_layers,
@@ -61,31 +60,6 @@ def transform_to_messages(examples):
         else:
             raise ValueError(f"{Colors.FAIL}Dataset must have a 'messages' column or both 'question' and 'answer' columns{Colors.ENDC}")
 
-    if args.apply_chat_template:
-        logger.info(f"{Colors.OKBLUE}Applying chat template to the dataset{Colors.ENDC}")
-        
-        def process_data(examples):
-            if config["model_type"] == "pixtral":
-                conversations = apply_chat_template(
-                    config=config,
-                    processor=processor,
-                    prompt=examples["messages"],
-                    return_messages=True,
-                )
-                examples["messages"] = [
-                    json.dumps(item, ensure_ascii=False) for item in conversations
-                ]
-            else:
-                examples["messages"] = apply_chat_template(
-                    config=config,
-                    processor=processor,
-                    prompt=examples["messages"],
-                    return_messages=True,
-                )
-            return examples
-        
-        dataset = dataset.map(process_data)
-    
     # Create Dataset objects
     train_dataset = Dataset(
         dataset,
@@ -95,6 +69,14 @@ def process_data(examples):
         image_resize_shape=args.image_resize_shape,
     )
 
+    if args.epochs is not None:
+        dataset_size = len(train_dataset)
+        steps_per_epoch = dataset_size // args.batch_size
+        total_steps = steps_per_epoch * args.epochs
+        iters = total_steps
+    else:
+        iters = args.iters
+    
     # Use train dataset for validation if no validation dataset is provided
     val_dataset = None
 
@@ -144,7 +126,7 @@ def process_data(examples):
     # Create TrainingArgs
     training_args = TrainingArgs(
         batch_size=args.batch_size,
-        iters=args.iters,
+        iters=iters,
         steps_per_report=args.steps_per_report,
         steps_per_eval=args.steps_per_eval,
         steps_per_save=args.steps_per_save,
@@ -232,6 +214,12 @@ def process_data(examples):
     parser.add_argument(
         "--iters", type=int, default=1000, help="Number of iterations to train for"
     )
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=None,
+        help="Number of epochs to train for. If provided, overrides --iters and computes steps from dataset size and batch size.",
+    )
     parser.add_argument(
         "--steps-per-report", type=int, default=10, help="Number of training steps between loss reporting"
     )

mlx_vlm/trainer/datasets.py

@@ -2,26 +2,6 @@
 import json
 
 
-def get_prompt(model_type, processor, conversation):
-    if model_type == "paligemma":
-        return conversation
-    
-    if "chat_template" in processor.__dict__.keys():
-        prompt = processor.apply_chat_template(
-            conversation,
-            tokenize=False,
-            add_generation_prompt=False,
-        )
-    elif "tokenizer" in processor.__dict__.keys():
-        prompt = processor.tokenizer.apply_chat_template(
-            conversation,
-            tokenize=False,
-            add_generation_prompt=False,
-        )
-    
-    return prompt
-
-
 class Dataset:
     def __init__(
         self,
@@ -53,9 +33,21 @@ def __getitem__(self, idx):
         item = self.dataset[idx]
 
         images = item.get("images", item.get("image", None))
-        conversations = item.get("messages", item.get("conversations"))
-        if images in (None, "", []):
+
+        if images is None or images == "" or images == []:
             images = []
+        elif not isinstance(images, list):
+            images = [images]
+
+        image_paths = []
+        image_data = []
+        for img in images:
+            if isinstance(img, str):
+                image_paths.append(img)
+            else:
+                image_data.append(img)
+        
+        conversations = item.get("messages", item.get("conversations"))
         prompts = []
 
         if isinstance(conversations, list) and isinstance(conversations[0], list):
@@ -67,27 +59,52 @@ def __getitem__(self, idx):
                             "Pixtral batch processing is not supported yet. Set batch size to 1."
                         )
 
-                prompt = get_prompt(
-                    self.config["model_type"], self.processor, conversation
-                )
+                if "chat_template" in self.processor.__dict__:
+                    prompt = self.processor.apply_chat_template(
+                        conversation,
+                        tokenize=False,
+                        add_generation_prompt=False,
+                        num_images=len(images),
+                        num_audios=0,
+                    )
+                else:
+                    prompt = self.processor.tokenizer.apply_chat_template(
+                        conversation,
+                        tokenize=False,
+                        add_generation_prompt=False,
+                        num_images=len(images),
+                        num_audios=0,
+                    )
                 prompts.append(prompt)
 
         else:
             if self.config["model_type"] == "pixtral":
                 conversations = [json.loads(i) for i in conversations]
-            prompt = get_prompt(
-                self.config["model_type"], self.processor, conversations
-            )
+            if "chat_template" in self.processor.__dict__:
+                prompt = self.processor.apply_chat_template(
+                    conversations,
+                    tokenize=False,
+                    add_generation_prompt=False,
+                    num_images=len(images),
+                    num_audios=0,
+                )
+            else:
+                prompt = self.processor.tokenizer.apply_chat_template(
+                    conversations,
+                    tokenize=False,
+                    add_generation_prompt=False,
+                    num_images=len(images),
+                    num_audios=0,
+                )
             prompts.append(prompt)
 
-        image_token_index = getattr(self.config, "image_token_index", "image_token_id")
 
         inputs = prepare_inputs(
             processor=self.processor,
-            images=images,
+            images=image_data,
             audio=None,
             prompts=prompts,
-            image_token_index=image_token_index,
+            image_token_index=getattr(self.config, "image_token_index", "image_token_id"),
             resize_shape=self.image_resize_shape
         )
 

mlx_vlm/trainer/utils.py

@@ -20,7 +20,7 @@ class Colors:
     BOLD = '\033[1m'
     UNDERLINE = '\033[4m'
 
-supported_for_training = {"qwen2_vl", "qwen2_5_vl"}
+supported_for_training = {"qwen2_vl", "qwen2_5_vl", "gemma3"}
 
 def grad_checkpoint(layer):
     """