mteb 2.7.11__py3-none-any.whl → 2.7.13__py3-none-any.whl

mteb/abstasks/abstask.py

@@ -116,7 +116,7 @@ class AbsTask(ABC):
             logger.warning(msg)
             warnings.warn(msg)
 
-    def dataset_transform(self, num_proc: int = 1):
+    def dataset_transform(self, num_proc: int = 1, **kwargs: Any) -> None:
         """A transform operations applied to the dataset after loading.
 
         This method is useful when the dataset from Huggingface is not in an `mteb` compatible format.
@@ -124,6 +124,7 @@ class AbsTask(ABC):
 
         Args:
             num_proc: Number of processes to use for the transformation.
+            kwargs: Additional keyword arguments passed to the load_dataset function. Keep for forward compatibility.
         """
         pass
 

mteb/models/model_implementations/nomic_models.py

@@ -7,6 +7,7 @@ import torch
 import torch.nn.functional as F
 from packaging.version import Version
 
+from mteb.models import sentence_transformers_loader
 from mteb.models.model_meta import ModelMeta, ScoringFunction
 from mteb.models.sentence_transformer_wrapper import SentenceTransformerEncoderWrapper
 from mteb.types import PromptType
@@ -509,3 +510,42 @@ nomic_embed_text_v2_moe = ModelMeta(
       url={https://arxiv.org/abs/2502.07972},
 }""",
 )
+
+nomic_embed_code = ModelMeta(
+    loader=sentence_transformers_loader,
+    loader_kwargs={
+        "trust_remote_code": True,
+        "model_prompts": model_prompts,
+    },
+    name="nomic-ai/nomic-embed-code",
+    revision="11114029805cee545ef111d5144b623787462a52",
+    release_date="2025-03-24",
+    languages=["eng-Latn"],
+    n_parameters=7_070_619_136,
+    n_embedding_parameters=None,
+    memory_usage_mb=26972.0,
+    max_tokens=32768,
+    embed_dim=3584,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code="https://github.com/gangiswag/cornstack/",
+    public_training_data="https://huggingface.co/collections/nomic-ai/cornstack",
+    framework=["PyTorch", "Sentence Transformers", "safetensors"],
+    reference="https://huggingface.co/nomic-ai/nomic-embed-code",
+    similarity_fn_name=ScoringFunction.COSINE,
+    use_instructions=True,
+    training_datasets={"CoRNStack"},
+    adapted_from=None,
+    superseded_by=None,
+    modalities=["text"],
+    model_type=["dense"],
+    citation="""@misc{suresh2025cornstackhighqualitycontrastivedata,
+      title={CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking},
+      author={Tarun Suresh and Revanth Gangi Reddy and Yifei Xu and Zach Nussbaum and Andriy Mulyar and Brandon Duderstadt and Heng Ji},
+      year={2025},
+      eprint={2412.01007},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2412.01007},
+}""",
+)

mteb/models/model_implementations/nvidia_llama_nemoretriever_colemb.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 
 import torch
-from packaging.version import Version
+from packaging.specifiers import SpecifierSet
 from torch.utils.data import DataLoader
 from transformers import __version__ as transformers_version
 
@@ -31,18 +31,20 @@ class LlamaNemoretrieverColembed(AbsEncoder):
         model_name_or_path: str,
         revision: str,
         trust_remote_code: bool,
+        transformers_version_constraint: str | None = None,
         device_map="cuda",
         torch_dtype=torch.bfloat16,
         attn_implementation="flash_attention_2",
         **kwargs,
     ):
-        required_transformers_version = "4.49.0"
-
-        if Version(transformers_version) != Version(required_transformers_version):
-            raise RuntimeError(
-                f"transformers version {transformers_version} is not match with required "
-                f"install version {required_transformers_version} to run `nvidia/llama-nemoretriever-colembed`"
-            )
+        if transformers_version_constraint is not None:
+            spec = SpecifierSet(transformers_version_constraint)
+            if transformers_version not in spec:
+                raise RuntimeError(
+                    f"Model `{model_name_or_path}` requires transformers{transformers_version_constraint}, "
+                    f"but {transformers_version} is installed. "
+                    f"Run: pip install 'transformers{transformers_version_constraint}'"
+                )
 
         from transformers import AutoModel
 
@@ -150,10 +152,24 @@ TRAINING_DATA = {
     "wiki-ss-nq",
 }
 
+
+TRAINING_DATA_v2 = {
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
+    "docmatix-ir",
+    "VDRMultilingualRetrieval",
+    "VisRAG-Ret-Train-Synthetic-data",
+    "VisRAG-Ret-Train-In-domain-data",
+    "wiki-ss-nq",
+}
+
 llama_nemoretriever_colembed_1b_v1 = ModelMeta(
     loader=LlamaNemoretrieverColembed,
     loader_kwargs=dict(
         trust_remote_code=True,
+        transformers_version_constraint="==4.49.0",
     ),
     name="nvidia/llama-nemoretriever-colembed-1b-v1",
     model_type=["late-interaction"],
@@ -168,7 +184,7 @@ llama_nemoretriever_colembed_1b_v1 = ModelMeta(
     embed_dim=2048,
     license="https://huggingface.co/nvidia/llama-nemoretriever-colembed-1b-v1/blob/main/LICENSE",
     open_weights=True,
-    public_training_code="Proprietary Code",
+    public_training_code=None,
     public_training_data="https://huggingface.co/nvidia/llama-nemoretriever-colembed-1b-v1#training-dataset",
     framework=["PyTorch", "Transformers", "safetensors"],
     reference="https://huggingface.co/nvidia/llama-nemoretriever-colembed-1b-v1",
@@ -182,6 +198,7 @@ llama_nemoretriever_colembed_3b_v1 = ModelMeta(
     loader=LlamaNemoretrieverColembed,
     loader_kwargs=dict(
         trust_remote_code=True,
+        transformers_version_constraint="==4.49.0",
     ),
     name="nvidia/llama-nemoretriever-colembed-3b-v1",
     model_type=["late-interaction"],
@@ -196,7 +213,7 @@ llama_nemoretriever_colembed_3b_v1 = ModelMeta(
     embed_dim=3072,
     license="https://huggingface.co/nvidia/llama-nemoretriever-colembed-1b-v1/blob/main/LICENSE",
     open_weights=True,
-    public_training_code="Proprietary Code",
+    public_training_code=None,
     public_training_data="https://huggingface.co/nvidia/llama-nemoretriever-colembed-1b-v1#training-dataset",
     framework=["PyTorch", "Transformers", "safetensors"],
     reference="https://huggingface.co/nvidia/llama-nemoretriever-colembed-3b-v1",
@@ -205,3 +222,86 @@ llama_nemoretriever_colembed_3b_v1 = ModelMeta(
     training_datasets=TRAINING_DATA,
     citation=LLAMA_NEMORETRIEVER_CITATION,
 )
+
+llama_nemotron_colembed_vl_3b_v2 = ModelMeta(
+    loader=LlamaNemoretrieverColembed,
+    loader_kwargs=dict(
+        trust_remote_code=True,
+        transformers_version_constraint="==4.49.0",
+    ),
+    name="nvidia/llama-nemotron-colembed-vl-3b-v2",
+    model_type=["late-interaction"],
+    languages=["eng-Latn"],
+    revision="75f03c712cb3a252e062295f9a0966e5d95d6156",
+    release_date="2026-01-21",
+    modalities=["image", "text"],
+    n_parameters=4_407_000_000,
+    memory_usage_mb=8403,
+    max_tokens=8192,
+    embed_dim=3072,
+    license="https://huggingface.co/nvidia/llama-nemotron-colembed-vl-3b-v2/blob/main/LICENSE",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data="https://huggingface.co/nvidia/llama-nemotron-colembed-vl-3b-v2#training-dataset",
+    framework=["PyTorch", "Transformers", "safetensors"],
+    reference="https://huggingface.co/nvidia/llama-nemotron-colembed-vl-3b-v2",
+    similarity_fn_name="MaxSim",
+    use_instructions=True,
+    training_datasets=TRAINING_DATA,
+    citation=LLAMA_NEMORETRIEVER_CITATION,
+)
+
+nemotron_colembed_vl_4b_v2 = ModelMeta(
+    loader=LlamaNemoretrieverColembed,
+    loader_kwargs=dict(
+        trust_remote_code=True,
+        transformers_version_constraint="==5.0.0rc0",
+    ),
+    name="nvidia/nemotron-colembed-vl-4b-v2",
+    revision="823b1625c15fe3da73fa094205e538a7a2301a2a",
+    languages=["eng-Latn"],
+    release_date="2026-01-07",
+    modalities=["image", "text"],
+    n_parameters=4_800_000_000,
+    memory_usage_mb=9206,
+    max_tokens=262144,
+    embed_dim=2560,
+    license="https://huggingface.co/nvidia/nemotron-colembed-vl-4b-v2/blob/main/LICENSE",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data="https://huggingface.co/nvidia/nemotron-colembed-vl-4b-v2#training-dataset",
+    framework=["PyTorch", "Transformers"],
+    reference="https://huggingface.co/nvidia/nemotron-colembed-vl-4b-v2",
+    similarity_fn_name="MaxSim",
+    use_instructions=True,
+    training_datasets=TRAINING_DATA_v2,
+    citation=LLAMA_NEMORETRIEVER_CITATION,
+)
+
+
+nemotron_colembed_vl_8b_v2 = ModelMeta(
+    loader=LlamaNemoretrieverColembed,
+    loader_kwargs=dict(
+        trust_remote_code=True,
+        transformers_version_constraint="==5.0.0rc0",
+    ),
+    name="nvidia/nemotron-colembed-vl-8b-v2",
+    revision="6cbe43579dda6237768fc373768ad372cc5cdfec",
+    languages=["eng-Latn"],
+    release_date="2026-01-07",
+    modalities=["image", "text"],
+    n_parameters=8_700_000_000,
+    memory_usage_mb=16722,
+    max_tokens=262144,
+    embed_dim=4096,
+    license="https://huggingface.co/nvidia/nemotron-colembed-vl-8b-v2/blob/main/LICENSE",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data="https://huggingface.co/nvidia/nemotron-colembed-vl-8b-v2#training-dataset",
+    framework=["PyTorch", "Transformers"],
+    reference="https://huggingface.co/nvidia/nemotron-colembed-vl-8b-v2",
+    similarity_fn_name="MaxSim",
+    use_instructions=True,
+    training_datasets=TRAINING_DATA_v2,
+    citation=LLAMA_NEMORETRIEVER_CITATION,
+)

mteb/models/model_implementations/ops_colqwen3_models.py

@@ -0,0 +1,267 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+from tqdm.auto import tqdm
+from transformers import AutoModel, AutoProcessor
+
+from mteb._requires_package import requires_image_dependencies
+from mteb.models.abs_encoder import AbsEncoder
+from mteb.models.model_meta import ModelMeta, ScoringFunction
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
+
+
+class OpsColQwen3Wrapper(AbsEncoder):
+    """Wrapper for OpsColQwen3 model."""
+
+    def __init__(
+        self,
+        model_name: str = "OpenSearch-AI/Ops-Colqwen3-4B",
+        revision: str | None = None,
+        device: str | None = None,
+        attn_implementation: str | None = None,
+        **kwargs,
+    ):
+        requires_image_dependencies()
+        from transformers.utils.import_utils import is_flash_attn_2_available
+
+        if attn_implementation is None:
+            attn_implementation = (
+                "flash_attention_2" if is_flash_attn_2_available() else None
+            )
+
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.mdl = AutoModel.from_pretrained(
+            model_name,
+            device_map=self.device,
+            attn_implementation=attn_implementation,
+            trust_remote_code=True,
+            revision=revision,
+            **kwargs,
+        )
+        self.mdl.eval()
+
+        self.processor = AutoProcessor.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+        )
+
+    def encode(
+        self,
+        inputs: DataLoader[BatchedInput],
+        *,
+        task_metadata: TaskMetadata,
+        hf_split: str,
+        hf_subset: str,
+        prompt_type: PromptType | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        text_embeddings = None
+        image_embeddings = None
+
+        if "text" in inputs.dataset.features:
+            text_embeddings = self.get_text_embeddings(inputs, **kwargs)
+        if "image" in inputs.dataset.features:
+            image_embeddings = self.get_image_embeddings(inputs, **kwargs)
+
+        if text_embeddings is not None and image_embeddings is not None:
+            if len(text_embeddings) != len(image_embeddings):
+                raise ValueError(
+                    "The number of texts and images must have the same length"
+                )
+            fused_embeddings = text_embeddings + image_embeddings
+            return fused_embeddings
+        elif text_embeddings is not None:
+            return text_embeddings
+        elif image_embeddings is not None:
+            return image_embeddings
+        raise ValueError("No text or image inputs found")
+
+    def encode_input(self, inputs):
+        return self.mdl(**inputs)
+
+    def get_image_embeddings(
+        self,
+        images: DataLoader,
+        batch_size: int = 32,
+        **kwargs,
+    ) -> torch.Tensor:
+        import torchvision.transforms.functional as F
+        from PIL import Image
+
+        all_embeds = []
+
+        with torch.no_grad():
+            for batch in tqdm(images, desc="Encoding images"):
+                # batch may be list of tensors or PIL
+                imgs = [
+                    F.to_pil_image(b.to(self.device))
+                    if not isinstance(b, Image.Image)
+                    else b
+                    for b in batch["image"]
+                ]
+                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: DataLoader,
+        batch_size: int = 32,
+        **kwargs,
+    ) -> torch.Tensor:
+        all_embeds = []
+
+        with torch.no_grad():
+            for batch in tqdm(texts, desc="Encoding texts"):
+                batch_texts = batch["text"]
+                inputs = self.processor.process_queries(batch_texts)
+                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 similarity(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return self.processor.score_multi_vector(a, b, device=self.device)
+
+
+OPS_COLQWEN3_TRAINING_DATA = {
+    "VDRMultilingualRetrieval",
+    # from https://huggingface.co/datasets/vidore/colpali_train_set
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
+    "docmatix-ir",
+    "HotpotQA",
+    "FEVER",
+    "NQ",
+    "MIRACLRetrieval",
+    "WebInstructSub",  # MathStackExchange and ScienceStackExchange only
+    "MrTyDi",
+}
+
+multilingual_langs = [
+    "afr-Latn",
+    "ara-Arab",
+    "aze-Latn",
+    "bel-Cyrl",
+    "bul-Cyrl",
+    "ben-Beng",
+    "cat-Latn",
+    "ceb-Latn",
+    "ces-Latn",
+    "cym-Latn",
+    "dan-Latn",
+    "deu-Latn",
+    "ell-Grek",
+    "eng-Latn",
+    "spa-Latn",
+    "est-Latn",
+    "eus-Latn",
+    "fas-Arab",
+    "fin-Latn",
+    "fra-Latn",
+    "glg-Latn",
+    "guj-Gujr",
+    "heb-Hebr",
+    "hin-Deva",
+    "hrv-Latn",
+    "hat-Latn",
+    "hun-Latn",
+    "hye-Armn",
+    "ind-Latn",
+    "isl-Latn",
+    "ita-Latn",
+    "jpn-Jpan",
+    "jav-Latn",
+    "kat-Geor",
+    "kaz-Cyrl",
+    "khm-Khmr",
+    "kan-Knda",
+    "kor-Hang",
+    "kir-Cyrl",
+    "lao-Laoo",
+    "lit-Latn",
+    "lav-Latn",
+    "mkd-Cyrl",
+    "mal-Mlym",
+    "mon-Cyrl",
+    "mar-Deva",
+    "msa-Latn",
+    "mya-Mymr",
+    "nep-Deva",
+    "nld-Latn",
+    "nor-Latn",
+    "nob-Latn",
+    "nno-Latn",
+    "pan-Guru",
+    "pol-Latn",
+    "por-Latn",
+    "que-Latn",
+    "ron-Latn",
+    "rus-Cyrl",
+    "sin-Sinh",
+    "slk-Latn",
+    "slv-Latn",
+    "swa-Latn",
+    "tam-Taml",
+    "tel-Telu",
+    "tha-Thai",
+    "tgl-Latn",
+    "tur-Latn",
+    "ukr-Cyrl",
+    "urd-Arab",
+    "vie-Latn",
+    "yor-Latn",
+    "zho-Hans",
+]
+
+OPS_COLQWEN3_CITATION = """
+@misc{ops_colqwen3_4b,
+  author       = {OpenSearch-AI},
+  title        = {Ops-ColQwen3: State-of-the-Art Multimodal Embedding Model for Visual Document Retrieval},
+  year         = {2026},
+  url          = {https://huggingface.co/OpenSearch-AI/Ops-ColQwen3-4B},
+}"""
+
+ops_colqwen3_4b = ModelMeta(
+    loader=OpsColQwen3Wrapper,
+    name="OpenSearch-AI/Ops-Colqwen3-4B",
+    loader_kwargs=dict(dtype=torch.float16, trust_remote_code=True),
+    languages=multilingual_langs,
+    revision="4894b7d451ff33981650acc693bb482dbef302d3",
+    release_date="2026-01-24",
+    modalities=["image", "text"],
+    n_parameters=4_800_000_000,
+    memory_usage_mb=9206,
+    max_tokens=32768,
+    embed_dim=2560,
+    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/OpenSearch-AI/Ops-Colqwen3-4B",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    training_datasets=OPS_COLQWEN3_TRAINING_DATA,
+    citation=OPS_COLQWEN3_CITATION,
+)