[mieb] Adding MOCOv3 models

mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py

@@ -62,7 +62,7 @@ def __call__(
             encode_kwargs["batch_size"] = 64
 
         label_list = list(
-            set([x for n in self.dataset[self.choices_column_name] for x in n])
+            {x for n in self.dataset[self.choices_column_name] for x in n}
         )
         label_embeddings = model.get_text_embeddings(label_list)
         label_embedding_dict = {}

mteb/models/__init__.py

@@ -25,6 +25,7 @@
     gte_models,
     jina_clip,
     llm2vec_models,
+    moco_models,
     mxbai_models,
     nomic_models,
     nomic_models_vision,
@@ -147,6 +148,7 @@ def model_meta_from_sentence_transformers(model: SentenceTransformer) -> ModelMe
     gte_models,
     jina_clip,
     llm2vec_models,
+    moco_models,
     mxbai_models,
     nomic_models,
     nomic_models_vision,

mteb/models/moco_models.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Any
+
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from mteb.model_meta import ModelMeta
+
+
+def mocov3_loader(**kwargs):
+    try:
+        import timm
+    except ImportError:
+        raise ImportError("Please install `pip install timm` to use MOCOv3 models.")
+
+    class MOCOv3Wrapper:
+        """A wrapper class for MOCOv3 models that supports image encoding.
+        Text encoding and text-image fusion are not supported.
+        """
+
+        def __init__(
+            self,
+            model_name: str = "nyu-visionx/moco-v3-vit-b",
+            device: str = "cuda" if torch.cuda.is_available() else "cpu",
+            **kwargs: Any,
+        ):
+            self.model_name = model_name
+            self.device = device
+            name = "vit_base_patch16_224"
+            if "vit-l" in model_name:
+                name = "vit_large_patch16_224"
+            model = timm.create_model(
+                name,
+                pretrained=True,
+                num_classes=0,
+                pretrained_cfg_overlay={"hf_hub_id": model_name},
+            )
+
+            self.model = model.eval()
+
+            # get model specific transforms (normalization, resize)
+            data_config = timm.data.resolve_model_data_config(self.model)
+            self.processor = timm.data.create_transform(
+                **data_config, is_training=False
+            )
+
+        @staticmethod
+        def get_text_embeddings(texts: list[str], batch_size: int = 32):
+            raise ValueError("MOCO models only support image encoding.")
+
+        def get_image_embeddings(
+            self,
+            images: list[Image.Image] | DataLoader,
+            batch_size: int = 32,
+        ):
+            all_image_embeddings = []
+
+            if isinstance(images, DataLoader):
+                import torchvision.transforms.functional as F
+
+                with torch.no_grad():
+                    for batch in tqdm(images):
+                        inputs = torch.vstack(
+                            [
+                                self.processor(F.to_pil_image(b.to("cpu"))).unsqueeze(0)
+                                for b in batch
+                            ]
+                        )
+                        output = self.model(
+                            inputs
+                        )  # output is (batch_size, num_features) shaped tensor
+                        all_image_embeddings.append(output)
+            else:
+                with torch.no_grad():
+                    for i in tqdm(range(0, len(images), batch_size)):
+                        batch_images = images[i : i + batch_size]
+                        output = self.model(
+                            self.processor(batch_images)
+                        )  # output is (batch_size, num_features) shaped tensor
+                        all_image_embeddings.append(output)
+
+            all_image_embeddings = torch.cat(all_image_embeddings, dim=0)
+            return all_image_embeddings
+
+        @staticmethod
+        def calculate_probs(text_embeddings, image_embeddings):
+            raise ValueError("MOCO models only support image encoding.")
+
+        def get_fused_embeddings(
+            self,
+            texts: list[str] = None,
+            images: list[Image.Image] | DataLoader = None,
+            fusion_mode="sum",
+            batch_size: int = 32,
+        ):
+            if texts is None and images is None:
+                raise ValueError("images must be provided for MOCO models")
+
+            text_embeddings = None
+            image_embeddings = None
+
+            if texts is not None:
+                text_embeddings = self.get_text_embeddings(texts, batch_size)
+
+            if images is not None:
+                image_embeddings = self.get_image_embeddings(images, batch_size)
+
+            if text_embeddings is not None and image_embeddings is not None:
+                raise ValueError("MOCO models only support image encoding.")
+            elif text_embeddings is not None:
+                return text_embeddings
+            elif image_embeddings is not None:
+                return image_embeddings
+
+    return MOCOv3Wrapper(**kwargs)
+
+
+mocov3_vit_base = ModelMeta(
+    loader=partial(
+        mocov3_loader,
+        model_name="nyu-visionx/moco-v3-vit-b",
+    ),
+    name="nyu-visionx/moco-v3-vit-b",
+    languages=["eng_Latn"],
+    open_source=True,
+    revision="7d091cd70772c5c0ecf7f00b5f12ca609a99d69d",
+    release_date="2024-06-03",
+)
+
+mocov3_vit_large = ModelMeta(
+    loader=partial(
+        mocov3_loader,
+        model_name="nyu-visionx/moco-v3-vit-l",
+    ),
+    name="nyu-visionx/moco-v3-vit-l",
+    languages=["eng_Latn"],
+    open_source=True,
+    revision="7bf75358d616f39b9716148bf4e3425f3bd35b47",
+    release_date="2024-06-03",
+)

results-mieb/nyu-visionx__moco-v3-vit-b/7d091cd70772c5c0ecf7f00b5f12ca609a99d69d/OxfordFlowersClassification.json

@@ -0,0 +1,48 @@
+{
+  "dataset_revision": "a37b1891609c0376fa81eced756e7863e1bd873b",
+  "evaluation_time": 347.9395191669464,
+  "kg_co2_emissions": null,
+  "mteb_version": "1.14.21",
+  "scores": {
+    "test": [
+      {
+        "accuracy": 0.8880392156862745,
+        "f1": 0.8858313070145011,
+        "f1_weighted": 0.8855993382917493,
+        "hf_subset": "default",
+        "languages": [
+          "eng-Latn"
+        ],
+        "main_score": 0.8880392156862745,
+        "scores_per_experiment": [
+          {
+            "accuracy": 0.8921568627450981,
+            "f1": 0.8908725994288734,
+            "f1_weighted": 0.890707976140322
+          },
+          {
+            "accuracy": 0.888235294117647,
+            "f1": 0.8862588054389066,
+            "f1_weighted": 0.8857789823937094
+          },
+          {
+            "accuracy": 0.8833333333333333,
+            "f1": 0.8807271651302992,
+            "f1_weighted": 0.8807519823317017
+          },
+          {
+            "accuracy": 0.884313725490196,
+            "f1": 0.8818842242107174,
+            "f1_weighted": 0.8815393805409137
+          },
+          {
+            "accuracy": 0.8921568627450981,
+            "f1": 0.8894137408637092,
+            "f1_weighted": 0.8892183700520992
+          }
+        ]
+      }
+    ]
+  },
+  "task_name": "OxfordFlowersClassification"
+}

results-mieb/nyu-visionx__moco-v3-vit-b/7d091cd70772c5c0ecf7f00b5f12ca609a99d69d/model_meta.json

@@ -0,0 +1 @@
+{"name": "nyu-visionx/moco-v3-vit-b", "revision": "7d091cd70772c5c0ecf7f00b5f12ca609a99d69d", "release_date": "2024-06-03", "languages": ["eng_Latn"], "n_parameters": null, "memory_usage": null, "max_tokens": null, "embed_dim": null, "license": null, "open_source": true, "similarity_fn_name": null, "framework": [], "loader": "mocov3_loader"}

results-mieb/nyu-visionx__moco-v3-vit-l/7bf75358d616f39b9716148bf4e3425f3bd35b47/OxfordFlowersClassification.json

@@ -0,0 +1,48 @@
+{
+  "dataset_revision": "a37b1891609c0376fa81eced756e7863e1bd873b",
+  "evaluation_time": 780.7572722434998,
+  "kg_co2_emissions": null,
+  "mteb_version": "1.14.21",
+  "scores": {
+    "test": [
+      {
+        "accuracy": 0.8990196078431373,
+        "f1": 0.8976575858316652,
+        "f1_weighted": 0.8973426503552833,
+        "hf_subset": "default",
+        "languages": [
+          "eng-Latn"
+        ],
+        "main_score": 0.8990196078431373,
+        "scores_per_experiment": [
+          {
+            "accuracy": 0.8960784313725491,
+            "f1": 0.8954744443526484,
+            "f1_weighted": 0.8949013011291541
+          },
+          {
+            "accuracy": 0.9009803921568628,
+            "f1": 0.89924527322374,
+            "f1_weighted": 0.89919388345273
+          },
+          {
+            "accuracy": 0.8960784313725491,
+            "f1": 0.8937528098526334,
+            "f1_weighted": 0.8936883759623736
+          },
+          {
+            "accuracy": 0.8990196078431373,
+            "f1": 0.8985132758082764,
+            "f1_weighted": 0.8977617544247366
+          },
+          {
+            "accuracy": 0.9029411764705882,
+            "f1": 0.901302125921028,
+            "f1_weighted": 0.901167936807422
+          }
+        ]
+      }
+    ]
+  },
+  "task_name": "OxfordFlowersClassification"
+}

results-mieb/nyu-visionx__moco-v3-vit-l/7bf75358d616f39b9716148bf4e3425f3bd35b47/model_meta.json

@@ -0,0 +1 @@
+{"name": "nyu-visionx/moco-v3-vit-l", "revision": "7bf75358d616f39b9716148bf4e3425f3bd35b47", "release_date": "2024-06-03", "languages": ["eng_Latn"], "n_parameters": null, "memory_usage": null, "max_tokens": null, "embed_dim": null, "license": null, "open_source": true, "similarity_fn_name": null, "framework": [], "loader": "mocov3_loader"}