[Add] transformers text and image models

clip_benchmark/models/__init__.py

@@ -2,11 +2,13 @@
 import torch
 from .open_clip import load_open_clip
 from .japanese_clip import load_japanese_clip
+from .transformers_clip import load_transformers_clip
 
 # loading function must return (model, transform, tokenizer)
 TYPE2FUNC = {
     "open_clip": load_open_clip,
-    "ja_clip": load_japanese_clip
+    "ja_clip": load_japanese_clip,
+    "transformers": load_transformers_clip,
 }
 MODEL_TYPES = list(TYPE2FUNC.keys())
 

clip_benchmark/models/transformers_clip.py

@@ -0,0 +1,29 @@
+import torch
+from torch import nn
+from transformers import AutoModel, AutoProcessor
+from functools import partial
+
+class TransformerWrapper(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def encode_text(self, text):
+        return self.model.get_text_features(**text)
+
+    def encode_image(self, image):
+        # 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.preprocess, return_tensors="pt")
+    tokenizer = partial(
+        processor.tokenizer, return_tensors="pt", padding="max_length", max_length=64
+    )
+    return model, transforms, tokenizer