mteb 2.5.3__py3-none-any.whl → 2.5.4__py3-none-any.whl

mteb/get_tasks.py

@@ -4,7 +4,7 @@ import difflib
 import logging
 import warnings
 from collections import Counter, defaultdict
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from typing import Any
 
 import pandas as pd
@@ -23,12 +23,11 @@ logger = logging.getLogger(__name__)
 def _gather_tasks() -> tuple[type[AbsTask], ...]:
     import mteb.tasks as tasks
 
-    tasks = [
+    return tuple(
         t
         for t in tasks.__dict__.values()
         if isinstance(t, type) and issubclass(t, AbsTask)
-    ]
-    return tuple(tasks)
+    )
 
 
 def _create_name_to_task_mapping(
@@ -44,7 +43,7 @@ def _create_name_to_task_mapping(
     return metadata_names
 
 
-def _create_similar_tasks(tasks: Sequence[type[AbsTask]]) -> dict[str, list[str]]:
+def _create_similar_tasks(tasks: Iterable[type[AbsTask]]) -> dict[str, list[str]]:
     """Create a dictionary of similar tasks.
 
     Returns:
@@ -195,9 +194,8 @@ class MTEBTasks(tuple[AbsTask]):
             string with a LaTeX table.
         """
         if include_citation_in_name and "name" in properties:
-            properties += ["intext_citation"]
-            df = self.to_dataframe(properties)
-            df["name"] = df["name"] + " " + df["intext_citation"]
+            df = self.to_dataframe(tuple(properties) + ("intext_citation",))
+            df["name"] = df["name"] + " " + df["intext_citation"]  # type: ignore[operator]
             df = df.drop(columns=["intext_citation"])
         else:
             df = self.to_dataframe(properties)
@@ -222,17 +220,17 @@ class MTEBTasks(tuple[AbsTask]):
 
 
 def get_tasks(
-    tasks: list[str] | None = None,
+    tasks: Sequence[str] | None = None,
     *,
-    languages: list[str] | None = None,
-    script: list[str] | None = None,
-    domains: list[TaskDomain] | None = None,
-    task_types: list[TaskType] | None = None,  # type: ignore
-    categories: list[TaskCategory] | None = None,
+    languages: Sequence[str] | None = None,
+    script: Sequence[str] | None = None,
+    domains: Sequence[TaskDomain] | None = None,
+    task_types: Sequence[TaskType] | None = None,
+    categories: Sequence[TaskCategory] | None = None,
     exclude_superseded: bool = True,
-    eval_splits: list[str] | None = None,
+    eval_splits: Sequence[str] | None = None,
     exclusive_language_filter: bool = False,
-    modalities: list[Modalities] | None = None,
+    modalities: Sequence[Modalities] | None = None,
     exclusive_modality_filter: bool = False,
     exclude_aggregate: bool = False,
     exclude_private: bool = True,
@@ -288,7 +286,7 @@ def get_tasks(
         ]
         return MTEBTasks(_tasks)
 
-    _tasks = filter_tasks(
+    tasks_: Sequence[type[AbsTask]] = filter_tasks(
         TASK_LIST,
         languages=languages,
         script=script,
@@ -301,12 +299,12 @@ def get_tasks(
         exclude_aggregate=exclude_aggregate,
         exclude_private=exclude_private,
     )
-    _tasks = [
-        cls().filter_languages(languages, script).filter_eval_splits(eval_splits)
-        for cls in _tasks
-    ]
-
-    return MTEBTasks(_tasks)
+    return MTEBTasks(
+        [
+            cls().filter_languages(languages, script).filter_eval_splits(eval_splits)
+            for cls in tasks_
+        ]
+    )
 
 
 _TASK_RENAMES = {"PersianTextTone": "SynPerTextToneClassification"}
@@ -314,10 +312,10 @@ _TASK_RENAMES = {"PersianTextTone": "SynPerTextToneClassification"}
 
 def get_task(
     task_name: str,
-    languages: list[str] | None = None,
-    script: list[str] | None = None,
-    eval_splits: list[str] | None = None,
-    hf_subsets: list[str] | None = None,
+    languages: Sequence[str] | None = None,
+    script: Sequence[str] | None = None,
+    eval_splits: Sequence[str] | None = None,
+    hf_subsets: Sequence[str] | None = None,
     exclusive_language_filter: bool = False,
 ) -> AbsTask:
     """Get a task by name.

mteb/languages/language_scripts.py

@@ -1,9 +1,9 @@
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 
 from typing_extensions import Self
 
-from mteb.languages import check_language_code
+from mteb.languages.check_language_code import check_language_code
 
 
 @dataclass
@@ -25,7 +25,9 @@ class LanguageScripts:
 
     @classmethod
     def from_languages_and_scripts(
-        cls, languages: list[str] | None = None, scripts: list[str] | None = None
+        cls,
+        languages: Sequence[str] | None = None,
+        scripts: Sequence[str] | None = None,
     ) -> Self:
         """Create a LanguageScripts object from lists of languages and scripts.
 

mteb/leaderboard/app.py

@@ -169,7 +169,7 @@ def _update_task_info(task_names: str) -> gr.DataFrame:
     df = df.drop(columns="reference")
     return gr.DataFrame(
         df,
-        datatype=["markdown"] + ["str"] * (len(df.columns) - 1),  # type: ignore
+        datatype=["markdown"] + ["str"] * (len(df.columns) - 1),
         buttons=["copy", "fullscreen"],
         show_search="filter",
     )

mteb/load_results.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import sys
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from pathlib import Path
 
 from mteb.abstasks.abstask import AbsTask
@@ -45,8 +45,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 +83,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

@@ -44,7 +44,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(
@@ -111,7 +111,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

mteb/models/cache_wrappers/cache_backend_protocol.py

@@ -5,8 +5,6 @@ from typing import Any, Protocol, runtime_checkable
 
 import numpy as np
 
-from mteb.types import BatchedInput
-
 
 @runtime_checkable
 class CacheBackendProtocol(Protocol):
@@ -26,7 +24,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 +32,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 +51,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,12 +1,13 @@
 import hashlib
+from collections.abc import Mapping
+from typing import Any
 
-from mteb.types import BatchedInput
 
-
-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

mteb/models/cache_wrappers/cache_backends/faiss_cache.py

@@ -2,6 +2,7 @@ import json
 import logging
 import warnings
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 
@@ -37,7 +38,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
 

mteb/models/cache_wrappers/cache_backends/numpy_cache.py

@@ -2,11 +2,10 @@ 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__)
@@ -15,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"
@@ -28,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:
@@ -39,7 +38,12 @@ 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:
                     msg = f"Hash collision or duplicate item for hash {item_hash}. Overwriting existing vector."
@@ -75,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),
         )
@@ -147,9 +159,11 @@ 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:
                     msg = "Vector dimension not set. Unable to load vectors file."
@@ -164,8 +178,11 @@ class NumpyCache:
             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:
@@ -177,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

@@ -93,7 +93,14 @@ def get_model(
     meta = get_model_meta(model_name, revision)
     model = meta.load_model(**kwargs)
 
-    model.mteb_model_meta = meta  # type: ignore
+    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[misc]
     return model
 
 

mteb/models/instruct_wrapper.py

@@ -17,7 +17,7 @@ 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,
     **kwargs,
 ):
     """Instruct wrapper for models. Uses GritLM to pass instructions to the model.
@@ -40,7 +40,9 @@ def instruct_wrapper(
             self,
             model_name_or_path: str,
             mode: str,
-            instruction_template: str | Callable[[str, PromptType], str] | None = None,
+            instruction_template: str
+            | Callable[[str, PromptType | None], str]
+            | None = None,
             **kwargs,
         ):
             if (
@@ -82,8 +84,11 @@ 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]
+                _inputs,  # type: ignore[arg-type]
+                instruction=instruction,
+                *args,
+                **kwargs,
             )
             if isinstance(embeddings, torch.Tensor):
                 # sometimes in kwargs can be return_tensors=True
@@ -141,7 +146,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:
@@ -193,6 +198,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/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"],
@@ -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"],

mteb/models/model_implementations/blip_models.py

@@ -128,7 +128,7 @@ class BLIPModel(AbsEncoder):
 
 # in descending order of usage (downloads from huggingface)
 blip_image_captioning_large = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-image-captioning-large",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -156,7 +156,7 @@ blip_image_captioning_large = ModelMeta(
 )
 
 blip_image_captioning_base = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-image-captioning-base",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -185,7 +185,7 @@ blip_image_captioning_base = ModelMeta(
 
 
 blip_vqa_base = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-vqa-base",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -212,7 +212,7 @@ blip_vqa_base = ModelMeta(
 )
 
 blip_vqa_capfilt_large = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-vqa-capfilt-large",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -239,7 +239,7 @@ blip_vqa_capfilt_large = ModelMeta(
 )
 
 blip_itm_base_coco = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-itm-base-coco",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -266,7 +266,7 @@ blip_itm_base_coco = ModelMeta(
 )
 
 blip_itm_large_coco = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-itm-large-coco",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -294,7 +294,7 @@ blip_itm_large_coco = ModelMeta(
 )
 
 blip_itm_base_flickr = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-itm-base-flickr",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -322,7 +322,7 @@ blip_itm_base_flickr = ModelMeta(
 )
 
 blip_itm_large_flickr = ModelMeta(
-    loader=BLIPModel,  # type: ignore
+    loader=BLIPModel,
     name="Salesforce/blip-itm-large-flickr",
     model_type=["dense"],
     languages=["eng-Latn"],

mteb/models/model_implementations/bm25.py

@@ -113,7 +113,7 @@ def bm25_loader(model_name, **kwargs) -> SearchProtocol:
 
         def encode(self, texts: list[str]):
             """Encode input text as term vectors"""
-            return bm25s.tokenize(texts, stopwords=self.stopwords, stemmer=self.stemmer)  # type: ignore
+            return bm25s.tokenize(texts, stopwords=self.stopwords, stemmer=self.stemmer)
 
     return BM25Search(**kwargs)
 

mteb/models/model_implementations/clip_models.py

@@ -115,7 +115,7 @@ CLIP_CITATION = """
 
 
 clip_vit_large_patch14 = ModelMeta(
-    loader=CLIPModel,  # type: ignore
+    loader=CLIPModel,
     name="openai/clip-vit-large-patch14",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -139,7 +139,7 @@ clip_vit_large_patch14 = ModelMeta(
 )
 
 clip_vit_base_patch32 = ModelMeta(
-    loader=CLIPModel,  # type: ignore
+    loader=CLIPModel,
     name="openai/clip-vit-base-patch32",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -163,7 +163,7 @@ clip_vit_base_patch32 = ModelMeta(
 )
 
 clip_vit_base_patch16 = ModelMeta(
-    loader=CLIPModel,  # type: ignore
+    loader=CLIPModel,
     name="openai/clip-vit-base-patch16",
     model_type=["dense"],
     languages=["eng-Latn"],

mteb/models/model_implementations/cohere_models.py

@@ -222,7 +222,7 @@ class CohereTextEmbeddingModel(AbsEncoder):
     ) -> None:
         requires_package(self, "cohere", model_name, "pip install 'mteb[cohere]'")
 
-        import cohere  # type: ignore
+        import cohere
 
         self.model_name = model_name.removeprefix("Cohere/Cohere-")
         self.sep = sep

mteb/models/model_implementations/cohere_v.py

@@ -378,7 +378,7 @@ def cohere_v_loader(model_name, **kwargs):
 
 
 cohere_mult_3 = ModelMeta(
-    loader=cohere_v_loader,  # type: ignore
+    loader=cohere_v_loader,
     loader_kwargs={"model_name": "embed-multilingual-v3.0"},
     name="cohere/embed-multilingual-v3.0",
     model_type=["dense"],
@@ -402,7 +402,7 @@ cohere_mult_3 = ModelMeta(
 )
 
 cohere_eng_3 = ModelMeta(
-    loader=cohere_v_loader,  # type: ignore
+    loader=cohere_v_loader,
     loader_kwargs={"model_name": "embed-english-v3.0"},
     name="cohere/embed-english-v3.0",
     model_type=["dense"],