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

mteb/leaderboard/table.py

@@ -1,3 +1,7 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
@@ -5,8 +9,9 @@ import pandas as pd
 from matplotlib.colors import LinearSegmentedColormap
 from pandas.api.types import is_numeric_dtype
 
-from mteb.benchmarks.benchmark import Benchmark
-from mteb.results.benchmark_results import BenchmarkResults
+if TYPE_CHECKING:
+    from mteb.benchmarks.benchmark import Benchmark
+    from mteb.results.benchmark_results import BenchmarkResults
 
 
 def _borda_count(scores: pd.Series) -> pd.Series:

mteb/load_results.py

@@ -1,13 +1,19 @@
+from __future__ import annotations
+
 import json
 import logging
 import sys
-from collections.abc import Sequence
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 from mteb.abstasks.abstask import AbsTask
 from mteb.models.model_meta import ModelMeta
 from mteb.results import BenchmarkResults, ModelResult, TaskResult
-from mteb.types import ModelName, Revision
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Sequence
+    from pathlib import Path
+
+    from mteb.types import ModelName, Revision
 
 if sys.version_info >= (3, 13):
     from warnings import deprecated
@@ -45,8 +51,8 @@ def _model_name_and_revision(
 def load_results(
     results_repo: str = "https://github.com/embeddings-benchmark/results",
     download_latest: bool = True,
-    models: Sequence[ModelMeta] | Sequence[str] | None = None,
-    tasks: Sequence[AbsTask] | Sequence[str] | None = None,
+    models: Iterable[ModelMeta] | Sequence[str] | None = None,
+    tasks: Iterable[AbsTask] | Sequence[str] | None = None,
     validate_and_filter: bool = True,
     require_model_meta: bool = True,
     only_main_score: bool = False,
@@ -83,21 +89,21 @@ def load_results(
 
     if models is not None:
         models_to_keep = {}
-        for model_path in models:
-            if isinstance(model_path, ModelMeta):
-                models_to_keep[model_path.name] = model_path.revision
+        for model in models:
+            if isinstance(model, ModelMeta):
+                models_to_keep[model.name] = model.revision
             else:
-                models_to_keep[model_path] = None
+                models_to_keep[model] = None
     else:
         models_to_keep = None
 
-    task_names = {}
+    task_names: dict[str, AbsTask | None] = {}
     if tasks is not None:
-        for task in tasks:
-            if isinstance(task, AbsTask):
-                task_names[task.metadata.name] = task
+        for task_ in tasks:
+            if isinstance(task_, AbsTask):
+                task_names[task_.metadata.name] = task_
             else:
-                task_names[task] = None
+                task_names[task_] = None
 
     model_results = []
     for model_path in model_paths:

mteb/models/abs_encoder.py

@@ -1,12 +1,12 @@
+from __future__ import annotations
+
 import logging
+import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Sequence
-from typing import Any, Literal, cast, get_args, overload
-
-from torch.utils.data import DataLoader
+from typing import TYPE_CHECKING, Any, Literal, cast, get_args, overload
 
 import mteb
-from mteb.abstasks.task_metadata import TaskMetadata, TaskType
+from mteb.abstasks.task_metadata import TaskType
 from mteb.similarity_functions import (
     cos_sim,
     dot_score,
@@ -16,12 +16,25 @@ from mteb.similarity_functions import (
     pairwise_max_sim,
 )
 from mteb.types import (
-    Array,
-    BatchedInput,
     PromptType,
 )
 
-from .model_meta import ModelMeta, ScoringFunction
+from .model_meta import ScoringFunction
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from torch.utils.data import DataLoader
+    from typing_extensions import Unpack
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import (
+        Array,
+        BatchedInput,
+        EncodeKwargs,
+    )
+
+    from .model_meta import ModelMeta
 
 logger = logging.getLogger(__name__)
 
@@ -43,7 +56,7 @@ class AbsEncoder(ABC):
     model: Any
     mteb_model_meta: ModelMeta | None = None
     model_prompts: dict[str, str] | None = None
-    instruction_template: str | Callable[[str, PromptType], str] | None = None
+    instruction_template: str | Callable[[str, PromptType | None], str] | None = None
     prompts_dict: dict[str, str] | None = None
 
     def get_prompt_name(
@@ -110,7 +123,7 @@ class AbsEncoder(ABC):
         if not self.model_prompts:
             return None
         prompt_name = self.get_prompt_name(task_metadata, prompt_type)
-        return self.model_prompts.get(prompt_name)
+        return self.model_prompts.get(prompt_name) if prompt_name else None
 
     @staticmethod
     @overload
@@ -187,6 +200,7 @@ class AbsEncoder(ABC):
                 except KeyError:
                     msg = f"Task name {task_name} is not valid. {valid_keys_msg}"
                     logger.warning(msg)
+                    warnings.warn(msg)
                     invalid_task_messages.add(msg)
                     invalid_keys.add(task_key)
 
@@ -232,9 +246,9 @@ class AbsEncoder(ABC):
         if isinstance(prompt, dict) and prompt_type:
             if prompt.get(prompt_type.value):
                 return prompt[prompt_type.value]
-            logger.warning(
-                f"Prompt type '{prompt_type}' not found in task metadata for task '{task_metadata.name}'."
-            )
+            msg = f"Prompt type '{prompt_type}' not found in task metadata for task '{task_metadata.name}'."
+            logger.warning(msg)
+            warnings.warn(msg)
             return ""
 
         if prompt:
@@ -310,7 +324,7 @@ class AbsEncoder(ABC):
             ):
                 arr = self.model.similarity(embeddings1, embeddings2)
                 # We assume that the model returns an Array-like object:
-                arr = cast(Array, arr)
+                arr = cast("Array", arr)
                 return arr
             return cos_sim(embeddings1, embeddings2)
         if self.mteb_model_meta.similarity_fn_name is ScoringFunction.COSINE:
@@ -348,7 +362,7 @@ class AbsEncoder(ABC):
             ):
                 arr = self.model.similarity_pairwise(embeddings1, embeddings2)
                 # We assume that the model returns an Array-like object:
-                arr = cast(Array, arr)
+                arr = cast("Array", arr)
                 return arr
             return pairwise_cos_sim(embeddings1, embeddings2)
         if self.mteb_model_meta.similarity_fn_name is ScoringFunction.COSINE:
@@ -368,7 +382,7 @@ class AbsEncoder(ABC):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Encodes the given sentences using the encoder.
 

mteb/models/cache_wrappers/cache_backend_protocol.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from pathlib import Path
-from typing import Any, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
-import numpy as np
+if TYPE_CHECKING:
+    from pathlib import Path
 
-from mteb.types import BatchedInput
+    import numpy as np
 
 
 @runtime_checkable
@@ -26,7 +26,7 @@ class CacheBackendProtocol(Protocol):
             **kwargs: Additional backend-specific arguments.
         """
 
-    def add(self, item: list[BatchedInput], vectors: np.ndarray) -> None:
+    def add(self, item: list[dict[str, Any]], vectors: np.ndarray) -> None:
         """Add a vector to the cache.
 
         Args:
@@ -34,7 +34,7 @@ class CacheBackendProtocol(Protocol):
             vectors: Embedding vector of shape (dim,) or (1, dim).
         """
 
-    def get_vector(self, item: BatchedInput) -> np.ndarray | None:
+    def get_vector(self, item: dict[str, Any]) -> np.ndarray | None:
         """Retrieve the cached vector for the given item.
 
         Args:
@@ -53,5 +53,5 @@ class CacheBackendProtocol(Protocol):
     def close(self) -> None:
         """Release resources or flush data."""
 
-    def __contains__(self, item: BatchedInput) -> bool:
+    def __contains__(self, item: dict[str, Any]) -> bool:
         """Check whether the cache contains an item."""

mteb/models/cache_wrappers/cache_backends/_hash_utils.py

@@ -1,16 +1,21 @@
+from __future__ import annotations
+
 import hashlib
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
 
-from mteb.types import BatchedInput
+    from PIL import Image
 
 
-def _hash_item(item: BatchedInput) -> str:
+def _hash_item(item: Mapping[str, Any]) -> str:
     item_hash = ""
     if "text" in item:
-        item_hash = hashlib.sha256(item["text"].encode()).hexdigest()
+        item_text: str = item["text"]
+        item_hash = hashlib.sha256(item_text.encode()).hexdigest()
 
     if "image" in item:
-        from PIL import Image
-
         image: Image.Image = item["image"]
         item_hash += hashlib.sha256(image.tobytes()).hexdigest()
 

mteb/models/cache_wrappers/cache_backends/faiss_cache.py

@@ -1,14 +1,22 @@
+from __future__ import annotations
+
 import json
 import logging
+import warnings
 from pathlib import Path
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 from mteb._requires_package import requires_package
-from mteb.types import BatchedInput
 
 from ._hash_utils import _hash_item
 
+if TYPE_CHECKING:
+    import faiss
+
+    from mteb.types import BatchedInput
+
 logger = logging.getLogger(__name__)
 
 
@@ -22,7 +30,6 @@ class FaissCache:
             "FAISS-based vector cache",
             install_instruction="pip install mteb[faiss-cpu]",
         )
-        import faiss
 
         self.directory = Path(directory)
         self.directory.mkdir(parents=True, exist_ok=True)
@@ -36,7 +43,7 @@ class FaissCache:
         logger.info(f"Initialized FAISS VectorCacheMap in {self.directory}")
         self.load()
 
-    def add(self, items: list[BatchedInput], vectors: np.ndarray) -> None:
+    def add(self, items: list[dict[str, Any]], vectors: np.ndarray) -> None:
         """Add vector to FAISS index."""
         import faiss
 
@@ -71,7 +78,9 @@ class FaissCache:
         try:
             return self.index.reconstruct(idx)
         except Exception:
-            logger.warning(f"Vector id {idx} missing for hash {item_hash}")
+            msg = f"Vector id {idx} missing for hash {item_hash}"
+            logger.warning(msg)
+            warnings.warn(msg)
             return None
 
     def save(self) -> None:

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

@@ -1,21 +1,26 @@
+from __future__ import annotations
+
 import logging
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
 from datasets import Dataset
-from torch.utils.data import DataLoader
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models.cache_wrappers.cache_backend_protocol import (
-    CacheBackendProtocol,
-)
 from mteb.models.cache_wrappers.cache_backends.numpy_cache import NumpyCache
-from mteb.models.model_meta import ModelMeta
-from mteb.models.models_protocols import EncoderProtocol
-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.models.cache_wrappers.cache_backend_protocol import (
+        CacheBackendProtocol,
+    )
+    from mteb.models.model_meta import ModelMeta
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import Array, BatchedInput, PromptType
 
 logger = logging.getLogger(__name__)
 
@@ -90,9 +95,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

@@ -1,15 +1,22 @@
+from __future__ import annotations
+
 import difflib
 import logging
-from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from mteb.abstasks import AbsTask
 from mteb.models import (
     ModelMeta,
-    MTEBModels,
 )
 from mteb.models.model_implementations import MODEL_REGISTRY
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from mteb.abstasks import AbsTask
+    from mteb.models import (
+        MTEBModels,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -21,6 +28,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 +41,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 +50,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 +67,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 +89,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,20 +102,31 @@ 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
 
 
 def get_model_meta(
-    model_name: str, revision: str | None = None, fetch_from_hf: bool = True
+    model_name: str,
+    revision: str | None = None,
+    fetch_from_hf: bool = True,
+    fill_missing: bool = False,
 ) -> ModelMeta:
     """A function to fetch a model metadata object by name.
 
@@ -106,6 +134,7 @@ def get_model_meta(
         model_name: Name of the model to fetch
         revision: Revision of the model to fetch
         fetch_from_hf: Whether to fetch the model from HuggingFace Hub if not found in the registry
+        fill_missing: Computes missing attributes from the metadata including number of parameters and memory usage.
 
     Returns:
         A model metadata object
@@ -117,10 +146,25 @@ def get_model_meta(
             raise ValueError(
                 f"Model revision {revision} not found for model {model_name}. Expected {model_meta.revision}."
             )
+
+        if fill_missing and fetch_from_hf:
+            original_meta_dict = model_meta.model_dump()
+            new_meta = ModelMeta.from_hub(model_name)
+            new_meta_dict = new_meta.model_dump(exclude_none=True)
+
+            updates = {
+                k: v
+                for k, v in new_meta_dict.items()
+                if original_meta_dict.get(k) is None
+            }
+
+            if updates:
+                return model_meta.model_copy(update=updates)
         return model_meta
+
     if fetch_from_hf:
         logger.info(
-            "Model not found in model registry. Attempting to extract metadata by loading the model ({model_name}) using HuggingFace."
+            f"Model not found in model registry. Attempting to extract metadata by loading the model ({model_name}) using HuggingFace."
         )
         meta = ModelMeta.from_hub(model_name, revision)
         return meta