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

mteb/models/instruct_wrapper.py

@@ -1,23 +1,32 @@
+from __future__ import annotations
+
 import logging
-from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.utils.data import DataLoader
 
 from mteb._requires_package import requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
 
 from .abs_encoder import AbsEncoder
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
+
+
 logger = logging.getLogger(__name__)
 
 
 def instruct_wrapper(
     model_name_or_path: str,
     mode: str,
-    instruction_template: str | Callable[[str], str] | None = None,
+    instruction_template: str | Callable[[str, PromptType | None], str] | None = None,
+    device: str | None = None,
     **kwargs,
 ):
     """Instruct wrapper for models. Uses GritLM to pass instructions to the model.
@@ -28,6 +37,7 @@ def instruct_wrapper(
         model_name_or_path: Model name or path.
         mode: Mode of the model. Either 'query' or 'passage'.
         instruction_template: Instruction template. Should contain the string '{instruction}'.
+        device: Device used to load the model.
         **kwargs: Additional arguments to pass to the model.
     """
     requires_package(
@@ -40,7 +50,10 @@ def instruct_wrapper(
             self,
             model_name_or_path: str,
             mode: str,
-            instruction_template: str | Callable[[str, PromptType], str] | None = None,
+            device: str | None = None,
+            instruction_template: str
+            | Callable[[str, PromptType | None], str]
+            | None = None,
             **kwargs,
         ):
             if (
@@ -61,7 +74,12 @@ def instruct_wrapper(
                 )
 
             self.instruction_template = instruction_template
-            super().__init__(model_name_or_path=model_name_or_path, mode=mode, **kwargs)
+            super().__init__(
+                model_name_or_path=model_name_or_path,
+                mode=mode,
+                device=device,
+                **kwargs,
+            )
 
         def encode(
             self,
@@ -82,15 +100,20 @@ def instruct_wrapper(
             logger.info(
                 f"Using instruction: '{instruction}' for task: '{task_metadata.name}'"
             )
-            embeddings = super().encode(
-                _inputs, instruction=instruction, *args, **kwargs
+            embeddings = super().encode(  # type: ignore[safe-super,call-arg]
+                _inputs,  # type: ignore[arg-type]
+                instruction=instruction,
+                *args,
+                **kwargs,
             )
             if isinstance(embeddings, torch.Tensor):
                 # sometimes in kwargs can be return_tensors=True
                 embeddings = embeddings.cpu().detach().float().numpy()
             return embeddings
 
-    return InstructGritLMModel(model_name_or_path, mode, instruction_template, **kwargs)
+    return InstructGritLMModel(
+        model_name_or_path, mode, instruction_template=instruction_template, **kwargs
+    )
 
 
 class InstructSentenceTransformerModel(AbsEncoder):
@@ -100,6 +123,7 @@ class InstructSentenceTransformerModel(AbsEncoder):
         self,
         model_name: str,
         revision: str,
+        device: str | None = None,
         instruction_template: str
         | Callable[[str, PromptType | None], str]
         | None = None,
@@ -117,6 +141,7 @@ class InstructSentenceTransformerModel(AbsEncoder):
         Arguments:
             model_name: Model name of the sentence transformers model.
             revision: Revision of the sentence transformers model.
+            device: Device used to load the model.
             instruction_template: Model template. Should contain the string '{instruction}'.
             max_seq_length: Maximum sequence length. If None, the maximum sequence length will be read from the model config.
             apply_instruction_to_passages: Whether to apply the instruction template to the passages.
@@ -141,7 +166,7 @@ class InstructSentenceTransformerModel(AbsEncoder):
             )
 
         self.instruction_template = instruction_template
-        tokenizer_params = {}
+        tokenizer_params: dict[str, Any] = {}
         if add_eos_token:
             tokenizer_params["add_eos_token"] = add_eos_token
         if max_seq_length is not None:
@@ -153,7 +178,9 @@ class InstructSentenceTransformerModel(AbsEncoder):
         kwargs.setdefault("tokenizer_kwargs", {}).update(tokenizer_params)
 
         self.model_name = model_name
-        self.model = SentenceTransformer(model_name, revision=revision, **kwargs)
+        self.model = SentenceTransformer(
+            model_name, revision=revision, device=device, **kwargs
+        )
         if max_seq_length:
             # https://github.com/huggingface/sentence-transformers/issues/3575
             self.model.max_seq_length = max_seq_length
@@ -193,6 +220,7 @@ class InstructSentenceTransformerModel(AbsEncoder):
             The encoded input in a numpy array or torch tensor of the shape (Number of sentences) x (Embedding dimension).
         """
         sentences = [text for batch in inputs for text in batch["text"]]
+        instruction: str | None
         instruction = self.get_task_instruction(task_metadata, prompt_type)
 
         # to passage prompts won't be applied to passages

mteb/models/model_implementations/align_models.py

@@ -1,13 +1,18 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
-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
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
 
 
 class ALIGNModel(AbsEncoder):
@@ -111,6 +116,7 @@ align_base = ModelMeta(
     release_date="2023-02-24",
     modalities=["image", "text"],
     n_parameters=176_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=671,
     max_tokens=64,
     embed_dim=768,
@@ -118,7 +124,7 @@ align_base = ModelMeta(
     open_weights=True,
     public_training_code="https://github.com/kakaobrain/coyo-align",
     public_training_data=True,
-    framework=["PyTorch"],
+    framework=["PyTorch", "Transformers"],
     reference="https://huggingface.co/kakaobrain/align-base",
     similarity_fn_name=ScoringFunction.COSINE,
     use_instructions=False,

mteb/models/model_implementations/amazon_models.py

@@ -8,6 +8,7 @@ amazon_titan_text_embeddings_v2 = ModelMeta(
     release_date="2024-04-30",
     languages=["eng-Latn"],
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     max_tokens=None,
     embed_dim=None,

mteb/models/model_implementations/andersborges.py

@@ -4,7 +4,7 @@ from mteb.models.model_implementations.model2vec_models import Model2VecModel
 from mteb.models.model_meta import ModelMeta, ScoringFunction
 
 model2vecdk = ModelMeta(
-    loader=Model2VecModel,  # type: ignore
+    loader=Model2VecModel,
     name="andersborges/model2vecdk",
     model_type=["dense"],
     languages=["dan-Latn"],
@@ -12,12 +12,13 @@ model2vecdk = ModelMeta(
     revision="cb576c78dcc1b729e4612645f61db59929d69e61",
     release_date="2025-11-21",
     n_parameters=48042496,
+    n_embedding_parameters=None,
     memory_usage_mb=183,
     max_tokens=np.inf,
     embed_dim=256,
     license="mit",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["NumPy", "Sentence Transformers"],
+    framework=["NumPy", "Sentence Transformers", "safetensors"],
     reference="https://huggingface.co/andersborges/model2vecdk",
     use_instructions=False,
     adapted_from="https://huggingface.co/jealk/TTC-L2V-supervised-2",
@@ -35,7 +36,7 @@ model2vecdk = ModelMeta(
 
 
 model2vecdk_stem = ModelMeta(
-    loader=Model2VecModel,  # type: ignore
+    loader=Model2VecModel,
     name="andersborges/model2vecdk-stem",
     model_type=["dense"],
     languages=["dan-Latn"],
@@ -43,12 +44,13 @@ model2vecdk_stem = ModelMeta(
     revision="cb576c78dcc1b729e4612645f61db59929d69e61",
     release_date="2025-11-21",
     n_parameters=48578560,
+    n_embedding_parameters=None,
     memory_usage_mb=185,
     max_tokens=np.inf,
     embed_dim=256,
     license="mit",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["NumPy", "Sentence Transformers"],
+    framework=["NumPy", "Sentence Transformers", "safetensors"],
     reference="https://huggingface.co/andersborges/model2vecdk",
     use_instructions=False,
     adapted_from="https://huggingface.co/jealk/TTC-L2V-supervised-2",

mteb/models/model_implementations/ara_models.py

@@ -10,13 +10,14 @@ arabic_triplet_matryoshka = ModelMeta(
     revision="ed357f222f0b6ea6670d2c9b5a1cb93950d34200",
     release_date="2024-07-28",
     n_parameters=135_000_000,
+    n_embedding_parameters=49_152_000,
     memory_usage_mb=516,
     embed_dim=768,
     license="apache-2.0",
     max_tokens=768,
     reference="https://huggingface.co/Omartificial-Intelligence-Space/Arabic-Triplet-Matryoshka-V2",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors", "Transformers"],
     use_instructions=False,
     public_training_code=None,
     adapted_from="aubmindlab/bert-base-arabertv02",

mteb/models/model_implementations/arctic_models.py

@@ -145,8 +145,9 @@ arctic_embed_xs = ModelMeta(
     release_date="2024-07-08",  # initial commit of hf model.
     languages=["eng-Latn"],
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     n_parameters=22_600_000,
+    n_embedding_parameters=11_720_448,
     memory_usage_mb=86,
     max_tokens=512,
     embed_dim=384,
@@ -171,8 +172,9 @@ arctic_embed_s = ModelMeta(
     release_date="2024-04-12",  # initial commit of hf model.
     languages=["eng-Latn"],
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     n_parameters=32_200_000,
+    n_embedding_parameters=11_720_448,
     memory_usage_mb=127,
     max_tokens=512,
     embed_dim=384,
@@ -197,8 +199,9 @@ arctic_embed_m = ModelMeta(
     release_date="2024-04-12",  # initial commit of hf model.
     languages=["eng-Latn"],
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     n_parameters=109_000_000,
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=415,
     max_tokens=512,
     embed_dim=768,
@@ -223,8 +226,9 @@ arctic_embed_m_long = ModelMeta(
     release_date="2024-04-12",  # initial commit of hf model.
     languages=["eng-Latn"],
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     n_parameters=137_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=522,
     max_tokens=2048,
     embed_dim=768,
@@ -248,8 +252,9 @@ arctic_embed_l = ModelMeta(
     release_date="2024-04-12",  # initial commit of hf model.
     languages=["eng-Latn"],
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     n_parameters=335_000_000,
+    n_embedding_parameters=31_254_528,
     memory_usage_mb=1274,
     max_tokens=512,
     embed_dim=1024,
@@ -278,8 +283,9 @@ arctic_embed_m_v1_5 = ModelMeta(
     release_date="2024-07-08",  # initial commit of hf model.
     languages=["eng-Latn"],
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors", "GGUF"],
     n_parameters=109_000_000,
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=415,
     max_tokens=512,
     embed_dim=768,
@@ -304,8 +310,9 @@ arctic_embed_m_v2_0 = ModelMeta(
     release_date="2024-12-04",  # initial commit of hf model.
     languages=LANGUAGES_V2_0,
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     n_parameters=305_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=1165,
     max_tokens=8192,
     embed_dim=768,
@@ -329,8 +336,9 @@ arctic_embed_l_v2_0 = ModelMeta(
     release_date="2024-12-04",  # initial commit of hf model.
     languages=LANGUAGES_V2_0,
     open_weights=True,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     n_parameters=568_000_000,
+    n_embedding_parameters=256_002_048,
     memory_usage_mb=2166,
     max_tokens=8192,
     embed_dim=1024,

mteb/models/model_implementations/b1ade_models.py

@@ -16,13 +16,14 @@ b1ade_embed = ModelMeta(
     open_weights=True,
     release_date="2025-03-10",
     n_parameters=335_000_000,
+    n_embedding_parameters=31_254_528,
     memory_usage_mb=1278,
     embed_dim=1024,
     license="mit",
     max_tokens=4096,
     reference="https://huggingface.co/w601sxs/b1ade-embed",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "Transformers", "safetensors"],
     use_instructions=False,
     public_training_code=None,
     public_training_data=None,

mteb/models/model_implementations/bedrock_models.py

@@ -1,20 +1,30 @@
+from __future__ import annotations
+
 import json
 import logging
 import re
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import requires_package
-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 .cohere_models import model_prompts as cohere_model_prompts
-from .cohere_models import supported_languages as cohere_supported_languages
+from .cohere_models import (
+    model_prompts as cohere_model_prompts,
+)
+from .cohere_models import (
+    supported_languages as cohere_supported_languages,
+)
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
+
 
 logger = logging.getLogger(__name__)
 
@@ -169,6 +179,7 @@ amazon_titan_embed_text_v1 = ModelMeta(
     embed_dim=1536,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     public_training_code=None,
     public_training_data=None,  # assumed
@@ -196,6 +207,7 @@ amazon_titan_embed_text_v2 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     public_training_code=None,
     public_training_data=None,  # assumed
@@ -225,6 +237,7 @@ cohere_embed_english_v3 = ModelMeta(
     revision="1",
     release_date="2023-11-02",
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     public_training_code=None,
     public_training_data=None,  # assumed
@@ -253,6 +266,7 @@ cohere_embed_multilingual_v3 = ModelMeta(
     revision="1",
     release_date="2023-11-02",
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     public_training_code=None,
     public_training_data=None,  # assumed