fix: Add colpali models family

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(

mteb/model_meta.py

@@ -39,6 +39,7 @@
     "NumPy",
     "PyLate",
     "ColBERT",
+    "ColPali",
 ]
 DISTANCE_METRICS = Literal["cosine", "max_sim", "dot"]
 

mteb/models/colpali_models.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import logging
+from functools import partial
+from typing import Any, Literal
+
+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,
+    suggest_package,
+)
+
+logger = logging.getLogger(__name__)
+
+EncodeTypes = Literal["query", "passage"]
+
+
+class ColPaliEngineWrapper:
+    """Base wrapper for ColPali 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,
+        device: str = None,
+        **kwargs,
+    ):
+        requires_image_dependencies()
+        if suggest_package(
+            self,
+            "flash_attn",
+            model_name,
+            "pip install flash-attn --no-build-isolation",
+        ):
+            import flash_attn  # noqa
+        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, trust_remote_code=True, **kwargs
+        )
+        self.mdl.eval().to(self.device)
+
+        # 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",
+        device: str = "cuda" if torch.cuda.is_available() else "cpu",
+        **kwargs,
+    ):
+        from colpali_engine.models import ColPali, ColPaliProcessor
+
+        super().__init__(
+            model_name=model_name,
+            model_class=ColPali,
+            processor_class=ColPaliProcessor,
+            device=device,
+            **kwargs,
+        )
+
+
+colpali_training_datasets = {
+    # TODO: Add the training datasets here
+}
+
+colpali_v1_3 = ModelMeta(
+    loader=partial(
+        ColPaliWrapper,
+        model_name="vidore/colpali-v1.3",
+    ),
+    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=2048,
+    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=False,
+    training_datasets=colpali_training_datasets,
+)
+
+colpali_v1_2 = ModelMeta(
+    loader=partial(
+        ColPaliWrapper,
+        model_name="vidore/colpali-v1.2",
+    ),
+    name="vidore/colpali-v1.2",
+    languages=["eng-Latn"],
+    revision="1b5c8929330df1a66de441a9b5409a878f0de5b0",
+    release_date="2024-08-26",
+    modalities=["image", "text"],
+    n_parameters=2_920_000_000,
+    memory_usage_mb=4700,
+    max_tokens=16384,
+    embed_dim=2048,
+    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=False,
+    training_datasets=colpali_training_datasets,
+)
+
+colpali_v1_1 = ModelMeta(
+    loader=partial(
+        ColPaliWrapper,
+        model_name="vidore/colpali-v1.1",
+    ),
+    name="vidore/colpali-v1.1",
+    languages=["eng-Latn"],
+    revision="1b5c8929330df1a66de441a9b5409a878f0de5b0",
+    release_date="2024-08-21",
+    modalities=["image", "text"],
+    n_parameters=2_920_000_000,
+    memory_usage_mb=4700,
+    max_tokens=16384,
+    embed_dim=2048,
+    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=False,
+    training_datasets=colpali_training_datasets,
+)
+
+colpali = ModelMeta(
+    loader=partial(
+        ColPaliWrapper,
+        model_name="vidore/colpali",
+    ),
+    name="vidore/colpali",
+    languages=["eng-Latn"],
+    revision="1b5c8929330df1a66de441a9b5409a878f0de5b0",
+    release_date="2024-06-25",
+    modalities=["image", "text"],
+    n_parameters=2_920_000_000,
+    memory_usage_mb=4700,
+    max_tokens=16384,
+    embed_dim=2048,
+    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",
+    similarity_fn_name="max_sim",
+    use_instructions=False,
+    training_datasets=colpali_training_datasets,
+)