mteb 2.7.20__py3-none-any.whl → 2.7.21__py3-none-any.whl

mteb/models/model_implementations/{nvidia_llama_nemoretriever_colemb.py → nvidia_nemotron_colembed_vl.py}

@@ -7,6 +7,7 @@ from packaging.specifiers import SpecifierSet
 from torch.utils.data import DataLoader
 from transformers import __version__ as transformers_version
 
+from mteb._requires_package import requires_package
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta
 
@@ -24,27 +25,51 @@ LLAMA_NEMORETRIEVER_CITATION = """@misc{xu2025llamanemoretrievercolembedtopperfo
       url={https://arxiv.org/abs/2507.05513}
 }"""
 
+# Transformers version constraints per extra.
+# Keep in sync with pyproject.toml [project.optional-dependencies]
+#
+# Note: The extra name reflects the transformers version requirement, not the model version.
+# For example, llama-nemotron-colembed-vl-3b-v2 uses "llama-nemotron-colembed-vl" because it
+# requires transformers==4.49.0, even though it's a "v2" model by name.
+_TRANSFORMERS_CONSTRAINTS: dict[str, str] = {
+    "llama-nemotron-colembed-vl": "==4.49.0",  # llama-nemoretriever-colembed-*
+    "nemotron-colembed-vl-v2": "==5.0.0rc0",  # nemotron-colembed-vl-4b-v2, nemotron-colembed-vl-8b-v2
+}
+
+
+class NemotronColEmbedVL(AbsEncoder):
+    """Encoder for the NemotronColEmbedVL family of models."""
 
-class LlamaNemoretrieverColembed(AbsEncoder):
     def __init__(
         self,
         model_name_or_path: str,
         revision: str,
         trust_remote_code: bool,
-        transformers_version_constraint: str | None = None,
+        extra_name: str = "llama-nemotron-colembed-vl",
         device_map="cuda",
         torch_dtype=torch.bfloat16,
         attn_implementation="flash_attention_2",
         **kwargs,
     ):
-        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}'"
-                )
+        install_hint = f"pip install 'mteb[{extra_name}]'"
+
+        # Check transformers version
+        constraint = _TRANSFORMERS_CONSTRAINTS.get(extra_name)
+        if constraint is None:
+            raise ValueError(
+                f"Unknown extra_name '{extra_name}'. "
+                f"Must be one of: {list(_TRANSFORMERS_CONSTRAINTS.keys())}"
+            )
+        if transformers_version not in SpecifierSet(constraint):
+            raise RuntimeError(
+                f"Model `{model_name_or_path}` requires transformers{constraint}, "
+                f"but {transformers_version} is installed. "
+                f"Run: {install_hint}"
+            )
+
+        # Check required packages
+        for package in ("torchvision", "accelerate", "flash_attn"):
+            requires_package(self, package, model_name_or_path, install_hint)
 
         from transformers import AutoModel
 
@@ -166,10 +191,10 @@ TRAINING_DATA_v2 = {
 }
 
 llama_nemoretriever_colembed_1b_v1 = ModelMeta(
-    loader=LlamaNemoretrieverColembed,
+    loader=NemotronColEmbedVL,
     loader_kwargs=dict(
+        extra_name="llama-nemotron-colembed-vl",
         trust_remote_code=True,
-        transformers_version_constraint="==4.49.0",
     ),
     name="nvidia/llama-nemoretriever-colembed-1b-v1",
     model_type=["late-interaction"],
@@ -195,10 +220,10 @@ llama_nemoretriever_colembed_1b_v1 = ModelMeta(
 )
 
 llama_nemoretriever_colembed_3b_v1 = ModelMeta(
-    loader=LlamaNemoretrieverColembed,
+    loader=NemotronColEmbedVL,
     loader_kwargs=dict(
+        extra_name="llama-nemotron-colembed-vl",
         trust_remote_code=True,
-        transformers_version_constraint="==4.49.0",
     ),
     name="nvidia/llama-nemoretriever-colembed-3b-v1",
     model_type=["late-interaction"],
@@ -224,10 +249,10 @@ llama_nemoretriever_colembed_3b_v1 = ModelMeta(
 )
 
 llama_nemotron_colembed_vl_3b_v2 = ModelMeta(
-    loader=LlamaNemoretrieverColembed,
+    loader=NemotronColEmbedVL,
     loader_kwargs=dict(
+        extra_name="llama-nemotron-colembed-vl",
         trust_remote_code=True,
-        transformers_version_constraint="==4.49.0",
     ),
     name="nvidia/llama-nemotron-colembed-vl-3b-v2",
     model_type=["late-interaction"],
@@ -251,11 +276,12 @@ llama_nemotron_colembed_vl_3b_v2 = ModelMeta(
     citation=LLAMA_NEMORETRIEVER_CITATION,
 )
 
+
 nemotron_colembed_vl_4b_v2 = ModelMeta(
-    loader=LlamaNemoretrieverColembed,
+    loader=NemotronColEmbedVL,
     loader_kwargs=dict(
+        extra_name="nemotron-colembed-vl-v2",
         trust_remote_code=True,
-        transformers_version_constraint="==5.0.0rc0",
     ),
     name="nvidia/nemotron-colembed-vl-4b-v2",
     revision="823b1625c15fe3da73fa094205e538a7a2301a2a",
@@ -280,10 +306,10 @@ nemotron_colembed_vl_4b_v2 = ModelMeta(
 
 
 nemotron_colembed_vl_8b_v2 = ModelMeta(
-    loader=LlamaNemoretrieverColembed,
+    loader=NemotronColEmbedVL,
     loader_kwargs=dict(
+        extra_name="nemotron-colembed-vl-v2",
         trust_remote_code=True,
-        transformers_version_constraint="==5.0.0rc0",
     ),
     name="nvidia/nemotron-colembed-vl-8b-v2",
     revision="6cbe43579dda6237768fc373768ad372cc5cdfec",

{mteb-2.7.20.dist-info → mteb-2.7.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mteb
-Version: 2.7.20
+Version: 2.7.21
 Summary: Massive Text Embedding Benchmark
 Author-email: MTEB Contributors <niklas@huggingface.co>, Kenneth Enevoldsen <kenneth.enevoldsen@cas.au.dk>, Nouamane Tazi <nouamane@huggingface.co>, Nils Reimers <info@nils-reimers.de>
 Maintainer-email: Kenneth Enevoldsen <kenneth.enevoldsen@cas.au.dk>, Roman Solomatin <risolomatin@gmail.com>, Isaac Chung <chungisaac1217@gmail.com>
@@ -104,6 +104,16 @@ Requires-Dist: tencentcloud-sdk-python-common>=3.0.1454; extra == "youtu"
 Requires-Dist: tencentcloud-sdk-python-lkeap>=3.0.1451; extra == "youtu"
 Provides-Extra: llama-embed-nemotron
 Requires-Dist: transformers==4.51.0; extra == "llama-embed-nemotron"
+Provides-Extra: llama-nemotron-colembed-vl
+Requires-Dist: transformers[torch]==4.49.0; extra == "llama-nemotron-colembed-vl"
+Requires-Dist: torchvision>=0.22.0; extra == "llama-nemotron-colembed-vl"
+Requires-Dist: flash-attn>=2.6.3; extra == "llama-nemotron-colembed-vl"
+Requires-Dist: accelerate; extra == "llama-nemotron-colembed-vl"
+Provides-Extra: nemotron-colembed-vl-v2
+Requires-Dist: transformers[torch]==5.0.0rc0; extra == "nemotron-colembed-vl-v2"
+Requires-Dist: torchvision>=0.22.0; extra == "nemotron-colembed-vl-v2"
+Requires-Dist: flash-attn>=2.6.3; extra == "nemotron-colembed-vl-v2"
+Requires-Dist: accelerate; extra == "nemotron-colembed-vl-v2"
 Provides-Extra: faiss-cpu
 Requires-Dist: faiss-cpu>=1.12.0; extra == "faiss-cpu"
 Provides-Extra: eager-embed

{mteb-2.7.20.dist-info → mteb-2.7.21.dist-info}/RECORD

@@ -1568,8 +1568,8 @@ mteb/models/model_implementations/nbailab.py,sha256=iv2xdqVM5HoTAlBR6e_UdzJu6rSP
 mteb/models/model_implementations/no_instruct_sentence_models.py,sha256=DTb-eHZYSY6lGJkkdkC0tZ_n0GHLQwVlUehVg59T5N4,4198
 mteb/models/model_implementations/nomic_models.py,sha256=BO6XQbX4PFa5By0opAYkxz95CcHmjxbG5DYcklxJ1l8,16986
 mteb/models/model_implementations/nomic_models_vision.py,sha256=AzTCWbXBonUAVub0TTxWCsBtg4WYex3vPiLlz3ULdHc,6916
-mteb/models/model_implementations/nvidia_llama_nemoretriever_colemb.py,sha256=-5_kd9jeDcgVv9gdwWuvX_-bNQdhAxInf1Mqo8_BdS8,10653
 mteb/models/model_implementations/nvidia_models.py,sha256=r-AW1dVQbteWjexjvZgFEt_90OHNRYer_5GLuqSXRS0,26924
+mteb/models/model_implementations/nvidia_nemotron_colembed_vl.py,sha256=BMb3NG8lYpabf6FDa5zheUajbMc2E0_dyvgO6eUSzAc,11654
 mteb/models/model_implementations/octen_models.py,sha256=5z-t2O-iIFiOOLdZ_AK9f7GrVRg-9_vx3JNAG9dJNPE,8562
 mteb/models/model_implementations/openai_models.py,sha256=fE8SfSAcl20GccR8D8s-7MR9w_kO6LlN5Pm80Iwx82c,9777
 mteb/models/model_implementations/openclip_models.py,sha256=z2gQum16O0QhJPyxqKor3oO-_uWfnep6wSXqOFQQ2Q8,11969
@@ -2647,9 +2647,9 @@ mteb/types/_metadata.py,sha256=NN-W0S6a5TDV7UkpRx1pyWtGF4TyyCyoPUfHOwdeci8,2290
 mteb/types/_result.py,sha256=UKNokV9pu3G74MGebocU512aU_fFU9I9nPKnrG9Q0iE,1035
 mteb/types/_string_validators.py,sha256=PY-dYq4E8O50VS3bLYdldPWp400fl_WzUjfVSkNWe8U,523
 mteb/types/statistics.py,sha256=gElgSShKBXpfcqaZHhU_d2UHln1CyzUj8FN8KFun_UA,4087
-mteb-2.7.20.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mteb-2.7.20.dist-info/METADATA,sha256=TlZtKy_JecJva-vrjFYjqLen3vuDp3zWw-RfSvuwAFI,14348
-mteb-2.7.20.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-mteb-2.7.20.dist-info/entry_points.txt,sha256=8IJoEJFKoDHmVnNev-qJ9pp4Ln7_1-ma9QsXnzVCzGU,39
-mteb-2.7.20.dist-info/top_level.txt,sha256=OLVIjcQAlWBz0bdmutKlWHLF42FF0hp4uVAg3ZyiG4U,5
-mteb-2.7.20.dist-info/RECORD,,
+mteb-2.7.21.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mteb-2.7.21.dist-info/METADATA,sha256=wsI05ECfej7Ql49z9pJnJo0Ompp8z3jaM3y-jaI6kz8,15007
+mteb-2.7.21.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+mteb-2.7.21.dist-info/entry_points.txt,sha256=8IJoEJFKoDHmVnNev-qJ9pp4Ln7_1-ma9QsXnzVCzGU,39
+mteb-2.7.21.dist-info/top_level.txt,sha256=OLVIjcQAlWBz0bdmutKlWHLF42FF0hp4uVAg3ZyiG4U,5
+mteb-2.7.21.dist-info/RECORD,,