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

mteb/models/cache_wrappers/cache_backends/numpy_cache.py

@@ -1,11 +1,11 @@
 import json
 import logging
+import warnings
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 
-from mteb.types import BatchedInput
-
 from ._hash_utils import _hash_item
 
 logger = logging.getLogger(__name__)
@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 class NumpyCache:
     """Generic vector cache for both text and images."""
 
-    def __init__(self, directory: str | Path, initial_vectors: int = 100000):
+    def __init__(self, directory: str | Path, initial_vectors: int = 100_000):
         self.directory = Path(directory)
         self.directory.mkdir(parents=True, exist_ok=True)
         self.vectors_file = self.directory / "vectors.npy"
@@ -27,7 +27,7 @@ class NumpyCache:
         logger.info(f"Initialized VectorCacheMap in directory: {self.directory}")
         self._initialize_vectors_file()
 
-    def add(self, item: list[BatchedInput], vectors: np.ndarray) -> None:
+    def add(self, items: list[dict[str, Any]], vectors: np.ndarray) -> None:
         """Add a vector to the cache."""
         try:
             if self.vector_dim is None:
@@ -38,12 +38,17 @@ class NumpyCache:
                 self._save_dimension()
                 logger.info(f"Initialized vector dimension to {self.vector_dim}")
 
-            for item, vec in zip(item, vectors):
+            if self.vectors is None:
+                raise RuntimeError(
+                    "Vectors file not initialized. Call _initialize_vectors_file() first."
+                )
+
+            for item, vec in zip(items, vectors):
                 item_hash = _hash_item(item)
                 if item_hash in self.hash_to_index:
-                    logger.warning(
-                        "Hash collision or duplicate item. Overwriting existing vector."
-                    )
+                    msg = f"Hash collision or duplicate item for hash {item_hash}. Overwriting existing vector."
+                    logger.warning(msg)
+                    warnings.warn(msg)
                     index = self.hash_to_index[item_hash]
                 else:
                     index = len(self.hash_to_index)
@@ -74,18 +79,26 @@ class NumpyCache:
                 shape=(self.initial_vectors, self.vector_dim),
             )
         else:
-            self.vectors = np.memmap(self.vectors_file, dtype="float32", mode="r+")
-            self.vectors = self.vectors.reshape(-1, self.vector_dim)
+            self.vectors = np.memmap(
+                self.vectors_file,
+                dtype="float32",
+                mode="r+",
+                shape=(-1, self.vector_dim),
+            )
         logger.info(f"Vectors file initialized with shape: {self.vectors.shape}")
 
     def _double_vectors_file(self) -> None:
+        if self.vectors is None or self.vector_dim is None:
+            raise RuntimeError(
+                "Vectors file not initialized. Call _initialize_vectors_file() first."
+            )
         current_size = len(self.vectors)
         new_size = current_size * 2
         logger.info(f"Doubling vectors file from {current_size} to {new_size} vectors")
         self.vectors.flush()
         new_vectors = np.memmap(
-            self.vectors_file,
-            dtype="float32",
+            str(self.vectors_file),
+            dtype=np.float32,
             mode="r+",
             shape=(new_size, self.vector_dim),
         )
@@ -107,9 +120,9 @@ class NumpyCache:
                 f"Loaded vector dimension {self.vector_dim} from {self.dimension_file}"
             )
         else:
-            logger.warning(
-                "Dimension file not found. Vector dimension remains uninitialized."
-            )
+            msg = "Dimension file not found. Vector dimension remains uninitialized."
+            logger.warning(msg)
+            warnings.warn(msg)
 
     def save(self) -> None:
         """Persist VectorCacheMap to disk."""
@@ -146,25 +159,30 @@ class NumpyCache:
 
                 if self.vector_dim is not None:
                     self.vectors = np.memmap(
-                        self.vectors_file, dtype="float32", mode="r+"
+                        self.vectors_file,
+                        dtype="float32",
+                        mode="r+",
+                        shape=(-1, self.vector_dim),
                     )
-                    self.vectors = self.vectors.reshape(-1, self.vector_dim)
                     logger.info(f"Loaded vectors file with shape: {self.vectors.shape}")
                 else:
-                    logger.warning(
-                        "Vector dimension not set. Unable to load vectors file."
-                    )
+                    msg = "Vector dimension not set. Unable to load vectors file."
+                    logger.warning(msg)
+                    warnings.warn(msg)
                 logger.info(f"Loaded VectorCacheMap from {self.directory}")
             else:
-                logger.warning(
-                    "No existing files found. Initialized empty VectorCacheMap."
-                )
+                msg = "No existing files found. Initialized empty VectorCacheMap."
+                logger.warning(msg)
+                warnings.warn(msg)
         except Exception as e:
             logger.error(f"Error loading VectorCacheMap: {str(e)}")
             raise
 
-    def get_vector(self, item: BatchedInput) -> np.ndarray | None:
+    def get_vector(self, item: dict[str, Any]) -> np.ndarray | None:
         """Retrieve vector from index by hash."""
+        if self.vectors is None:
+            return None
+
         try:
             item_hash = _hash_item(item)
             if item_hash not in self.hash_to_index:
@@ -176,7 +194,7 @@ class NumpyCache:
             logger.error(f"Error retrieving vector for item: {str(e)}")
             raise
 
-    def __contains__(self, item: BatchedInput) -> bool:
+    def __contains__(self, item: dict[str, Any]) -> bool:
         return _hash_item(item) in self.hash_to_index
 
     def __del__(self):

mteb/models/cache_wrappers/cache_wrapper.py

@@ -90,9 +90,9 @@ class CachedEmbeddingWrapper:
         try:
             cache = self._get_or_create_cache(task_name)
 
-            uncached_items: list[BatchedInput] = []
+            uncached_items: list[dict[str, Any]] = []
             uncached_indices: list[int] = []
-            all_items = inputs.dataset
+            all_items: Dataset = inputs.dataset
             cached_vectors: dict[int, np.ndarray] = {}
 
             for i, item in enumerate(all_items):

mteb/models/get_model_meta.py

@@ -21,6 +21,7 @@ def get_model_metas(
     n_parameters_range: tuple[int | None, int | None] = (None, None),
     use_instructions: bool | None = None,
     zero_shot_on: list[AbsTask] | None = None,
+    model_types: Iterable[str] | None = None,
 ) -> list[ModelMeta]:
     """Load all models' metadata that fit the specified criteria.
 
@@ -33,6 +34,7 @@ def get_model_metas(
             If (None, None), this filter is ignored.
         use_instructions: Whether to filter by models that use instructions. If None, all models are included.
         zero_shot_on: A list of tasks on which the model is zero-shot. If None this filter is ignored.
+        model_types: A list of model types to filter by. If None, all model types are included.
 
     Returns:
         A list of model metadata objects that fit the specified criteria.
@@ -41,6 +43,7 @@ def get_model_metas(
     model_names = set(model_names) if model_names is not None else None
     languages = set(languages) if languages is not None else None
     frameworks = set(frameworks) if frameworks is not None else None
+    model_types_set = set(model_types) if model_types is not None else None
     for model_meta in MODEL_REGISTRY.values():
         if (model_names is not None) and (model_meta.name not in model_names):
             continue
@@ -57,6 +60,10 @@ def get_model_metas(
             model_meta.use_instructions != use_instructions
         ):
             continue
+        if model_types_set is not None and not model_types_set.intersection(
+            model_meta.model_type
+        ):
+            continue
 
         lower, upper = n_parameters_range
         n_parameters = model_meta.n_parameters
@@ -75,7 +82,10 @@ def get_model_metas(
 
 
 def get_model(
-    model_name: str, revision: str | None = None, **kwargs: Any
+    model_name: str,
+    revision: str | None = None,
+    device: str | None = None,
+    **kwargs: Any,
 ) -> MTEBModels:
     """A function to fetch and load model object by name.
 
@@ -85,15 +95,23 @@ def get_model(
     Args:
         model_name: Name of the model to fetch
         revision: Revision of the model to fetch
+        device: Device used to load the model
         **kwargs: Additional keyword arguments to pass to the model loader
 
     Returns:
         A model object
     """
     meta = get_model_meta(model_name, revision)
-    model = meta.load_model(**kwargs)
+    model = meta.load_model(device=device, **kwargs)
+
+    if kwargs:
+        logger.info(
+            f"Model '{model_name}' loaded with additional arguments: {list(kwargs.keys())}"
+        )
+        meta = meta.model_copy(deep=True)
+        meta.loader_kwargs |= kwargs
 
-    model.mteb_model_meta = meta  # type: ignore
+    model.mteb_model_meta = meta  # type: ignore[misc]
     return model
 
 

mteb/models/instruct_wrapper.py

@@ -17,7 +17,8 @@ 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 +29,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 +42,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 +66,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 +92,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 +115,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 +133,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 +158,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 +170,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 +212,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

@@ -118,7 +118,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/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"],
@@ -17,7 +17,7 @@ model2vecdk = ModelMeta(
     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 +35,7 @@ model2vecdk = ModelMeta(
 
 
 model2vecdk_stem = ModelMeta(
-    loader=Model2VecModel,  # type: ignore
+    loader=Model2VecModel,
     name="andersborges/model2vecdk-stem",
     model_type=["dense"],
     languages=["dan-Latn"],
@@ -48,7 +48,7 @@ model2vecdk_stem = ModelMeta(
     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

@@ -16,7 +16,7 @@ arabic_triplet_matryoshka = ModelMeta(
     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,7 +145,7 @@ 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,
     memory_usage_mb=86,
     max_tokens=512,
@@ -171,7 +171,7 @@ 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,
     memory_usage_mb=127,
     max_tokens=512,
@@ -197,7 +197,7 @@ 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,
     memory_usage_mb=415,
     max_tokens=512,
@@ -223,7 +223,7 @@ 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,
     memory_usage_mb=522,
     max_tokens=2048,
@@ -248,7 +248,7 @@ 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,
     memory_usage_mb=1274,
     max_tokens=512,
@@ -278,7 +278,7 @@ 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,
     memory_usage_mb=415,
     max_tokens=512,
@@ -304,7 +304,7 @@ 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,
     memory_usage_mb=1165,
     max_tokens=8192,
@@ -329,7 +329,7 @@ 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,
     memory_usage_mb=2166,
     max_tokens=8192,

mteb/models/model_implementations/b1ade_models.py

@@ -22,7 +22,7 @@ b1ade_embed = ModelMeta(
     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,