fix: make torchvision optional

Makefile

@@ -1,12 +1,12 @@
 install:
 	@echo "--- 🚀 Installing project dependencies ---"
-	pip install -e ".[dev]"
+	pip install -e ".[dev,image]"
 	pre-commit install
 
 install-for-tests:
 	@echo "--- 🚀 Installing project dependencies for test ---"
 	@echo "This ensures that the project is not installed in editable mode"
-	pip install ".[dev,speedtask]"
+	pip install ".[dev,speedtask,image]"
 
 lint:
 	@echo "--- 🧹 Running linters ---"

mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py

@@ -15,9 +15,9 @@
 from datasets import Dataset
 from PIL import Image
 from torch.utils.data import DataLoader
-from torchvision import transforms
 
 from mteb.encoder_interface import Encoder
+from mteb.requires_package import requires_image_dependencies
 
 from ..Evaluator import Evaluator
 from ..utils import (
@@ -36,7 +36,13 @@
 
 logger = logging.getLogger(__name__)
 
-transform = transforms.Compose([transforms.PILToTensor()])
+
+def get_default_transform():
+    requires_image_dependencies()
+    from torchvision import transforms
+
+    return transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
+
 
 
 class ImageDataset(torch.utils.data.Dataset):
@@ -121,6 +127,8 @@ def search(
 
         q_modality = queries[0]["modality"]
 
+        default_transform = get_default_transform()
+
         if q_modality == "text":
             query_texts = queries["text"]
             query_embeddings = self.model.get_text_embeddings(
@@ -130,7 +138,7 @@ def search(
             )
         else:
             queries_dataset = ImageDataset(
-                queries, image_column_name="image", transform=transform
+                queries, image_column_name="image", transform=default_transform
             )
             query_image_dataloader = DataLoader(
                 queries_dataset,
@@ -182,7 +190,7 @@ def search(
                 )
             else:
                 corpus_dataset = ImageDataset(
-                    chunk, image_column_name="image", transform=transform
+                    chunk, image_column_name="image", transform=default_transform
                 )
                 corpus_image_dataloader = DataLoader(
                     corpus_dataset,

mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py

@@ -15,9 +15,9 @@
 from datasets import Dataset
 from PIL import Image
 from torch.utils.data import DataLoader
-from torchvision import transforms
 
 from mteb.encoder_interface import Encoder, PromptType
+from mteb.requires_package import requires_image_dependencies
 
 from ..Evaluator import Evaluator
 from ..utils import (
@@ -36,7 +36,12 @@
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_TRANSFORM = transforms.Compose([transforms.PILToTensor()])
+
+def get_default_transform():
+    requires_image_dependencies()
+    from torchvision import transforms
+
+    return transforms.Compose([transforms.PILToTensor()])
 
 
 class ImageDataset(torch.utils.data.Dataset):
@@ -74,13 +79,14 @@ def __init__(
         encode_kwargs: dict[str, Any] = {},
         corpus_chunk_size: int = 20000,
         previous_results: str | None = None,
-        transform=DEFAULT_TRANSFORM,
+        transform=None,
         **kwargs: Any,
     ):
         # Model is class that provides get_text_embeddings() and get_image_embeddings()
         self.model = model
         self.encode_kwargs = encode_kwargs
-        self.transform = transform
+        if transform is None:
+            self.transform = get_default_transform()
 
         if "batch_size" not in encode_kwargs:
             encode_kwargs["batch_size"] = 128

mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py

@@ -7,15 +7,13 @@
 import torch
 from sklearn.metrics import accuracy_score
 from sklearn.metrics.pairwise import cosine_similarity
-from torchvision import transforms
 from tqdm import tqdm
 
 from mteb.encoder_interface import Encoder, EncoderWithSimilarity
 from mteb.evaluation.evaluators.Evaluator import Evaluator
 
 logger = logging.getLogger(__name__)
 
-transform = transforms.Compose([transforms.PILToTensor()])
 
 
 class Any2TextMultipleChoiceEvaluator(Evaluator):
@@ -51,7 +49,6 @@ def __init__(
         self.label_column_name = label_column_name
         self.choices_column_name = choices_column_name
         self.task_name = task_name
-        self.transform = transform
 
     def __call__(
         self,

mteb/evaluation/evaluators/Image/ClassificationEvaluator.py

@@ -16,9 +16,9 @@
 from sklearn.neighbors import KNeighborsClassifier
 from torch import Tensor
 from torch.utils.data import DataLoader
-from torchvision import transforms
 
 from mteb.encoder_interface import Encoder
+from mteb.requires_package import requires_image_dependencies
 
 from ..Evaluator import Evaluator
 
@@ -29,7 +29,11 @@ def dot_distance(a: np.ndarray, b: np.ndarray) -> float:
     return -np.dot(a, b)
 
 
-transform = transforms.Compose([transforms.PILToTensor()])
+def get_default_transform():
+    requires_image_dependencies()
+    from torchvision import transforms
+
+    return transforms.Compose([transforms.PILToTensor()])
 
 
 class ImageDataset(torch.utils.data.Dataset):
@@ -71,13 +75,18 @@ def __init__(
         if limit is not None:
             dataset_train = dataset_train.select(list(range(limit)))
 
+        default_transform = get_default_transform()
         self.dataset_train = ImageDataset(
-            dataset_train, image_column_name=image_column_name, transform=transform
+            dataset_train,
+            image_column_name=image_column_name,
+            transform=default_transform,
         )
         self.y_train = dataset_train[label_column_name]
 
         self.dataset_test = ImageDataset(
-            dataset_test, image_column_name=image_column_name, transform=transform
+            dataset_test,
+            image_column_name=image_column_name,
+            transform=default_transform,
         )
         self.y_test = dataset_test[label_column_name]
         self.task_name = task_name
@@ -155,13 +164,18 @@ def __init__(
         if limit is not None:
             dataset_train = dataset_train.select(list(range(limit)))
 
+        default_transform = get_default_transform()
         self.dataset_train = ImageDataset(
-            dataset_train, image_column_name=image_column_name, transform=transform
+            dataset_train,
+            image_column_name=image_column_name,
+            transform=default_transform,
         )
         self.y_train = dataset_train[label_column_name]
 
         self.dataset_test = ImageDataset(
-            dataset_test, image_column_name=image_column_name, transform=transform
+            dataset_test,
+            image_column_name=image_column_name,
+            transform=default_transform,
         )
         self.y_test = dataset_test[label_column_name]
         self.task_name = task_name
@@ -322,12 +336,17 @@ def __init__(
         if limit is not None:
             dataset_train = dataset_train.select(list(range(limit)))
 
+        default_transform = get_default_transform()
         self.dataset_train = ImageDataset(
-            dataset_train, image_column_name=image_column_name, transform=transform
+            dataset_train,
+            image_column_name=image_column_name,
+            transform=default_transform,
         )
         self.y_train = dataset_train[label_column_name]
         self.dataset_test = ImageDataset(
-            dataset_test, image_column_name=image_column_name, transform=transform
+            dataset_test,
+            image_column_name=image_column_name,
+            transform=default_transform,
         )
         self.y_test = dataset_test[label_column_name]
 

mteb/evaluation/evaluators/Image/ImageTextPairClassificationEvaluator.py

@@ -8,15 +8,12 @@
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
-from torchvision import transforms
 
 from mteb.encoder_interface import Encoder, EncoderWithSimilarity
 from mteb.evaluation.evaluators.Evaluator import Evaluator
 
 logger = logging.getLogger(__name__)
 
-transform = transforms.Compose([transforms.PILToTensor()])
-
 
 class ImageTextDataset(torch.utils.data.Dataset):
     def __init__(

mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py

@@ -14,14 +14,19 @@
     paired_manhattan_distances,
 )
 from torch.utils.data import DataLoader
-from torchvision import transforms
+
+from mteb.requires_package import requires_image_dependencies
 
 from ..Evaluator import Evaluator
 
 logger = logging.getLogger(__name__)
 
-transform = transforms.Compose([transforms.PILToTensor()])
 
+def get_default_transform():
+    requires_image_dependencies()
+    from torchvision import transforms
+
+    return transforms.Compose([transforms.PILToTensor()])
 
 class ImageDataset(torch.utils.data.Dataset):
     def __init__(self, hf_dataset, image_column_name: str = "image", transform=None):
@@ -54,11 +59,13 @@ def __init__(
         **kwargs,
     ):
         super().__init__(**kwargs)
+
+        default_transform = get_default_transform()
         self.sentence1_dataset = ImageDataset(
-            dataset, image_column_name=sentences_column_names[0], transform=transform
+            dataset, image_column_name=sentences_column_names[0], transform=default_transform
         )
         self.sentence2_dataset = ImageDataset(
-            dataset, image_column_name=sentences_column_names[1], transform=transform
+            dataset, image_column_name=sentences_column_names[1], transform=default_transform
         )
         self.gold_scores = gold_scores
         self.task_name = task_name

mteb/evaluation/evaluators/Image/ZeroShotClassificationEvaluator.py

@@ -8,15 +8,20 @@
 import torch
 from sklearn import metrics
 from torch.utils.data import DataLoader
-from torchvision import transforms
 
 from mteb.encoder_interface import Encoder
+from mteb.requires_package import requires_image_dependencies
 
 from ..Evaluator import Evaluator
 
 logger = logging.getLogger(__name__)
 
-transform = transforms.Compose([transforms.PILToTensor()])
+
+def get_default_transform():
+    requires_image_dependencies()
+    from torchvision import transforms
+
+    return transforms.Compose([transforms.PILToTensor()])
 
 
 class ImageDataset(torch.utils.data.Dataset):
@@ -52,7 +57,9 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.dataset = ImageDataset(
-            dataset, image_column_name=image_column_name, transform=transform
+            dataset,
+            image_column_name=image_column_name,
+            transform=get_default_transform(),
         )
         self.image_column_name = image_column_name
         self.labels = labels

mteb/model_meta.py

@@ -15,6 +15,7 @@
 
 from mteb.abstasks.AbsTask import AbsTask
 from mteb.encoder_interface import Encoder
+from mteb.requires_package import requires_image_dependencies
 
 from .custom_validators import LICENSES, MODALITIES, STR_DATE, STR_URL
 from .languages import ISO_LANGUAGE_SCRIPT

mteb/models/cohere_v.py

@@ -10,14 +10,11 @@
 import torch
 from PIL import Image
 from torch.utils.data import DataLoader
-from torchvision import transforms
 from tqdm import tqdm
 
 from mteb.encoder_interface import PromptType
 from mteb.model_meta import ModelMeta
-
-api_key = os.getenv("COHERE_API_KEY")
-tensor_to_image = transforms.Compose([transforms.ToPILImage()])
+from mteb.requires_package import requires_image_dependencies
 
 
 def cohere_v_loader(**kwargs):
@@ -32,15 +29,20 @@ def __init__(
             model_name: str,
             **kwargs: Any,
         ):
-            self.model_name = model_name
-            self.client = cohere.ClientV2(api_key)
-            self.image_format = "JPEG"
-            """ Wrapper for Cohere multimodal embedding model,
+            """Wrapper for Cohere multimodal embedding model,
 
             do `export COHERE_API_KEY=<Your_Cohere_API_KEY>` before running eval scripts.
             Cohere currently supports 40 images/min, thus time.sleep(1.5) is applied after each image.
             Remove or adjust this after Cohere API changes capacity.
             """
+            requires_image_dependencies()
+            from torchvision import transforms
+
+            self.model_name = model_name
+            api_key = os.getenv("COHERE_API_KEY")
+            self.client = cohere.ClientV2(api_key)
+            self.image_format = "JPEG"
+            self.transform = transforms.Compose([transforms.PILToTensor()])
 
         def get_text_embeddings(
             self,

mteb/models/evaclip_models.py

@@ -10,6 +10,7 @@
 
 from mteb.encoder_interface import PromptType
 from mteb.model_meta import ModelMeta
+from mteb.requires_package import requires_image_dependencies
 
 
 def evaclip_loader(**kwargs):
@@ -36,6 +37,8 @@ def __init__(
             device: str = "cuda" if torch.cuda.is_available() else "cpu",
             **kwargs: Any,
         ):
+            requires_image_dependencies()
+
             self.model_name = model_name
             self.device = device
             pretrained = "eva_clip"  # or "/path/to/EVA02_CLIP_B_psz16_s8B.pt"
@@ -86,10 +89,10 @@ def get_image_embeddings(
             batch_size: int = 32,
             **kwargs: Any,
         ):
+            import torchvision.transforms.functional as F
+
             all_image_embeddings = []
             if isinstance(images, DataLoader):
-                import torchvision.transforms.functional as F
-
                 with torch.no_grad(), torch.cuda.amp.autocast():
                     for batch in tqdm(images):
                         # import pdb; pdb.set_trace()