diff --git a/mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py b/mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py index 74a41fb1a3..06fb66c0a9 100644 --- a/mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py +++ b/mteb/evaluation/evaluators/Image/Any2AnyRetrievalEvaluator.py @@ -117,10 +117,18 @@ def search( return_sorted: bool = False, **kwargs, ) -> dict[str, dict[str, float]]: - if score_function not in self.score_functions: - raise ValueError( - f"score function: {score_function} must be either (cos_sim) for cosine similarity or (dot) for dot product" + if hasattr(self.model, "similarity"): + score_function = self.model.similarity + logger.info("Scoring Function: from model") + else: + if score_function not in self.score_functions: + raise ValueError( + f"score function: {score_function} must be either (cos_sim) for cosine similarity or (dot) for dot product" + ) + logger.info( + f"Scoring Function: {self.score_function_desc[score_function]} ({score_function})" ) + score_function = self.score_functions[score_function] logger.info("Encoding Queries.") query_ids = list(queries["id"]) @@ -172,9 +180,6 @@ def search( corpus_modality = corpus[0]["modality"] logger.info("Encoding Corpus in batches... Warning: This might take a while!") - logger.info( - f"Scoring Function: {self.score_function_desc[score_function]} ({score_function})" - ) result_heaps = {qid: [] for qid in query_ids} for chunk_start in range(0, len(corpus), self.corpus_chunk_size): @@ -223,9 +228,7 @@ def search( else: raise ValueError(f"Unsupported modality: {corpus_modality}") - cos_scores = self.score_functions[score_function]( - query_embeddings, sub_corpus_embeddings - ) + cos_scores = score_function(query_embeddings, sub_corpus_embeddings) cos_scores[torch.isnan(cos_scores)] = -1 cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk( diff --git a/mteb/model_meta.py b/mteb/model_meta.py index 2aeb1ac6d6..baafe7139d 100644 --- a/mteb/model_meta.py +++ b/mteb/model_meta.py @@ -39,6 +39,7 @@ "NumPy", "PyLate", "ColBERT", + "ColPali", ] DISTANCE_METRICS = Literal["cosine", "max_sim", "dot"] diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py new file mode 100644 index 0000000000..04f0c21755 --- /dev/null +++ b/mteb/models/colpali_models.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +import logging +from functools import partial +from typing import Any + +import torch +from PIL import Image +from torch.utils.data import DataLoader + +from mteb.encoder_interface import PromptType +from mteb.model_meta import ModelMeta +from mteb.requires_package import ( + requires_image_dependencies, + requires_package, +) + +logger = logging.getLogger(__name__) + + +class ColPaliEngineWrapper: + """Base wrapper for `colpali_engine` models. Adapted from https://github.com/illuin-tech/colpali/tree/bebcdd6715dba42624acd8d7f7222a16a5daf848/colpali_engine/models""" + + def __init__( + self, + model_name: str, + model_class: type, + processor_class: type, + revision: str | None = None, + device: str | None = None, + **kwargs, + ): + requires_image_dependencies() + requires_package( + self, "colpali_engine", model_name, "pip install mteb[colpali_engine]" + ) + + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + + # Load model + self.mdl = model_class.from_pretrained( + model_name, revision=revision, device_map=self.device, **kwargs + ) + self.mdl.eval() + + # Load processor + self.processor = processor_class.from_pretrained(model_name) + + def encode(self, sentences, **kwargs): + return self.get_text_embeddings(texts=sentences, **kwargs) + + def encode_input(self, inputs): + return self.mdl(**inputs) + + def get_image_embeddings( + self, + images, + batch_size: int = 32, + **kwargs, + ): + import torchvision.transforms.functional as F + + all_embeds = [] + + if isinstance(images, DataLoader): + iterator = images + else: + iterator = DataLoader(images, batch_size=batch_size) + + with torch.no_grad(): + for batch in iterator: + # batch may be list of tensors or PIL + imgs = [ + F.to_pil_image(b.to("cpu")) if not isinstance(b, Image.Image) else b + for b in batch + ] + inputs = self.processor.process_images(imgs) + inputs = {k: v.to(self.device) for k, v in inputs.items()} + outs = self.encode_input(inputs) + all_embeds.extend(outs.cpu().to(torch.float32)) + + padded = torch.nn.utils.rnn.pad_sequence( + all_embeds, batch_first=True, padding_value=0 + ) + return padded + + def get_text_embeddings( + self, + texts, + batch_size: int = 32, + **kwargs, + ): + all_embeds = [] + with torch.no_grad(): + for i in range(0, len(texts), batch_size): + batch = texts[i : i + batch_size] + inputs = self.processor.process_queries(batch) + inputs = {k: v.to(self.device) for k, v in inputs.items()} + outs = self.encode_input(inputs) + all_embeds.extend(outs.cpu().to(torch.float32)) + + padded = torch.nn.utils.rnn.pad_sequence( + all_embeds, batch_first=True, padding_value=0 + ) + return padded + + def get_fused_embeddings( + self, + texts: list[str] | None = None, + images: list[Image.Image] | DataLoader | None = None, + *, + task_name: str | None = None, + prompt_type: PromptType | None = None, + batch_size: int = 32, + fusion_mode="sum", + **kwargs: Any, + ): + raise NotImplementedError( + "Fused embeddings are not supported yet. Please use get_text_embeddings or get_image_embeddings." + ) + + def calculate_probs(self, text_embeddings, image_embeddings): + scores = self.similarity(text_embeddings, image_embeddings) + return (scores * 100).softmax(dim=-1) + + def similarity(self, a, b): + return self.processor.score_multi_vector(a, b) + + +class ColPaliWrapper(ColPaliEngineWrapper): + """Wrapper for ColPali models.""" + + def __init__( + self, + model_name: str = "vidore/colpali-v1.3", + revision: str | None = None, + device: str | None = None, + **kwargs, + ): + requires_package( + self, "colpali_engine", model_name, "pip install mteb[colpali_engine]" + ) + from colpali_engine.models import ColPali, ColPaliProcessor + + super().__init__( + model_name=model_name, + model_class=ColPali, + processor_class=ColPaliProcessor, + revision=revision, + device=device, + **kwargs, + ) + + +COLPALI_TRAINING_DATA = { + # from https://huggingface.co/datasets/vidore/colpali_train_set + "DocVQA": ["train"], + "InfoVQA": ["train"], + "TATDQA": ["train"], + "arXivQA": ["train"], +} + +colpali_v1_1 = ModelMeta( + loader=partial( + ColPaliWrapper, + model_name="vidore/colpali-v1.1", + torch_dtype=torch.float16, + ), + name="vidore/colpali-v1.1", + languages=["eng-Latn"], + revision="a0f15e3bcf97110e7ac1bb4be4bcd30eeb31992a", + release_date="2024-08-21", + modalities=["image", "text"], + n_parameters=2_920_000_000, + memory_usage_mb=4700, + max_tokens=16384, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/illuin-tech/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/vidore/colpali-v1.1", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) + +colpali_v1_2 = ModelMeta( + loader=partial( + ColPaliWrapper, + model_name="vidore/colpali-v1.2", + torch_dtype=torch.float16, + ), + name="vidore/colpali-v1.2", + languages=["eng-Latn"], + revision="6b89bc63c16809af4d111bfe412e2ac6bc3c9451", + release_date="2024-08-26", + modalities=["image", "text"], + n_parameters=2_920_000_000, + memory_usage_mb=4700, + max_tokens=16384, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/illuin-tech/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/vidore/colpali-v1.2", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) + +colpali_v1_3 = ModelMeta( + loader=partial( + ColPaliWrapper, + model_name="vidore/colpali-v1.3", + torch_dtype=torch.float16, + ), + name="vidore/colpali-v1.3", + languages=["eng-Latn"], + revision="1b5c8929330df1a66de441a9b5409a878f0de5b0", + release_date="2024-11-01", + modalities=["image", "text"], + n_parameters=2_920_000_000, + memory_usage_mb=4700, + max_tokens=16384, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/illuin-tech/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/vidore/colpali-v1.3", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py new file mode 100644 index 0000000000..88724e2c73 --- /dev/null +++ b/mteb/models/colqwen_models.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +import logging +from functools import partial + +import torch +from transformers.utils.import_utils import is_flash_attn_2_available + +from mteb.model_meta import ModelMeta +from mteb.models.colpali_models import COLPALI_TRAINING_DATA, ColPaliEngineWrapper +from mteb.requires_package import ( + requires_package, +) + +logger = logging.getLogger(__name__) + + +class ColQwen2Wrapper(ColPaliEngineWrapper): + """Wrapper for ColQwen2 model.""" + + def __init__( + self, + model_name: str = "vidore/colqwen2-v1.0", + revision: str | None = None, + device: str | None = None, + **kwargs, + ): + requires_package( + self, "colpali_engine", model_name, "pip install mteb[colpali_engine]" + ) + from colpali_engine.models import ColQwen2, ColQwen2Processor + + super().__init__( + model_name=model_name, + model_class=ColQwen2, + processor_class=ColQwen2Processor, + revision=revision, + device=device, + **kwargs, + ) + + +class ColQwen2_5Wrapper(ColPaliEngineWrapper): + """Wrapper for ColQwen2.5 model.""" + + def __init__( + self, + model_name: str = "vidore/colqwen2.5-v0.2", + revision: str | None = None, + device: str | None = None, + **kwargs, + ): + requires_package( + self, "colpali_engine", model_name, "pip install mteb[colpali_engine]" + ) + from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor + + super().__init__( + model_name=model_name, + model_class=ColQwen2_5, + processor_class=ColQwen2_5_Processor, + revision=revision, + device=device, + **kwargs, + ) + + +colqwen2 = ModelMeta( + loader=partial( + ColQwen2Wrapper, + model_name="vidore/colqwen2-v1.0", + torch_dtype=torch.float16, + attn_implementation="flash_attention_2" + if is_flash_attn_2_available() + else None, + ), + name="vidore/colqwen2-v1.0", + languages=["eng-Latn"], + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-11-03", + modalities=["image", "text"], + n_parameters=2_210_000_000, + memory_usage_mb=7200, + max_tokens=32768, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/illuin-tech/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/vidore/colqwen2-v1.0", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) + +colqwen2_5 = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="vidore/colqwen2.5-v0.2", + torch_dtype=torch.float16, + attn_implementation="flash_attention_2" + if is_flash_attn_2_available() + else None, + ), + name="vidore/colqwen2.5-v0.2", + languages=["eng-Latn"], + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-01-31", + modalities=["image", "text"], + n_parameters=3_000_000_000, + memory_usage_mb=7200, + max_tokens=128000, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/illuin-tech/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/vidore/colqwen2.5-v0.2", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) + +colnomic_7b = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="nomic-ai/colnomic-embed-multimodal-7b", + torch_dtype=torch.float16, + attn_implementation="flash_attention_2" + if is_flash_attn_2_available() + else None, + ), + name="nomic-ai/colnomic-embed-multimodal-7b", + languages=["eng-Latn"], + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-03-31", + modalities=["image", "text"], + n_parameters=7_000_000_000, + memory_usage_mb=14400, + max_tokens=128000, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/nomic-ai/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-7b", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) + +COLNOMIC_TRAINING_DATA = {"VDRMultilingual": ["Train"], **COLPALI_TRAINING_DATA} +COLNOMIC_LANGUAGES = [ + "deu-Latn", # German + "spa-Latn", # Spanish + "eng-Latn", # English + "fra-Latn", # French + "ita-Latn", # Italian +] + +colnomic_3b = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="nomic-ai/colnomic-embed-multimodal-3b", + torch_dtype=torch.float16, + attn_implementation="flash_attention_2" + if is_flash_attn_2_available() + else None, + ), + name="nomic-ai/colnomic-embed-multimodal-3b", + languages=COLNOMIC_LANGUAGES, + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-03-31", + modalities=["image", "text"], + n_parameters=3_000_000_000, + memory_usage_mb=7200, + max_tokens=128000, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/nomic-ai/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-3b", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLNOMIC_TRAINING_DATA, +) + +colnomic_7b = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="nomic-ai/colnomic-embed-multimodal-7b", + torch_dtype=torch.float16, + attn_implementation="flash_attention_2" + if is_flash_attn_2_available() + else None, + ), + name="nomic-ai/colnomic-embed-multimodal-7b", + languages=COLNOMIC_LANGUAGES, + revision="09dbc9502b66605d5be56d2226019b49c9fd3293", + release_date="2025-03-31", + modalities=["image", "text"], + n_parameters=7_000_000_000, + memory_usage_mb=14400, + max_tokens=128000, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/nomic-ai/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-7b", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLNOMIC_TRAINING_DATA, +) diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py new file mode 100644 index 0000000000..c02fe365d5 --- /dev/null +++ b/mteb/models/colsmol_models.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import logging +from functools import partial + +import torch +from transformers.utils.import_utils import is_flash_attn_2_available + +from mteb.model_meta import ModelMeta +from mteb.models.colpali_models import COLPALI_TRAINING_DATA, ColPaliEngineWrapper +from mteb.requires_package import ( + requires_package, +) + +logger = logging.getLogger(__name__) + + +class ColSmolWrapper(ColPaliEngineWrapper): + """Wrapper for ColQwen2 model.""" + + def __init__( + self, + model_name: str = "vidore/colqwen2-v1.0", + revision: str | None = None, + device: str | None = None, + **kwargs, + ): + requires_package( + self, "colpali_engine", model_name, "pip install mteb[colpali_engine]" + ) + from colpali_engine.models import ColIdefics3, ColIdefics3Processor + + super().__init__( + model_name=model_name, + model_class=ColIdefics3, + processor_class=ColIdefics3Processor, + revision=revision, + device=device, + **kwargs, + ) + + +colsmol_256m = ModelMeta( + loader=partial( + ColSmolWrapper, + model_name="vidore/colSmol-256M", + torch_dtype=torch.float16, + attn_implementation="flash_attention_2" + if is_flash_attn_2_available() + else None, + ), + name="vidore/colSmol-256M", + languages=["eng-Latn"], + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-01-22", + modalities=["image", "text"], + n_parameters=256_000_000, + memory_usage_mb=800, + max_tokens=8192, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/illuin-tech/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/vidore/colSmol-256M", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) + +colsmol_500m = ModelMeta( + loader=partial( + ColSmolWrapper, + model_name="vidore/colSmol-500M", + torch_dtype=torch.float16, + attn_implementation="flash_attention_2" + if is_flash_attn_2_available() + else None, + ), + name="vidore/colSmol-500M", + languages=["eng-Latn"], + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-01-22", + modalities=["image", "text"], + n_parameters=500_000_000, + memory_usage_mb=1200, + max_tokens=8192, + embed_dim=128, + license="apache-2.0", + open_weights=True, + public_training_code="https://github.com/illuin-tech/colpali", + public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", + framework=["ColPali"], + reference="https://huggingface.co/vidore/colSmol-500M", + similarity_fn_name="max_sim", + use_instructions=True, + training_datasets=COLPALI_TRAINING_DATA, +) diff --git a/mteb/models/overview.py b/mteb/models/overview.py index 61d0042098..e7dbf541e2 100644 --- a/mteb/models/overview.py +++ b/mteb/models/overview.py @@ -27,6 +27,9 @@ cohere_models, cohere_v, colbert_models, + colpali_models, + colqwen_models, + colsmol_models, conan_models, dino_models, e5_instruct, @@ -163,6 +166,9 @@ ara_models, b1ade_models, nb_sbert, + colpali_models, + colqwen_models, + colsmol_models, ] MODEL_REGISTRY = {} diff --git a/pyproject.toml b/pyproject.toml index f3f5497f8e..bbba817621 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,7 @@ llm2vec = ["llm2vec>=0.2.3,<0.3.0"] timm = ["timm>=1.0.15,<1.1.0"] open_clip_torch = ["open_clip_torch==2.31.0"] ark = ["volcengine-python-sdk[ark]==3.0.2", "tiktoken>=0.8.0"] +colpali_engine = ["colpali_engine>=0.3.10"] [tool.coverage.report]