Update train.py

train.py

@@ -4,6 +4,9 @@ import datasets
 from datasets import Dataset
 from typing import cast
 import os
+import shutil
+import multiprocessing as mp
+from PIL import Image
 
 
 def load_model(model_name, device_id=0):
@@ -29,8 +32,16 @@ def load_model(model_name, device_id=0):
 def caption_batch(batch, processor, model):
     images = batch["image"]
 
-    encoded_list = []
+    processed_images = []
     for image in images:
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        processed_images.append(image)
+
+    encoded_list = []
+    for image in processed_images:
         msg = [
             {
                 "role": "user",
@@ -79,11 +90,16 @@ def caption_batch(batch, processor, model):
     for d in decoded:
         if "<|im_start|>assistant" in d:
             d = d.split("<|im_start|>assistant")[-1].strip()
-        d = d.replace("<|im_end|>", "").strip()
+
+        special_tokens = set(processor.tokenizer.all_special_tokens)
+        for token in special_tokens:
+            d = d.replace(token, "")
+
+        d = d.strip()
         captions.append(d)
 
     return {
-        "image": images,
+        "image": processed_images,
         "text": captions,
     }
 
@@ -120,12 +136,10 @@ def process_shard_worker(
 
 
 def main():
-    import multiprocessing as mp
-
-    input_dataset = "none-yet/anime-captions"
-    output_dataset = "none-yet/anime-captions"
+    input_dataset = "nroggendorff/fries"
+    output_dataset = "nroggendorff/fries"
     model_name = "datalab-to/chandra"
-    batch_size = 32
+    batch_size = 2
 
     print("Loading dataset info...")
     loaded = datasets.load_dataset(input_dataset, split="train")
@@ -174,8 +188,6 @@ def main():
     final_ds.push_to_hub(output_dataset, create_pr=True)
 
     print("Cleaning up temporary files...")
-    import shutil
-
     for f in temp_files:
         if os.path.exists(f):
             shutil.rmtree(f)