Add the cleaned dataset

data/train_dataset.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b6f8f9ecea3f4c6f8196f194510159fccde43ee7f2192b259a11d6bc9ad684cb
-size 13426560
+oid sha256:898792e7bb3f1b4b390d35a6e2bad326a1fd5db44e169dd33b02bce5f1d6a4dc
+size 14451628

data/valid_dataset.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7dbb940f0dee7cb4a85959dc6018aafc824a988b46e3ae8ca2fea6500251ee0a
-size 4132661
+oid sha256:5de7c2611565819c960f93f007a0015355686fc0c76ab496729474981c441a3d
+size 1834366

run_medclip.sh

@@ -1,15 +1,15 @@
 python src/hybrid_clip/run_hybrid_clip.py \
-    --output_dir ./snapshots/final \
+    --output_dir ./snapshots/vision_augmented \
     --text_model_name_or_path="roberta-base" \
     --vision_model_name_or_path="openai/clip-vit-base-patch32" \
     --tokenizer_name="roberta-base" \
     --train_file="data/train_dataset.json" \
     --validation_file="data/valid_dataset.json" \
     --do_train --do_eval \
-    --num_train_epochs="40" --max_seq_length 96 \
+    --num_train_epochs="40" --max_seq_length 128 \
     --per_device_train_batch_size="64" \
     --per_device_eval_batch_size="64" \
-    --learning_rate="5e-4" --warmup_steps="0" --weight_decay 0.1 \
+    --learning_rate="5e-5" --warmup_steps="0" --weight_decay 0.1 \
     --overwrite_output_dir \
     --preprocessing_num_workers 32 \
 #    --push_to_hub

src/hybrid_clip/run_hybrid_clip.py

@@ -37,6 +37,7 @@ from torchvision.datasets import VisionDataset
 from torchvision.io import ImageReadMode, read_image
 from torchvision.transforms import CenterCrop, ConvertImageDtype, Normalize, Resize
 from torchvision.transforms.functional import InterpolationMode
+from torchvision.transforms.transforms import GaussianBlur, RandomAutocontrast, RandomHorizontalFlip
 from tqdm import tqdm
 
 import jax
@@ -178,6 +179,9 @@ class Transform(torch.nn.Module):
         self.transforms = torch.nn.Sequential(
             Resize([image_size], interpolation=InterpolationMode.BICUBIC),
             CenterCrop(image_size),
+            GaussianBlur(3, sigma=(0.05, 0.2)),
+            RandomAutocontrast(p=0.5),
+            RandomHorizontalFlip(p=0.5),
             ConvertImageDtype(torch.float),
             Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
         )
@@ -224,10 +228,10 @@ class ImageTextDataset(VisionDataset):
         self.image_paths = []
 
         for example in examples:
-            self.captions.extend([example["captions"]])
-            self.image_paths.extend([example["image_path"]])
-            # self.captions.extend(example["captions"][:captions_per_image])
-            # self.image_paths.extend([example["image_path"]] * captions_per_image)
+            # self.captions.extend(example["captions"])
+            # self.image_paths.append(example["image_path"])
+            self.captions.extend(example["captions"][:captions_per_image])
+            self.image_paths.extend([example["image_path"]] * captions_per_image)
 
     def _load_image(self, idx: int):
         path = self.image_paths[idx]
@@ -374,8 +378,10 @@ def main():
     # Use collate function to tokenizer the text and convert the processed images to numpy
     def collate_fn(examples):
         pixel_values = torch.stack([example[0] for example in examples]).permute(0, 2, 3, 1).numpy()
+        # pixel_values = torch.stack([example[0] for example in examples]).numpy()
         captions = [example[1] for example in examples]
-        inputs = tokenizer(captions, max_length=data_args.max_seq_length, padding="max_length", return_tensors="np")
+        inputs = tokenizer(captions, max_length=data_args.max_seq_length, padding="max_length", return_tensors="np",
+                           truncation=True)
 
         batch = {
             "pixel_values": pixel_values,