From f00f31eee76ecb7bd93aa129ea86ca9e96e51036 Mon Sep 17 00:00:00 2001 From: Kenneth Enevoldsen Date: Wed, 19 Mar 2025 21:30:54 +0100 Subject: [PATCH 1/5] fix: make torchvision optional --- Makefile | 4 +-- .../Image/Any2AnyMultiChoiceEvaluator.py | 16 ++++++--- .../Image/Any2AnyRetrievalEvaluator.py | 14 +++++--- .../Image/Any2TextMultipleChoiceEvaluator.py | 3 -- .../Image/ClassificationEvaluator.py | 35 ++++++++++++++----- .../ImageTextPairClassificationEvaluator.py | 3 -- .../evaluators/Image/VisualSTSEvaluator.py | 15 +++++--- .../Image/ZeroShotClassificationEvaluator.py | 13 +++++-- mteb/model_meta.py | 1 + mteb/models/cohere_v.py | 18 +++++----- mteb/models/evaclip_models.py | 7 ++-- mteb/models/jina_clip.py | 6 +++- mteb/models/llm2clip_models.py | 7 ++-- mteb/models/moco_models.py | 11 +++--- mteb/models/openclip_models.py | 11 +++--- mteb/models/vista_models.py | 14 ++++---- mteb/models/vlm2vec_models.py | 14 ++++---- mteb/models/voyage_v.py | 22 +++++++----- mteb/requires_package.py | 8 +++++ pyproject.toml | 2 +- 20 files changed, 150 insertions(+), 74 deletions(-) diff --git a/Makefile b/Makefile index 5e6d700407..0f0937f08d 100644 --- a/Makefile +++ b/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 ---" diff --git a/mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py b/mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py index c69f0153a2..c8812cb04b 100644 --- a/mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py +++ b/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, diff --git a/mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py b/mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py index 777e3b545f..74a41fb1a3 100644 --- a/mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py +++ b/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 diff --git a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py index a93714e770..df35349a1f 100644 --- a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py +++ b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py @@ -7,7 +7,6 @@ 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 @@ -15,7 +14,6 @@ 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, diff --git a/mteb/evaluation/evaluators/Image/ClassificationEvaluator.py b/mteb/evaluation/evaluators/Image/ClassificationEvaluator.py index a0d84d5714..a3f7c022fa 100644 --- a/mteb/evaluation/evaluators/Image/ClassificationEvaluator.py +++ b/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] diff --git a/mteb/evaluation/evaluators/Image/ImageTextPairClassificationEvaluator.py b/mteb/evaluation/evaluators/Image/ImageTextPairClassificationEvaluator.py index f3188f7753..72cefbcfde 100644 --- a/mteb/evaluation/evaluators/Image/ImageTextPairClassificationEvaluator.py +++ b/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__( diff --git a/mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py b/mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py index a042d22f5a..170c498abb 100644 --- a/mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py +++ b/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 diff --git a/mteb/evaluation/evaluators/Image/ZeroShotClassificationEvaluator.py b/mteb/evaluation/evaluators/Image/ZeroShotClassificationEvaluator.py index da3e8f5f97..b25232d28e 100644 --- a/mteb/evaluation/evaluators/Image/ZeroShotClassificationEvaluator.py +++ b/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 diff --git a/mteb/model_meta.py b/mteb/model_meta.py index 96ef821021..74efd5a495 100644 --- a/mteb/model_meta.py +++ b/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 diff --git a/mteb/models/cohere_v.py b/mteb/models/cohere_v.py index a3856627b1..fd6406dc1c 100644 --- a/mteb/models/cohere_v.py +++ b/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=` 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, diff --git a/mteb/models/evaclip_models.py b/mteb/models/evaclip_models.py index 0b9e0e19bc..8cbd184447 100644 --- a/mteb/models/evaclip_models.py +++ b/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() diff --git a/mteb/models/jina_clip.py b/mteb/models/jina_clip.py index 94d498802f..7e11f4496b 100644 --- a/mteb/models/jina_clip.py +++ b/mteb/models/jina_clip.py @@ -11,6 +11,7 @@ from mteb.encoder_interface import PromptType from mteb.model_meta import ModelMeta +from mteb.requires_package import requires_image_dependencies class JinaCLIPModelWrapper: @@ -20,6 +21,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 self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to( @@ -63,11 +66,12 @@ def get_image_embeddings( convert_to_tensor=True, **kwargs: Any, ): + import torchvision.transforms.functional as F + all_image_embeddings = [] if isinstance(images, DataLoader): with torch.no_grad(): - import torchvision.transforms.functional as F for batch in tqdm(images): image_outputs = self.model.encode_image( diff --git a/mteb/models/llm2clip_models.py b/mteb/models/llm2clip_models.py index 25ed3c6808..5c2a17cfe8 100644 --- a/mteb/models/llm2clip_models.py +++ b/mteb/models/llm2clip_models.py @@ -12,6 +12,7 @@ from mteb.encoder_interface import PromptType from mteb.model_meta import ModelMeta +from mteb.requires_package import requires_image_dependencies MODEL2PROCESSOR = { "microsoft/LLM2CLIP-Openai-L-14-336": "openai/clip-vit-large-patch14-336", @@ -36,6 +37,8 @@ def __init__( device: str = "cuda" if torch.cuda.is_available() else "cpu", **kwargs: Any, ): + requires_image_dependencies() + if model_name not in MODEL2PROCESSOR: raise Exception( f"This model {model_name} is not in the supported mode list: {list(MODEL2PROCESSOR.keys())}." @@ -119,10 +122,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.amp.autocast("cuda"): for batch in tqdm(images): input_pixels = self.processor( diff --git a/mteb/models/moco_models.py b/mteb/models/moco_models.py index cb2ee875da..b88e9805c7 100644 --- a/mteb/models/moco_models.py +++ b/mteb/models/moco_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 mocov3_loader(**kwargs): @@ -29,6 +30,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 name = "vit_base_patch16_224" @@ -69,11 +72,11 @@ 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(): for batch in tqdm(images): inputs = torch.vstack( @@ -107,8 +110,8 @@ def calculate_probs(text_embeddings, image_embeddings): def get_fused_embeddings( self, - texts: list[str] = None, - images: list[Image.Image] | DataLoader = None, + texts: list[str] | None = None, + images: list[Image.Image] | DataLoader | None = None, *, task_name: str | None = None, prompt_type: PromptType | None = None, diff --git a/mteb/models/openclip_models.py b/mteb/models/openclip_models.py index 3079ff6933..8399fd4f64 100644 --- a/mteb/models/openclip_models.py +++ b/mteb/models/openclip_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 openclip_loader(**kwargs): @@ -25,6 +26,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 self.model, _, self.img_preprocess = open_clip.create_model_and_transforms( @@ -71,10 +74,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() @@ -112,8 +115,8 @@ def calculate_probs(self, text_embeddings, image_embeddings): def get_fused_embeddings( self, - texts: list[str] = None, - images: list[Image.Image] | DataLoader = None, + texts: list[str] | None = None, + images: list[Image.Image] | DataLoader | None = None, fusion_mode="sum", **kwargs: Any, ): diff --git a/mteb/models/vista_models.py b/mteb/models/vista_models.py index d2ae2d05c3..0905e649ab 100644 --- a/mteb/models/vista_models.py +++ b/mteb/models/vista_models.py @@ -6,13 +6,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 - -tensor_to_image = transforms.Compose([transforms.ToPILImage()]) +from mteb.requires_package import requires_image_dependencies def vista_loader(**kwargs): @@ -53,16 +51,19 @@ class VisualizedBGEWrapper(Visualized_BGE): def __init__( self, - model_name_bge: str = None, + model_name_bge: str | None = None, model_weight=None, normlized: bool = True, sentence_pooling_method: str = "cls", negatives_cross_device: bool = False, temperature: float = 0.02, from_pretrained=None, - image_tokens_num: int = None, + image_tokens_num: int | None = None, **kwargs: Any, ): + requires_image_dependencies() + from torchvision import transforms + super().__init__( model_name_bge=model_name_bge, model_weight=model_weight, @@ -76,6 +77,7 @@ def __init__( self.max_text_len_with_image = ( self.tokenizer.model_max_length - image_tokens_num ) + self.transform = transforms.Compose([transforms.ToPILImage()]) self.eval() def encode_text(self, texts): @@ -147,7 +149,7 @@ def encode( ] else: images = [ - self.preprocess_val(tensor_to_image(image)) + self.preprocess_val(self.tensor_to_image(image)) for image in images ] images = torch.stack(images) diff --git a/mteb/models/vlm2vec_models.py b/mteb/models/vlm2vec_models.py index fbf7bf9f0a..1d629b86c9 100644 --- a/mteb/models/vlm2vec_models.py +++ b/mteb/models/vlm2vec_models.py @@ -12,6 +12,7 @@ from mteb.encoder_interface import PromptType from mteb.model_meta import ModelMeta +from mteb.requires_package import requires_image_dependencies logging.basicConfig(level=logging.WARNING) logger = logging.getLogger(__name__) @@ -28,6 +29,7 @@ def __init__( device: str = "cuda" if torch.cuda.is_available() else "cpu", **kwargs, ): + requires_image_dependencies() try: import flash_attn # noqa from peft import LoraConfig, PeftModel # noqa @@ -119,11 +121,11 @@ def get_image_embeddings( batch_size: int = 32, **kwargs: Any, ): + import torchvision.transforms.functional as F + text = "<|image_1|> Represent the given image." all_image_embeddings = [] if isinstance(images, DataLoader): - import torchvision.transforms.functional as F - with torch.no_grad(): for batch in tqdm(images): input_ids, pixel_values, image_sizes = [], [], [] @@ -253,8 +255,8 @@ def calculate_probs(self, text_embeddings, image_embeddings): def get_fused_embeddings( self, - texts: list[str] = None, - images: list[Image.Image] | DataLoader = None, + texts: list[str] | None = None, + images: list[Image.Image] | DataLoader | None = None, *, task_name: str | None = None, prompt_type: PromptType | None = None, @@ -262,6 +264,8 @@ def get_fused_embeddings( fusion_mode="sum", **kwargs: Any, ): + import torchvision.transforms.functional as F + if texts is None and images is None: raise ValueError("Either texts or images must be provided") @@ -283,8 +287,6 @@ def get_fused_embeddings( texts = iter(texts) all_fused_embeddings = [] if isinstance(images, DataLoader): - import torchvision.transforms.functional as F - with torch.no_grad(): for batch in images: input_ids, pixel_values, image_sizes = [], [], [] diff --git a/mteb/models/voyage_v.py b/mteb/models/voyage_v.py index fc880347c5..5f93fda5d6 100644 --- a/mteb/models/voyage_v.py +++ b/mteb/models/voyage_v.py @@ -1,21 +1,19 @@ from __future__ import annotations import logging -import os from functools import partial from typing import Any 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 +from mteb.requires_package import requires_image_dependencies + -api_key = os.getenv("VOYAGE_API_KEY") -tensor_to_image = transforms.Compose([transforms.ToPILImage()]) def downsample_image( @@ -37,15 +35,15 @@ def downsample_image( logging.info( f"Downsampling image from {width}x{height} to {new_width}x{new_height}" ) - return image.resize(new_size, Image.LANCZOS) + return image.resize(new_size, Image.LANCZOS) # type: ignore if width > height: if width > 10000: logging.error("Processing extremely wide images.") - return image.resize((10000, height), Image.LANCZOS) + return image.resize((10000, height), Image.LANCZOS) # type: ignore else: if height > 10000: logging.error("Processing extremely high images.") - return image.resize((width, 10000), Image.LANCZOS) + return image.resize((width, 10000), Image.LANCZOS) # type: ignore return image @@ -67,8 +65,12 @@ def __init__( model_name: str, **kwargs: Any, ): + requires_image_dependencies() + from torchvision import transforms + self.model_name = model_name self.vo = voyageai.Client() + self.tensor_to_image = transforms.Compose([transforms.PILToTensor()]) @retry( stop=stop_after_attempt(6), # Stop after 6 attempts @@ -132,7 +134,8 @@ def get_image_embeddings( if index == 0: assert len(batch) == batch_size batch_images = [ - [downsample_image(tensor_to_image(image))] for image in batch + [downsample_image(self.tensor_to_image(image))] + for image in batch ] embeddings = self._multimodal_embed( batch_images, model=self.model_name, input_type=input_type @@ -190,7 +193,8 @@ def get_fused_embeddings( if index == 0: assert len(batch) == batch_size batch_images = [ - downsample_image(tensor_to_image(image)) for image in batch + downsample_image(self.tensor_to_image(image)) + for image in batch ] batch_texts = texts[ index * batch_size : (index + 1) * batch_size diff --git a/mteb/requires_package.py b/mteb/requires_package.py index e73dcb5f20..d261acdffb 100644 --- a/mteb/requires_package.py +++ b/mteb/requires_package.py @@ -22,3 +22,11 @@ def requires_package( f"{name} requires the `{package_name}` library but it was not found in your environment. " + f"If you want to load {model_name} models, please `{install_instruction}` to install the package." ) + + +def requires_image_dependencies() -> None: + if not _is_package_available("torchvision"): + raise ImportError( + "You are trying to running the image subset of mteb without having installed the required dependencies (`torchvision`). " + + "You can install the required dependencies using `pip install 'mteb[image]'` to install the required dependencies." + ) diff --git a/pyproject.toml b/pyproject.toml index 810416ef72..66e4ab66df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,6 @@ dependencies = [ "typing_extensions>=0.0.0", "eval_type_backport>=0.0.0", "polars>=0.20.22", - "torchvision>0.0.0", ] @@ -53,6 +52,7 @@ homepage = "https://github.com/embeddings-benchmark/mteb" mteb = "mteb.cli:main" [project.optional-dependencies] +image = ["torchvision>0.0.0"] dev = [ "ruff==0.9.7", # locked so we don't get PRs which fail only due to a lint update "pytest>=8.3.4", From 7b832b41e83985577116ca788e52593b089620c7 Mon Sep 17 00:00:00 2001 From: Kenneth Enevoldsen Date: Thu, 20 Mar 2025 08:36:10 +0100 Subject: [PATCH 2/5] format --- .../evaluators/Image/Any2AnyMultiChoiceEvaluator.py | 1 - .../evaluators/Image/Any2TextMultipleChoiceEvaluator.py | 1 - mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py | 9 +++++++-- mteb/model_meta.py | 1 - mteb/models/cohere_v.py | 6 +++--- mteb/models/jina_clip.py | 1 - mteb/models/voyage_v.py | 2 -- 7 files changed, 10 insertions(+), 11 deletions(-) diff --git a/mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py b/mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py index c8812cb04b..dd41e5f23c 100644 --- a/mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py +++ b/mteb/evaluation/evaluators/Image/Any2AnyMultiChoiceEvaluator.py @@ -44,7 +44,6 @@ def get_default_transform(): return transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()]) - class ImageDataset(torch.utils.data.Dataset): def __init__(self, hf_dataset, image_column_name: str = "image", transform=None): self.dataset = hf_dataset diff --git a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py index df35349a1f..89b1110261 100644 --- a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py +++ b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py @@ -15,7 +15,6 @@ logger = logging.getLogger(__name__) - class Any2TextMultipleChoiceEvaluator(Evaluator): """Evaluate a model based on the similarity of queries (can be interleaved) and candidate answers. The goal is to find the correct text in multiple candidates that diff --git a/mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py b/mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py index 170c498abb..69d203711f 100644 --- a/mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py +++ b/mteb/evaluation/evaluators/Image/VisualSTSEvaluator.py @@ -28,6 +28,7 @@ def get_default_transform(): return transforms.Compose([transforms.PILToTensor()]) + class ImageDataset(torch.utils.data.Dataset): def __init__(self, hf_dataset, image_column_name: str = "image", transform=None): self.dataset = hf_dataset @@ -62,10 +63,14 @@ def __init__( default_transform = get_default_transform() self.sentence1_dataset = ImageDataset( - dataset, image_column_name=sentences_column_names[0], transform=default_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=default_transform + dataset, + image_column_name=sentences_column_names[1], + transform=default_transform, ) self.gold_scores = gold_scores self.task_name = task_name diff --git a/mteb/model_meta.py b/mteb/model_meta.py index 74efd5a495..96ef821021 100644 --- a/mteb/model_meta.py +++ b/mteb/model_meta.py @@ -15,7 +15,6 @@ 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 diff --git a/mteb/models/cohere_v.py b/mteb/models/cohere_v.py index fd6406dc1c..a5d88ae9aa 100644 --- a/mteb/models/cohere_v.py +++ b/mteb/models/cohere_v.py @@ -83,7 +83,7 @@ def get_image_embeddings( for image in batch: # cohere only supports 1 image per call buffered = io.BytesIO() - image = tensor_to_image(image) + image = self.transform(image) image.save(buffered, format=self.image_format) image_bytes = buffered.getvalue() stringified_buffer = base64.b64encode(image_bytes).decode( @@ -144,8 +144,8 @@ def calculate_probs(self, text_embeddings, image_embeddings): def get_fused_embeddings( self, - texts: list[str] = None, - images: list[Image.Image] | DataLoader = None, + texts: list[str] | None = None, + images: list[Image.Image] | DataLoader | None = None, fusion_mode="sum", **kwargs: Any, ): diff --git a/mteb/models/jina_clip.py b/mteb/models/jina_clip.py index 7e11f4496b..208b77e44a 100644 --- a/mteb/models/jina_clip.py +++ b/mteb/models/jina_clip.py @@ -72,7 +72,6 @@ def get_image_embeddings( if isinstance(images, DataLoader): with torch.no_grad(): - for batch in tqdm(images): image_outputs = self.model.encode_image( [F.to_pil_image(b.to("cpu")) for b in batch], diff --git a/mteb/models/voyage_v.py b/mteb/models/voyage_v.py index 5f93fda5d6..96e7ff9997 100644 --- a/mteb/models/voyage_v.py +++ b/mteb/models/voyage_v.py @@ -14,8 +14,6 @@ from mteb.requires_package import requires_image_dependencies - - def downsample_image( image: Image.Image, max_pixels: int = 16000000, target_longest_side: int = 4000 ) -> Image.Image: From 0b40251fbdee5e8a6c927675b4ed163fdd93b869 Mon Sep 17 00:00:00 2001 From: Kenneth Enevoldsen Date: Thu, 20 Mar 2025 08:37:26 +0100 Subject: [PATCH 3/5] add docs --- docs/usage/usage.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/usage/usage.md b/docs/usage/usage.md index dafa72b50d..a315277952 100644 --- a/docs/usage/usage.md +++ b/docs/usage/usage.md @@ -44,10 +44,8 @@ results = evaluation.run(model) ### Evaluating on Different Modalities MTEB is not only text evaluating, but also allow you to evaluate image and image-text embeddings. - To evaluate image embeddings you can follows the same approach for any other task in `mteb`. Simply ensuring that the task contains the modality "image": From 075926b0a3e9b76ad5ecb68a01f7dac7511b4a49 Mon Sep 17 00:00:00 2001 From: Kenneth Enevoldsen Date: Thu, 20 Mar 2025 08:39:39 +0100 Subject: [PATCH 4/5] minor fix --- .../evaluators/Image/Any2TextMultipleChoiceEvaluator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py index 89b1110261..3cc9d2c058 100644 --- a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py +++ b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py @@ -48,6 +48,7 @@ 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, From a1aac3cbd4199928e8ed4185708c398a6923d3f8 Mon Sep 17 00:00:00 2001 From: Isaac Chung Date: Sat, 22 Mar 2025 11:56:18 +0200 Subject: [PATCH 5/5] remove transform from Any2TextMultipleChoiceEvaluator --- .../evaluators/Image/Any2TextMultipleChoiceEvaluator.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py index 3cc9d2c058..3aaa4a14ff 100644 --- a/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py +++ b/mteb/evaluation/evaluators/Image/Any2TextMultipleChoiceEvaluator.py @@ -35,7 +35,6 @@ def __init__( label_column_name: str, choices_column_name: str, task_name: str | None = None, - transform=None, limit: int | None = None, **kwargs, ): @@ -48,7 +47,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,