mteb 2.1.4__py3-none-any.whl → 2.5.2__py3-none-any.whl

mteb/models/model_implementations/colqwen_models.py

@@ -1,11 +1,18 @@
 import logging
+from typing import Any
 
 import torch
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
 
 from mteb._requires_package import (
+    requires_image_dependencies,
     requires_package,
 )
-from mteb.models.model_meta import ModelMeta
+from mteb.abstasks.task_metadata import TaskMetadata
+from mteb.models.abs_encoder import AbsEncoder
+from mteb.models.model_meta import ModelMeta, ScoringFunction
+from mteb.types import Array, BatchedInput, PromptType
 
 from .colpali_models import (
     COLPALI_CITATION,
@@ -73,12 +80,140 @@ class ColQwen2_5Wrapper(ColPaliEngineWrapper):  # noqa: N801
         )
 
 
+class ColQwen3Wrapper(AbsEncoder):
+    """Wrapper for the ColQwen3 vision-language retrieval model."""
+
+    def __init__(
+        self,
+        model_name: str,
+        *,
+        revision: str | None = None,
+        device: str | None = None,
+        dtype: torch.dtype | str | None = torch.bfloat16,
+        **kwargs: Any,
+    ):
+        requires_image_dependencies()
+        requires_package(self, "transformers", model_name, "pip install mteb[colqwen3]")
+        from transformers import AutoModel, AutoProcessor
+
+        self.device = device or (
+            "cuda"
+            if torch.cuda.is_available()
+            else "mps"
+            if torch.backends.mps.is_available()
+            else "cpu"
+        )
+        self.model = AutoModel.from_pretrained(
+            model_name,
+            revision=revision,
+            dtype=dtype,
+            trust_remote_code=True,
+            **kwargs,
+        ).to(self.device)
+        self.model.eval()
+
+        self.processor = AutoProcessor.from_pretrained(
+            model_name,
+            revision=revision,
+            trust_remote_code=True,
+            max_num_visual_tokens=1280,
+        )
+
+    def encode(
+        self,
+        inputs: DataLoader[BatchedInput],
+        *,
+        task_metadata: TaskMetadata,
+        hf_split: str,
+        hf_subset: str,
+        prompt_type: PromptType | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        if (
+            "text" not in inputs.dataset.features
+            and "image" not in inputs.dataset.features
+        ):
+            raise ValueError("No text or image features found in inputs.")
+        return self.get_fused_embeddings(inputs, **kwargs)
+
+    def _encode_inputs(self, encoded_inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        outputs = self.model(**encoded_inputs)
+        # Avoid boolean casting of tensors when checking for custom attributes.
+        embeddings = getattr(outputs, "embeddings", None)
+        if embeddings is None:
+            embeddings = outputs[0]
+        return embeddings
+
+    def get_fused_embeddings(
+        self,
+        image_texts_pairs: DataLoader[BatchedInput] | None = None,
+        batch_size: int = 32,
+        show_progress_bar: bool = True,
+        fusion_mode="concat",
+        **kwargs: Any,
+    ):
+        import torchvision.transforms.functional as F
+        from PIL import Image
+
+        contains_image = "image" in image_texts_pairs.dataset.features
+        contains_text = "text" in image_texts_pairs.dataset.features
+        contains_both = contains_image and contains_text
+
+        if contains_both:
+            progress_desc = "Encoding images+texts"
+        elif contains_image:
+            progress_desc = "Encoding images"
+        elif contains_text:
+            progress_desc = "Encoding texts"
+        else:
+            raise ValueError("No text or image features found in inputs.")
+
+        all_embeds: list[torch.Tensor] = []
+        with torch.no_grad():
+            for batch in tqdm(
+                image_texts_pairs,
+                disable=not show_progress_bar,
+                desc=progress_desc,
+            ):
+                if contains_image:
+                    imgs = [
+                        F.to_pil_image(b.to(self.device))
+                        if not isinstance(b, Image.Image)
+                        else b
+                        for b in batch["image"]
+                    ]
+                else:
+                    imgs = None
+                if contains_text:
+                    texts = batch["text"]
+                else:
+                    texts = None
+                if contains_both:
+                    assert len(imgs) == len(texts), (
+                        f"The number of texts and images must have the same length, got {len(imgs)} and {len(texts)}"
+                    )
+
+                inputs = self.processor(images=imgs, text=texts)
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                outs = self._encode_inputs(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 similarity(self, a, b):
+        return self.processor.score_multi_vector(a, b, device=self.device)
+
+
 colqwen2 = ModelMeta(
     loader=ColQwen2Wrapper,
     loader_kwargs=dict(
         torch_dtype=torch.float16,
     ),
     name="vidore/colqwen2-v1.0",
+    model_type=["late-interaction"],
     languages=["eng-Latn"],
     revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f",
     release_date="2025-11-03",
@@ -105,6 +240,7 @@ colqwen2_5 = ModelMeta(
         torch_dtype=torch.float16,
     ),
     name="vidore/colqwen2.5-v0.2",
+    model_type=["late-interaction"],
     languages=["eng-Latn"],
     revision="6f6fcdfd1a114dfe365f529701b33d66b9349014",
     release_date="2025-01-31",
@@ -125,12 +261,81 @@ colqwen2_5 = ModelMeta(
     citation=COLPALI_CITATION,
 )
 
+TOMORO_TRAINING_DATA = {
+    "VDRMultilingualRetrieval",
+    # from https://huggingface.co/datasets/vidore/colpali_train_set
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
+    "VisRAG-Ret-Train-Synthetic-data",
+    "VisRAG-Ret-Train-In-domain-data",
+}
+
+TOMORO_CITATION = """
+@misc{huang2025tomoro_colqwen3_embed,
+  title={TomoroAI/tomoro-colqwen3-embed},
+  author={Xin Huang and Kye Min Tan and Albert Phelps},
+  year={2025},
+  url={https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-8b}
+}
+"""
+
+colqwen3_8b = ModelMeta(
+    loader=ColQwen3Wrapper,
+    name="TomoroAI/tomoro-colqwen3-embed-8b",
+    model_type=["late-interaction"],
+    languages=["eng-Latn"],
+    revision="0b9fe28142910e209bbac15b1efe85507c27644f",
+    release_date="2025-11-26",
+    modalities=["image", "text"],
+    n_parameters=8_000_000_000,
+    memory_usage_mb=16724,
+    max_tokens=262144,
+    embed_dim=320,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code="https://github.com/illuin-tech/colpali",
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-8b",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    training_datasets=TOMORO_TRAINING_DATA,
+    citation=TOMORO_CITATION,
+)
+
+colqwen3_4b = ModelMeta(
+    loader=ColQwen3Wrapper,
+    name="TomoroAI/tomoro-colqwen3-embed-4b",
+    model_type=["late-interaction"],
+    languages=["eng-Latn"],
+    revision="6a32fb68598730bf5620fbf18d832c784235c59c",
+    release_date="2025-11-26",
+    modalities=["image", "text"],
+    n_parameters=4_000_000_000,
+    memory_usage_mb=8466,
+    max_tokens=262144,
+    embed_dim=320,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code="https://github.com/illuin-tech/colpali",
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-4b",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    training_datasets=TOMORO_TRAINING_DATA,
+    citation=TOMORO_CITATION,
+)
+
 colnomic_7b = ModelMeta(
     loader=ColQwen2_5Wrapper,
     loader_kwargs=dict(
         torch_dtype=torch.float16,
     ),
     name="nomic-ai/colnomic-embed-multimodal-7b",
+    model_type=["late-interaction"],
     languages=["eng-Latn"],
     revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f",
     release_date="2025-03-31",
@@ -175,6 +380,7 @@ colnomic_3b = ModelMeta(
         torch_dtype=torch.float16, attn_implementation="flash_attention_2"
     ),
     name="nomic-ai/colnomic-embed-multimodal-3b",
+    model_type=["late-interaction"],
     languages=COLNOMIC_LANGUAGES,
     revision="86627b4a9b0cade577851a70afa469084f9863a4",
     release_date="2025-03-31",
@@ -201,6 +407,7 @@ colnomic_7b = ModelMeta(
         torch_dtype=torch.float16,
     ),
     name="nomic-ai/colnomic-embed-multimodal-7b",
+    model_type=["late-interaction"],
     languages=COLNOMIC_LANGUAGES,
     revision="09dbc9502b66605d5be56d2226019b49c9fd3293",
     release_date="2025-03-31",
@@ -220,3 +427,66 @@ colnomic_7b = ModelMeta(
     training_datasets=COLNOMIC_TRAINING_DATA,
     citation=COLNOMIC_CITATION,
 )
+
+
+EVOQWEN_TRAINING_DATA = {
+    # "colpali_train_set",
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
+    "VisRAG-Ret-Train-Synthetic-data",
+    "VisRAG-Ret-Train-In-domain-data",
+}
+
+evoqwen25_vl_retriever_3b_v1 = ModelMeta(
+    loader=ColQwen2_5Wrapper,
+    loader_kwargs=dict(
+        torch_dtype=torch.float16, attn_implementation="flash_attention_2"
+    ),
+    name="ApsaraStackMaaS/EvoQwen2.5-VL-Retriever-3B-v1",
+    model_type=["late-interaction"],
+    languages=["eng-Latn"],
+    revision="aeacaa2775f2758d82721eb1cf2f5daf1a392da9",
+    release_date="2025-11-04",
+    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/ApsaraStackMaaS/EvoQwen2.5-VL-Retriever-3B-v1",
+    similarity_fn_name="MaxSim",
+    use_instructions=True,
+    training_datasets=EVOQWEN_TRAINING_DATA,
+)
+
+evoqwen25_vl_retriever_7b_v1 = ModelMeta(
+    loader=ColQwen2_5Wrapper,
+    loader_kwargs=dict(
+        torch_dtype=torch.float16, attn_implementation="flash_attention_2"
+    ),
+    name="ApsaraStackMaaS/EvoQwen2.5-VL-Retriever-7B-v1",
+    model_type=["late-interaction"],
+    languages=["eng-Latn"],
+    revision="8952ac6ee0e7de2e9211b165921518caf9202110",
+    release_date="2025-11-04",
+    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/illuin-tech/colpali",
+    public_training_data="https://huggingface.co/datasets/vidore/colpali_train_set",
+    framework=["ColPali"],
+    reference="https://huggingface.co/ApsaraStackMaaS/EvoQwen2.5-VL-Retriever-7B-v1",
+    similarity_fn_name="MaxSim",
+    use_instructions=True,
+    training_datasets=EVOQWEN_TRAINING_DATA,
+)

mteb/models/model_implementations/colsmol_models.py

@@ -54,6 +54,7 @@ colsmol_256m = ModelMeta(
         torch_dtype=torch.float16,
     ),
     name="vidore/colSmol-256M",
+    model_type=["late-interaction"],
     languages=["eng-Latn"],
     revision="530094e83a40ca4edcb5c9e5ddfa61a4b5ea0d2f",
     release_date="2025-01-22",
@@ -80,6 +81,7 @@ colsmol_500m = ModelMeta(
         torch_dtype=torch.float16, attn_implementation="flash_attention_2"
     ),
     name="vidore/colSmol-500M",
+    model_type=["late-interaction"],
     languages=["eng-Latn"],
     revision="1aa9325cba7ed2b3b9b97ede4d55026322504902",
     release_date="2025-01-22",

mteb/models/model_implementations/conan_models.py

@@ -190,6 +190,7 @@ class ConanWrapper(AbsEncoder):
 
 Conan_embedding_v2 = ModelMeta(
     name="TencentBAC/Conan-embedding-v2",
+    model_type=["dense"],
     revision="e5c87c63889630bca87486f6a2645ed97c5ddb17",
     release_date="2025-04-10",
     languages=[