adding comments

Author: ariG23498

clip_benchmark/models/transformers_clip.py

@@ -1,3 +1,4 @@
+import torch
 from torch import nn
 from transformers import AutoModel, AutoProcessor
 from functools import partial
@@ -11,14 +12,19 @@ def encode_text(self, text):
         return self.model.get_text_features(**text)
 
     def encode_image(self, image):
-        return self.model.get_image_features(image["pixel_values"].squeeze(1))
+        # we get an extended dimension possibly due to the collation in dataloader
+        image = {key: value.squeeze(1) for key, value in image.items()}
+        return self.model.get_image_features(**image)
 
 def load_transformers_clip(model_name, pretrained, cache_dir, device):
     ckpt = f"{model_name}/{pretrained}"
     model = AutoModel.from_pretrained(ckpt, cache_dir=cache_dir, device_map=device)
     model = TransformerWrapper(model)
+    
     processor = AutoProcessor.from_pretrained(ckpt)
-
-    transforms = partial(processor.image_processor, return_tensors="pt") 
-    tokenizer = partial(processor.tokenizer, return_tensors="pt", padding="max_length")
+    transforms = partial(processor.image_processor.preprocess, return_tensors="pt")
+    tokenizer = partial(
+        processor.tokenizer, return_tensors="pt", padding="max_length",
+        max_length=64 # very specific to SG2
+        )
     return model, transforms, tokenizer