Update train.py

train.py

@@ -6,9 +6,7 @@ from typing import cast
 import os
 import shutil
 import multiprocessing as mp
-from torch.utils.data import DataLoader
 from PIL import Image
-from functools import partial
 
 
 def load_model(model_name, device_id=0):
@@ -33,17 +31,16 @@ def load_model(model_name, device_id=0):
     return processor, model
 
 
-def prepare_image(image):
-    if isinstance(image, Image.Image):
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-        return image
-    return image
+def caption_batch(batch, processor, model):
+    images = batch["image"]
 
+    pil_images = []
+    for image in images:
+        if isinstance(image, Image.Image):
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            pil_images.append(image)
 
-def collate_fn(batch, processor):
-    images = [prepare_image(item["image"]) for item in batch]
-    
     msg = [
         {
             "role": "user",
@@ -60,9 +57,36 @@ def collate_fn(batch, processor):
     text = processor.apply_chat_template(
         msg, add_generation_prompt=True, tokenize=False
     )
-    texts = [text] * len(images)
+    texts = [text] * len(pil_images)
+
+    inputs = processor(text=texts, images=pil_images, return_tensors="pt", padding=True)
+
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+    with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.bfloat16):
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            do_sample=False,
+        )
 
-    return processor(text=texts, images=images, return_tensors="pt", padding=True)
+    decoded = processor.batch_decode(generated, skip_special_tokens=False)
+
+    captions = []
+    special_tokens = set(processor.tokenizer.all_special_tokens)
+    for d in decoded:
+        if "<|im_start|>assistant" in d:
+            d = d.split("<|im_start|>assistant")[-1]
+
+        for token in special_tokens:
+            d = d.replace(token, "")
+
+        d = d.strip()
+        captions.append(d)
+
+    return {
+        "text": captions,
+    }
 
 
 def process_shard(gpu_id, start, end, model_name, batch_size, input_dataset, output_file):
@@ -80,51 +104,16 @@ def process_shard(gpu_id, start, end, model_name, batch_size, input_dataset, out
         else:
             shard = cast(Dataset, loaded)
 
-        shard.set_format(type="torch", columns=["image"])
-        
-        dataloader = DataLoader(
-            shard,
+        print(f"[GPU {gpu_id}] Processing {len(shard)} examples...", flush=True)
+        result = shard.map(
+            lambda batch: caption_batch(batch, processor, model),
+            batched=True,
             batch_size=batch_size,
-            num_workers=4,
-            pin_memory=True,
-            collate_fn=partial(collate_fn, processor=processor),
-            prefetch_factor=2,
+            remove_columns=[col for col in shard.column_names if col != "image"],
         )
 
-        all_captions = []
-        special_tokens = set(processor.tokenizer.all_special_tokens)
-
-        print(f"[GPU {gpu_id}] Processing {len(shard)} examples...", flush=True)
-        
-        for batch_idx, inputs in enumerate(dataloader):
-            inputs = {k: v.to(model.device, non_blocking=True) for k, v in inputs.items()}
-
-            with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.bfloat16):
-                generated = model.generate(
-                    **inputs,
-                    max_new_tokens=128,
-                    do_sample=False,
-                )
-
-            decoded = processor.batch_decode(generated, skip_special_tokens=False)
-
-            for d in decoded:
-                if "<|im_start|>assistant" in d:
-                    d = d.split("<|im_start|>assistant")[-1]
-
-                for token in special_tokens:
-                    d = d.replace(token, "")
-
-                d = d.strip()
-                all_captions.append(d)
-
-            if (batch_idx + 1) % 10 == 0:
-                print(f"[GPU {gpu_id}] Processed {(batch_idx + 1) * batch_size}/{len(shard)} examples", flush=True)
-
-        result_ds = Dataset.from_dict({"text": all_captions})
-
         print(f"[GPU {gpu_id}] Saving results to {output_file}...", flush=True)
-        result_ds.save_to_disk(output_file)
+        result.save_to_disk(output_file)
 
         print(f"[GPU {gpu_id}] Done!", flush=True)
         return output_file