From 4e7b09c66b14fdedb7b0e459a0f2ff1be5607dc0 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Fri, 23 May 2025 15:51:23 +0200 Subject: [PATCH 01/11] add colpali models --- .../Image/Any2AnyRetrievalEvaluator.py | 21 +- mteb/models/colpali_models.py | 286 ++++++++++++++++++ mteb/models/colqwen_models.py | 151 +++++++++ mteb/models/overview.py | 4 + pyproject.toml | 1 + 5 files changed, 454 insertions(+), 9 deletions(-) create mode 100644 mteb/models/colpali_models.py create mode 100644 mteb/models/colqwen_models.py 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/models/colpali_models.py b/mteb/models/colpali_models.py new file mode 100644 index 0000000000..38b6f21420 --- /dev/null +++ b/mteb/models/colpali_models.py @@ -0,0 +1,286 @@ +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=819, + 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=["PyTorch"], + 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=819, + 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=["PyTorch"], + 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=819, + 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=["PyTorch"], + 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=819, + 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=["PyTorch"], + reference="https://huggingface.co/vidore/colpali", + similarity_fn_name="max_sim", + use_instructions=False, + training_datasets=colpali_training_datasets, +) + +colqwen2 = ModelMeta( + loader=partial( + ColPaliWrapper, + model_name="vidore/colqwen2-v1.0-merged", + ), + name="vidore/colqwen2-v1.0", + languages=["eng-Latn"], + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-02-11", + modalities=["image", "text"], + n_parameters=2_210_000_000, + memory_usage_mb=819, + max_tokens=32768, + embed_dim=1536, + 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=["PyTorch"], + reference="https://huggingface.co/vidore/colqwen2-v1.0", + similarity_fn_name="max_sim", + use_instructions=False, + training_datasets=colpali_training_datasets, +) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py new file mode 100644 index 0000000000..48f9e215d5 --- /dev/null +++ b/mteb/models/colqwen_models.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import logging +from functools import partial +from typing import Literal + +from mteb.model_meta import ModelMeta +from mteb.models.colpali_models import ColPaliEngineWrapper + +logger = logging.getLogger(__name__) + +EncodeTypes = Literal["query", "passage"] + + +class ColQwen2Wrapper(ColPaliEngineWrapper): + """Wrapper for ColQwen2 model.""" + + def __init__( + self, model_name: str = "vidore/colqwen2-v1.0", device: str = None, **kwargs + ): + from colpali_engine.models import ColQwen2, ColQwen2Processor + + super().__init__( + model_name=model_name, + model_class=ColQwen2, + processor_class=ColQwen2Processor, + device=device, + **kwargs, + ) + + +class ColQwen2_5Wrapper(ColPaliEngineWrapper): + """Wrapper for ColQwen2.5 model.""" + + def __init__( + self, model_name: str = "vidore/colqwen2.5-v0.2", device: str = None, **kwargs + ): + from colpali_engine.models import ColQwen2_5, ColQwen2_5Processor + + super().__init__( + model_name=model_name, + model_class=ColQwen2_5, + processor_class=ColQwen2_5Processor, + device=device, + **kwargs, + ) + + +colpali_training_datasets = { + # TODO: Add the training datasets here +} + +colqwen2 = ModelMeta( + loader=partial( + ColQwen2Wrapper, + model_name="vidore/colqwen2-v1.0-merged", # TODO: Understand why merged works but not peft one + ), + 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=819, + max_tokens=32768, + embed_dim=1536, + 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=["PyTorch"], + reference="https://huggingface.co/vidore/colqwen2-v1.0", + similarity_fn_name="max_sim", + use_instructions=False, + training_datasets=colpali_training_datasets, +) + +colqwen2_5 = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="vidore/colqwen2.5-v0.2", + ), + 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=819, + max_tokens=128000, + embed_dim=1536, + 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=["PyTorch"], + reference="https://huggingface.co/vidore/colqwen2.5-v0.2", + similarity_fn_name="max_sim", + use_instructions=False, + training_datasets=colpali_training_datasets, +) + +colnomic_7b = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="nomic-ai/colnomic-embed-multimodal-7b", + ), + 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=819, + max_tokens=128000, + embed_dim=1536, + 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=["PyTorch"], + reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-7b", + similarity_fn_name="max_sim", + use_instructions=False, + training_datasets=colpali_training_datasets, +) + +colnomic_3b = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="nomic-ai/colnomic-embed-multimodal-3b", + ), + name="nomic-ai/colnomic-embed-multimodal-3b", + languages=["eng-Latn"], + revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", + release_date="2025-03-31", + modalities=["image", "text"], + n_parameters=3_000_000_000, + memory_usage_mb=819, + max_tokens=128000, + embed_dim=1536, + 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=["PyTorch"], + reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-3b", + similarity_fn_name="max_sim", + use_instructions=False, + training_datasets=colpali_training_datasets, +) diff --git a/mteb/models/overview.py b/mteb/models/overview.py index 39aef262de..cc696bd2e5 100644 --- a/mteb/models/overview.py +++ b/mteb/models/overview.py @@ -27,6 +27,8 @@ cohere_models, cohere_v, colbert_models, + colpali_models, + colqwen_models, conan_models, dino_models, e5_instruct, @@ -161,6 +163,8 @@ ara_models, b1ade_models, nb_sbert, + colpali_models, + colqwen_models, ] MODEL_REGISTRY = {} diff --git a/pyproject.toml b/pyproject.toml index e0b081ff27..9778518093 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,7 @@ vertexai = ["vertexai==1.71.1"] 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"] +colpali = ["colpali_engine>=0.3.10"] [tool.coverage.report] From d328d48c9e8637e6a8f9e369bf98e186dfe6f96f Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Fri, 23 May 2025 15:57:45 +0200 Subject: [PATCH 02/11] add colpali as framework --- mteb/model_meta.py | 1 + mteb/models/colpali_models.py | 10 +++++----- mteb/models/colqwen_models.py | 8 ++++---- 3 files changed, 10 insertions(+), 9 deletions(-) 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 index 38b6f21420..aa0da4ca1b 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -178,7 +178,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/vidore/colpali-v1.3", similarity_fn_name="max_sim", use_instructions=False, @@ -203,7 +203,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/vidore/colpali-v1.2", similarity_fn_name="max_sim", use_instructions=False, @@ -228,7 +228,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/vidore/colpali-v1.1", similarity_fn_name="max_sim", use_instructions=False, @@ -253,7 +253,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/vidore/colpali", similarity_fn_name="max_sim", use_instructions=False, @@ -278,7 +278,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/vidore/colqwen2-v1.0", similarity_fn_name="max_sim", use_instructions=False, diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index 48f9e215d5..8aa6e23251 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -68,7 +68,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/vidore/colqwen2-v1.0", similarity_fn_name="max_sim", use_instructions=False, @@ -93,7 +93,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["colpali"], reference="https://huggingface.co/vidore/colqwen2.5-v0.2", similarity_fn_name="max_sim", use_instructions=False, @@ -118,7 +118,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/nomic-ai/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-7b", similarity_fn_name="max_sim", use_instructions=False, @@ -143,7 +143,7 @@ def __init__( open_weights=True, public_training_code="https://github.com/nomic-ai/colpali", public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set", - framework=["PyTorch"], + framework=["ColPali"], reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-3b", similarity_fn_name="max_sim", use_instructions=False, From 76892adf1cdb0a16dbd8f4662556c882c34cf1ab Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Fri, 23 May 2025 15:59:20 +0200 Subject: [PATCH 03/11] add colpali as framework --- mteb/models/colqwen_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index 8aa6e23251..a061d20188 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -93,7 +93,7 @@ def __init__( 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"], + framework=["ColPali"], reference="https://huggingface.co/vidore/colqwen2.5-v0.2", similarity_fn_name="max_sim", use_instructions=False, From 58010f7ee99d8096056f703db92458cfeb01ca47 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Fri, 23 May 2025 16:25:20 +0200 Subject: [PATCH 04/11] update metadata and add colsmol --- mteb/models/colpali_models.py | 10 ++--- mteb/models/colqwen_models.py | 8 ++-- mteb/models/colsmol_models.py | 83 +++++++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 9 deletions(-) create mode 100644 mteb/models/colsmol_models.py diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py index aa0da4ca1b..f65ce6d60f 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -171,7 +171,7 @@ def __init__( release_date="2024-11-01", modalities=["image", "text"], n_parameters=2_920_000_000, - memory_usage_mb=819, + memory_usage_mb=4700, max_tokens=16384, embed_dim=2048, license="apache-2.0", @@ -196,7 +196,7 @@ def __init__( release_date="2024-08-26", modalities=["image", "text"], n_parameters=2_920_000_000, - memory_usage_mb=819, + memory_usage_mb=4700, max_tokens=16384, embed_dim=2048, license="apache-2.0", @@ -221,7 +221,7 @@ def __init__( release_date="2024-08-21", modalities=["image", "text"], n_parameters=2_920_000_000, - memory_usage_mb=819, + memory_usage_mb=4700, max_tokens=16384, embed_dim=2048, license="apache-2.0", @@ -246,7 +246,7 @@ def __init__( release_date="2024-06-25", modalities=["image", "text"], n_parameters=2_920_000_000, - memory_usage_mb=819, + memory_usage_mb=4700, max_tokens=16384, embed_dim=2048, license="apache-2.0", @@ -271,7 +271,7 @@ def __init__( release_date="2025-02-11", modalities=["image", "text"], n_parameters=2_210_000_000, - memory_usage_mb=819, + memory_usage_mb=4700, max_tokens=32768, embed_dim=1536, license="apache-2.0", diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index a061d20188..0ebd7d8373 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -61,7 +61,7 @@ def __init__( release_date="2025-11-03", modalities=["image", "text"], n_parameters=2_210_000_000, - memory_usage_mb=819, + memory_usage_mb=7200, max_tokens=32768, embed_dim=1536, license="apache-2.0", @@ -86,7 +86,7 @@ def __init__( release_date="2025-01-31", modalities=["image", "text"], n_parameters=3_000_000_000, - memory_usage_mb=819, + memory_usage_mb=7200, max_tokens=128000, embed_dim=1536, license="apache-2.0", @@ -111,7 +111,7 @@ def __init__( release_date="2025-03-31", modalities=["image", "text"], n_parameters=7_000_000_000, - memory_usage_mb=819, + memory_usage_mb=14400, max_tokens=128000, embed_dim=1536, license="apache-2.0", @@ -136,7 +136,7 @@ def __init__( release_date="2025-03-31", modalities=["image", "text"], n_parameters=3_000_000_000, - memory_usage_mb=819, + memory_usage_mb=7200, max_tokens=128000, embed_dim=1536, license="apache-2.0", diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py new file mode 100644 index 0000000000..905858e7df --- /dev/null +++ b/mteb/models/colsmol_models.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import logging +from functools import partial +from typing import Literal + +from mteb.model_meta import ModelMeta +from mteb.models.colpali_models import ColPaliEngineWrapper + +logger = logging.getLogger(__name__) + +EncodeTypes = Literal["query", "passage"] + + +class ColSmolWrapper(ColPaliEngineWrapper): + """Wrapper for ColQwen2 model.""" + + def __init__( + self, model_name: str = "vidore/colqwen2-v1.0", device: str = None, **kwargs + ): + from colpali_engine.models import ColIdefics3, ColIdefics3Processor + + super().__init__( + model_name=model_name, + model_class=ColIdefics3, + processor_class=ColIdefics3Processor, + device=device, + **kwargs, + ) + +colpali_training_datasets = { + # TODO: Add the training datasets here +} + +colsmol_256m = ModelMeta( + loader=partial( + ColSmolWrapper, + model_name="vidore/colSmol-256M", + ), + 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=576, + 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=False, + training_datasets=colpali_training_datasets, +) + +colsmol_500m = ModelMeta( + loader=partial( + ColSmolWrapper, + model_name="vidore/colSmol-500M", + ), + 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=576, + 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=False, + training_datasets=colpali_training_datasets, +) \ No newline at end of file From 41dcc9c975e7aefc1364dd83e3f9e3cc1247a9d0 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Fri, 23 May 2025 16:38:37 +0200 Subject: [PATCH 05/11] ix typos --- mteb/models/colpali_models.py | 25 ------------------------- mteb/models/colqwen_models.py | 2 +- mteb/models/colsmol_models.py | 3 ++- mteb/models/overview.py | 2 ++ 4 files changed, 5 insertions(+), 27 deletions(-) diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py index f65ce6d60f..c896439ce8 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -259,28 +259,3 @@ def __init__( use_instructions=False, training_datasets=colpali_training_datasets, ) - -colqwen2 = ModelMeta( - loader=partial( - ColPaliWrapper, - model_name="vidore/colqwen2-v1.0-merged", - ), - name="vidore/colqwen2-v1.0", - languages=["eng-Latn"], - revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", - release_date="2025-02-11", - modalities=["image", "text"], - n_parameters=2_210_000_000, - memory_usage_mb=4700, - max_tokens=32768, - embed_dim=1536, - 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=False, - training_datasets=colpali_training_datasets, -) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index 0ebd7d8373..b291561f5c 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -53,7 +53,7 @@ def __init__( colqwen2 = ModelMeta( loader=partial( ColQwen2Wrapper, - model_name="vidore/colqwen2-v1.0-merged", # TODO: Understand why merged works but not peft one + model_name="vidore/colqwen2-v1.0", ), name="vidore/colqwen2-v1.0", languages=["eng-Latn"], diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py index 905858e7df..b846ca23fb 100644 --- a/mteb/models/colsmol_models.py +++ b/mteb/models/colsmol_models.py @@ -28,6 +28,7 @@ def __init__( **kwargs, ) + colpali_training_datasets = { # TODO: Add the training datasets here } @@ -80,4 +81,4 @@ def __init__( similarity_fn_name="max_sim", use_instructions=False, training_datasets=colpali_training_datasets, -) \ No newline at end of file +) diff --git a/mteb/models/overview.py b/mteb/models/overview.py index cc696bd2e5..5afe2e50a0 100644 --- a/mteb/models/overview.py +++ b/mteb/models/overview.py @@ -29,6 +29,7 @@ colbert_models, colpali_models, colqwen_models, + colsmol_models, conan_models, dino_models, e5_instruct, @@ -165,6 +166,7 @@ nb_sbert, colpali_models, colqwen_models, + colsmol_models, ] MODEL_REGISTRY = {} From 26f03fdbf7bcbbd3e04eb01de883dc112fa97754 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Fri, 23 May 2025 18:23:21 +0200 Subject: [PATCH 06/11] account for revision --- mteb/models/colpali_models.py | 8 +++++++- mteb/models/colqwen_models.py | 23 +++++++++++++++++++++-- mteb/models/colsmol_models.py | 13 ++++++++++++- 3 files changed, 40 insertions(+), 4 deletions(-) diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py index c896439ce8..de00680937 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -29,6 +29,7 @@ def __init__( model_name: str, model_class: type, processor_class: type, + revision: str | None = None, device: str = None, **kwargs, ): @@ -48,7 +49,7 @@ def __init__( # Load model self.mdl = model_class.from_pretrained( - model_name, trust_remote_code=True, **kwargs + model_name, trust_remote_code=True, revision=revision, **kwargs ) self.mdl.eval().to(self.device) @@ -142,15 +143,20 @@ class ColPaliWrapper(ColPaliEngineWrapper): def __init__( self, model_name: str = "vidore/colpali-v1.3", + revision: str | None = None, device: str = "cuda" if torch.cuda.is_available() else "cpu", **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, ) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index b291561f5c..6a7cb771dd 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -6,6 +6,9 @@ from mteb.model_meta import ModelMeta from mteb.models.colpali_models import ColPaliEngineWrapper +from mteb.requires_package import ( + requires_package, +) logger = logging.getLogger(__name__) @@ -16,14 +19,22 @@ class ColQwen2Wrapper(ColPaliEngineWrapper): """Wrapper for ColQwen2 model.""" def __init__( - self, model_name: str = "vidore/colqwen2-v1.0", device: str = None, **kwargs + self, + model_name: str = "vidore/colqwen2-v1.0", + revision: str | None = None, + device: str = 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, ) @@ -33,14 +44,22 @@ class ColQwen2_5Wrapper(ColPaliEngineWrapper): """Wrapper for ColQwen2.5 model.""" def __init__( - self, model_name: str = "vidore/colqwen2.5-v0.2", device: str = None, **kwargs + self, + model_name: str = "vidore/colqwen2.5-v0.2", + revision: str | None = None, + device: str = None, + **kwargs, ): + requires_package( + self, "colpali_engine", model_name, "pip install mteb[colpali_engine]" + ) from colpali_engine.models import ColQwen2_5, ColQwen2_5Processor super().__init__( model_name=model_name, model_class=ColQwen2_5, processor_class=ColQwen2_5Processor, + revision=revision, device=device, **kwargs, ) diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py index b846ca23fb..59cc74cb7e 100644 --- a/mteb/models/colsmol_models.py +++ b/mteb/models/colsmol_models.py @@ -6,6 +6,9 @@ from mteb.model_meta import ModelMeta from mteb.models.colpali_models import ColPaliEngineWrapper +from mteb.requires_package import ( + requires_package, +) logger = logging.getLogger(__name__) @@ -16,14 +19,22 @@ class ColSmolWrapper(ColPaliEngineWrapper): """Wrapper for ColQwen2 model.""" def __init__( - self, model_name: str = "vidore/colqwen2-v1.0", device: str = None, **kwargs + self, + model_name: str = "vidore/colqwen2-v1.0", + revision: str | None = None, + device: str = 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, ) From 9905c65767adeb07df326f1b616d1d18ae947389 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Mon, 26 May 2025 17:52:11 +0200 Subject: [PATCH 07/11] add training data info and lint --- mteb/models/colpali_models.py | 20 +++++++++++--------- mteb/models/colqwen_models.py | 19 +++++++------------ mteb/models/colsmol_models.py | 13 +++---------- 3 files changed, 21 insertions(+), 31 deletions(-) diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py index de00680937..d5805185bb 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -2,7 +2,7 @@ import logging from functools import partial -from typing import Any, Literal +from typing import Any import torch from PIL import Image @@ -18,8 +18,6 @@ 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""" @@ -162,8 +160,12 @@ def __init__( ) -colpali_training_datasets = { - # TODO: Add the training datasets here +COLPALI_TRAINING_DATA = { + # from https://huggingface.co/datasets/vidore/colpali_train_set + "DocVQA": ["train"], + "InfoVQA": ["train"], + "TATDQA": ["train"], + "arXivQA": ["train"], } colpali_v1_3 = ModelMeta( @@ -188,7 +190,7 @@ def __init__( reference="https://huggingface.co/vidore/colpali-v1.3", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) colpali_v1_2 = ModelMeta( @@ -213,7 +215,7 @@ def __init__( reference="https://huggingface.co/vidore/colpali-v1.2", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) colpali_v1_1 = ModelMeta( @@ -238,7 +240,7 @@ def __init__( reference="https://huggingface.co/vidore/colpali-v1.1", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) colpali = ModelMeta( @@ -263,5 +265,5 @@ def __init__( reference="https://huggingface.co/vidore/colpali", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index 6a7cb771dd..56539c7d5d 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -2,18 +2,15 @@ import logging from functools import partial -from typing import Literal from mteb.model_meta import ModelMeta -from mteb.models.colpali_models import ColPaliEngineWrapper +from mteb.models.colpali_models import COLPALI_TRAINING_DATA, ColPaliEngineWrapper from mteb.requires_package import ( requires_package, ) logger = logging.getLogger(__name__) -EncodeTypes = Literal["query", "passage"] - class ColQwen2Wrapper(ColPaliEngineWrapper): """Wrapper for ColQwen2 model.""" @@ -65,10 +62,6 @@ def __init__( ) -colpali_training_datasets = { - # TODO: Add the training datasets here -} - colqwen2 = ModelMeta( loader=partial( ColQwen2Wrapper, @@ -91,7 +84,7 @@ def __init__( reference="https://huggingface.co/vidore/colqwen2-v1.0", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) colqwen2_5 = ModelMeta( @@ -116,7 +109,7 @@ def __init__( reference="https://huggingface.co/vidore/colqwen2.5-v0.2", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) colnomic_7b = ModelMeta( @@ -141,9 +134,11 @@ def __init__( reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-7b", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) +COLNOMIC_TRAINING_DATA = {"VDRMultilingual": ["Train"], **COLPALI_TRAINING_DATA} + colnomic_3b = ModelMeta( loader=partial( ColQwen2_5Wrapper, @@ -166,5 +161,5 @@ def __init__( reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-3b", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLNOMIC_TRAINING_DATA, ) diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py index 59cc74cb7e..5ae99d16ea 100644 --- a/mteb/models/colsmol_models.py +++ b/mteb/models/colsmol_models.py @@ -2,18 +2,15 @@ import logging from functools import partial -from typing import Literal from mteb.model_meta import ModelMeta -from mteb.models.colpali_models import ColPaliEngineWrapper +from mteb.models.colpali_models import COLPALI_TRAINING_DATA, ColPaliEngineWrapper from mteb.requires_package import ( requires_package, ) logger = logging.getLogger(__name__) -EncodeTypes = Literal["query", "passage"] - class ColSmolWrapper(ColPaliEngineWrapper): """Wrapper for ColQwen2 model.""" @@ -40,10 +37,6 @@ def __init__( ) -colpali_training_datasets = { - # TODO: Add the training datasets here -} - colsmol_256m = ModelMeta( loader=partial( ColSmolWrapper, @@ -66,7 +59,7 @@ def __init__( reference="https://huggingface.co/vidore/colSmol-256M", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) colsmol_500m = ModelMeta( @@ -91,5 +84,5 @@ def __init__( reference="https://huggingface.co/vidore/colSmol-500M", similarity_fn_name="max_sim", use_instructions=False, - training_datasets=colpali_training_datasets, + training_datasets=COLPALI_TRAINING_DATA, ) From b6d62a2d062a3e23cb343a874900c5e8fba16a26 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Mon, 26 May 2025 17:56:56 +0200 Subject: [PATCH 08/11] modify meta --- mteb/models/colpali_models.py | 8 ++++---- mteb/models/colqwen_models.py | 8 ++++---- mteb/models/colsmol_models.py | 4 ++-- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py index d5805185bb..666c5600fe 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -189,7 +189,7 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colpali-v1.3", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) @@ -214,7 +214,7 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colpali-v1.2", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) @@ -239,7 +239,7 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colpali-v1.1", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) @@ -264,6 +264,6 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colpali", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index 56539c7d5d..3371e63de0 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -83,7 +83,7 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colqwen2-v1.0", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) @@ -108,7 +108,7 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colqwen2.5-v0.2", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) @@ -133,7 +133,7 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-7b", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) @@ -160,6 +160,6 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/nomic-ai/colnomic-embed-multimodal-3b", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLNOMIC_TRAINING_DATA, ) diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py index 5ae99d16ea..526b8ac946 100644 --- a/mteb/models/colsmol_models.py +++ b/mteb/models/colsmol_models.py @@ -58,7 +58,7 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colSmol-256M", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) @@ -83,6 +83,6 @@ def __init__( framework=["ColPali"], reference="https://huggingface.co/vidore/colSmol-500M", similarity_fn_name="max_sim", - use_instructions=False, + use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, ) From 1721e7fde76368369fc804145c71d65b25c77c65 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Tue, 27 May 2025 11:00:23 +0200 Subject: [PATCH 09/11] correct colmodels meta and add colnomic 7b --- mteb/models/colpali_models.py | 53 +++++++++-------------------------- mteb/models/colqwen_models.py | 37 ++++++++++++++++++++---- mteb/models/colsmol_models.py | 4 +-- 3 files changed, 47 insertions(+), 47 deletions(-) diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py index 666c5600fe..30434b4202 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -168,26 +168,26 @@ def __init__( "arXivQA": ["train"], } -colpali_v1_3 = ModelMeta( +colpali_v1_1 = ModelMeta( loader=partial( ColPaliWrapper, - model_name="vidore/colpali-v1.3", + model_name="vidore/colpali-v1.1", ), - name="vidore/colpali-v1.3", + name="vidore/colpali-v1.1", languages=["eng-Latn"], revision="1b5c8929330df1a66de441a9b5409a878f0de5b0", - release_date="2024-11-01", + release_date="2024-08-21", modalities=["image", "text"], n_parameters=2_920_000_000, memory_usage_mb=4700, max_tokens=16384, - embed_dim=2048, + 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", + reference="https://huggingface.co/vidore/colpali-v1.1", similarity_fn_name="max_sim", use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, @@ -206,7 +206,7 @@ def __init__( n_parameters=2_920_000_000, memory_usage_mb=4700, max_tokens=16384, - embed_dim=2048, + embed_dim=128, license="apache-2.0", open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", @@ -218,52 +218,27 @@ def __init__( training_datasets=COLPALI_TRAINING_DATA, ) -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=True, - training_datasets=COLPALI_TRAINING_DATA, -) - -colpali = ModelMeta( +colpali_v1_3 = ModelMeta( loader=partial( ColPaliWrapper, - model_name="vidore/colpali", + model_name="vidore/colpali-v1.3", ), - name="vidore/colpali", + name="vidore/colpali-v1.3", languages=["eng-Latn"], revision="1b5c8929330df1a66de441a9b5409a878f0de5b0", - release_date="2024-06-25", + release_date="2024-11-01", modalities=["image", "text"], n_parameters=2_920_000_000, memory_usage_mb=4700, max_tokens=16384, - embed_dim=2048, + 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", + reference="https://huggingface.co/vidore/colpali-v1.3", similarity_fn_name="max_sim", use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, -) +) \ No newline at end of file diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index 3371e63de0..a08edbcaad 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -50,12 +50,12 @@ def __init__( requires_package( self, "colpali_engine", model_name, "pip install mteb[colpali_engine]" ) - from colpali_engine.models import ColQwen2_5, ColQwen2_5Processor + from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor super().__init__( model_name=model_name, model_class=ColQwen2_5, - processor_class=ColQwen2_5Processor, + processor_class=ColQwen2_5_Processor, revision=revision, device=device, **kwargs, @@ -75,7 +75,7 @@ def __init__( n_parameters=2_210_000_000, memory_usage_mb=7200, max_tokens=32768, - embed_dim=1536, + embed_dim=128, license="apache-2.0", open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", @@ -100,7 +100,7 @@ def __init__( n_parameters=3_000_000_000, memory_usage_mb=7200, max_tokens=128000, - embed_dim=1536, + embed_dim=128, license="apache-2.0", open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", @@ -125,7 +125,7 @@ def __init__( n_parameters=7_000_000_000, memory_usage_mb=14400, max_tokens=128000, - embed_dim=1536, + embed_dim=128, license="apache-2.0", open_weights=True, public_training_code="https://github.com/nomic-ai/colpali", @@ -152,7 +152,7 @@ def __init__( n_parameters=3_000_000_000, memory_usage_mb=7200, max_tokens=128000, - embed_dim=1536, + embed_dim=128, license="apache-2.0", open_weights=True, public_training_code="https://github.com/nomic-ai/colpali", @@ -163,3 +163,28 @@ def __init__( use_instructions=True, training_datasets=COLNOMIC_TRAINING_DATA, ) + +colnomic_7b = ModelMeta( + loader=partial( + ColQwen2_5Wrapper, + model_name="nomic-ai/colnomic-embed-multimodal-7b", + ), + 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=COLNOMIC_TRAINING_DATA, +) \ No newline at end of file diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py index 526b8ac946..c6c394e778 100644 --- a/mteb/models/colsmol_models.py +++ b/mteb/models/colsmol_models.py @@ -50,7 +50,7 @@ def __init__( n_parameters=256_000_000, memory_usage_mb=800, max_tokens=8192, - embed_dim=576, + embed_dim=128, license="apache-2.0", open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", @@ -75,7 +75,7 @@ def __init__( n_parameters=500_000_000, memory_usage_mb=1200, max_tokens=8192, - embed_dim=576, + embed_dim=128, license="apache-2.0", open_weights=True, public_training_code="https://github.com/illuin-tech/colpali", From c07da0b0827996011229cd5b87f10dc3fcf02230 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Tue, 27 May 2025 11:00:56 +0200 Subject: [PATCH 10/11] fix typo in toml (colpali subdeps) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 1ce6b4ff13..bbba817621 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,7 +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 = ["colpali_engine>=0.3.10"] +colpali_engine = ["colpali_engine>=0.3.10"] [tool.coverage.report] From 60d610a1a4fdf13a73f89264c4d0c2f1085ca672 Mon Sep 17 00:00:00 2001 From: Paul Teiletche Date: Tue, 27 May 2025 14:48:13 +0200 Subject: [PATCH 11/11] refine colmodel loading and metadata --- mteb/models/colpali_models.py | 27 +++++++++------------- mteb/models/colqwen_models.py | 42 ++++++++++++++++++++++++++++++----- mteb/models/colsmol_models.py | 13 ++++++++++- 3 files changed, 59 insertions(+), 23 deletions(-) diff --git a/mteb/models/colpali_models.py b/mteb/models/colpali_models.py index 30434b4202..04f0c21755 100644 --- a/mteb/models/colpali_models.py +++ b/mteb/models/colpali_models.py @@ -13,14 +13,13 @@ from mteb.requires_package import ( requires_image_dependencies, requires_package, - suggest_package, ) logger = logging.getLogger(__name__) class ColPaliEngineWrapper: - """Base wrapper for ColPali models. Adapted from https://github.com/illuin-tech/colpali/tree/bebcdd6715dba42624acd8d7f7222a16a5daf848/colpali_engine/models""" + """Base wrapper for `colpali_engine` models. Adapted from https://github.com/illuin-tech/colpali/tree/bebcdd6715dba42624acd8d7f7222a16a5daf848/colpali_engine/models""" def __init__( self, @@ -28,17 +27,10 @@ def __init__( model_class: type, processor_class: type, revision: str | None = None, - device: str = None, + device: str | None = 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]" ) @@ -47,9 +39,9 @@ def __init__( # Load model self.mdl = model_class.from_pretrained( - model_name, trust_remote_code=True, revision=revision, **kwargs + model_name, revision=revision, device_map=self.device, **kwargs ) - self.mdl.eval().to(self.device) + self.mdl.eval() # Load processor self.processor = processor_class.from_pretrained(model_name) @@ -142,7 +134,7 @@ def __init__( self, model_name: str = "vidore/colpali-v1.3", revision: str | None = None, - device: str = "cuda" if torch.cuda.is_available() else "cpu", + device: str | None = None, **kwargs, ): requires_package( @@ -172,10 +164,11 @@ def __init__( loader=partial( ColPaliWrapper, model_name="vidore/colpali-v1.1", + torch_dtype=torch.float16, ), name="vidore/colpali-v1.1", languages=["eng-Latn"], - revision="1b5c8929330df1a66de441a9b5409a878f0de5b0", + revision="a0f15e3bcf97110e7ac1bb4be4bcd30eeb31992a", release_date="2024-08-21", modalities=["image", "text"], n_parameters=2_920_000_000, @@ -197,10 +190,11 @@ def __init__( loader=partial( ColPaliWrapper, model_name="vidore/colpali-v1.2", + torch_dtype=torch.float16, ), name="vidore/colpali-v1.2", languages=["eng-Latn"], - revision="1b5c8929330df1a66de441a9b5409a878f0de5b0", + revision="6b89bc63c16809af4d111bfe412e2ac6bc3c9451", release_date="2024-08-26", modalities=["image", "text"], n_parameters=2_920_000_000, @@ -222,6 +216,7 @@ def __init__( loader=partial( ColPaliWrapper, model_name="vidore/colpali-v1.3", + torch_dtype=torch.float16, ), name="vidore/colpali-v1.3", languages=["eng-Latn"], @@ -241,4 +236,4 @@ def __init__( similarity_fn_name="max_sim", use_instructions=True, training_datasets=COLPALI_TRAINING_DATA, -) \ No newline at end of file +) diff --git a/mteb/models/colqwen_models.py b/mteb/models/colqwen_models.py index a08edbcaad..88724e2c73 100644 --- a/mteb/models/colqwen_models.py +++ b/mteb/models/colqwen_models.py @@ -3,6 +3,9 @@ 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 ( @@ -19,7 +22,7 @@ def __init__( self, model_name: str = "vidore/colqwen2-v1.0", revision: str | None = None, - device: str = None, + device: str | None = None, **kwargs, ): requires_package( @@ -44,7 +47,7 @@ def __init__( self, model_name: str = "vidore/colqwen2.5-v0.2", revision: str | None = None, - device: str = None, + device: str | None = None, **kwargs, ): requires_package( @@ -66,6 +69,10 @@ def __init__( 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"], @@ -91,6 +98,10 @@ def __init__( 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"], @@ -116,6 +127,10 @@ def __init__( 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"], @@ -138,14 +153,25 @@ def __init__( ) 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=["eng-Latn"], + languages=COLNOMIC_LANGUAGES, revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f", release_date="2025-03-31", modalities=["image", "text"], @@ -168,10 +194,14 @@ def __init__( 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", + languages=COLNOMIC_LANGUAGES, + revision="09dbc9502b66605d5be56d2226019b49c9fd3293", release_date="2025-03-31", modalities=["image", "text"], n_parameters=7_000_000_000, @@ -187,4 +217,4 @@ def __init__( similarity_fn_name="max_sim", use_instructions=True, training_datasets=COLNOMIC_TRAINING_DATA, -) \ No newline at end of file +) diff --git a/mteb/models/colsmol_models.py b/mteb/models/colsmol_models.py index c6c394e778..c02fe365d5 100644 --- a/mteb/models/colsmol_models.py +++ b/mteb/models/colsmol_models.py @@ -3,6 +3,9 @@ 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 ( @@ -19,7 +22,7 @@ def __init__( self, model_name: str = "vidore/colqwen2-v1.0", revision: str | None = None, - device: str = None, + device: str | None = None, **kwargs, ): requires_package( @@ -41,6 +44,10 @@ def __init__( 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"], @@ -66,6 +73,10 @@ def __init__( 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"],