mteb 2.7.16__py3-none-any.whl → 2.7.18__py3-none-any.whl

mteb/abstasks/retrieval_dataset_loaders.py

@@ -78,7 +78,7 @@ class RetrievalDatasetLoader:
 
     def load(
         self,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> RetrievalSplitData:
         """Loads the dataset split for the specified configuration.
 
@@ -128,7 +128,11 @@ class RetrievalDatasetLoader:
             f"Split {self.split} not found in {splits}. Please specify a valid split."
         )
 
-    def _load_dataset_split(self, config: str, num_proc: int) -> Dataset:
+    def _load_dataset_split(
+        self,
+        config: str,
+        num_proc: int | None,
+    ) -> Dataset:
         return load_dataset(
             self.hf_repo,
             config,
@@ -138,7 +142,10 @@ class RetrievalDatasetLoader:
             num_proc=num_proc,
         )
 
-    def _load_corpus(self, num_proc: int) -> CorpusDatasetType:
+    def _load_corpus(
+        self,
+        num_proc: int | None,
+    ) -> CorpusDatasetType:
         config = f"{self.config}-corpus" if self.config is not None else "corpus"
         logger.info("Loading corpus subset: %s", config)
 
@@ -151,7 +158,10 @@ class RetrievalDatasetLoader:
         logger.debug("Doc Example: %s", corpus_ds[0])
         return corpus_ds
 
-    def _load_queries(self, num_proc: int) -> QueryDatasetType:
+    def _load_queries(
+        self,
+        num_proc: int | None,
+    ) -> QueryDatasetType:
         config = f"{self.config}-queries" if self.config is not None else "queries"
         logger.info("Loading queries subset: %s", config)
 
@@ -168,7 +178,10 @@ class RetrievalDatasetLoader:
 
         return queries_ds
 
-    def _load_qrels(self, num_proc: int) -> RelevantDocumentsType:
+    def _load_qrels(
+        self,
+        num_proc: int | None,
+    ) -> RelevantDocumentsType:
         config = f"{self.config}-qrels" if self.config is not None else "default"
 
         logger.info("Loading qrels subset: %s", config)
@@ -203,7 +216,10 @@ class RetrievalDatasetLoader:
         logger.info("Loaded %d %s qrels.", len(qrels_dict), self.split.upper())
         return qrels_dict
 
-    def _load_top_ranked(self, num_proc: int) -> TopRankedDocumentsType:
+    def _load_top_ranked(
+        self,
+        num_proc: int | None,
+    ) -> TopRankedDocumentsType:
         config = (
             f"{self.config}-top_ranked" if self.config is not None else "top_ranked"
         )
@@ -226,7 +242,10 @@ class RetrievalDatasetLoader:
         logger.info(f"Top ranked loaded: {len(top_ranked_ds)}")
         return top_ranked_dict
 
-    def _load_instructions(self, num_proc: int) -> InstructionDatasetType:
+    def _load_instructions(
+        self,
+        num_proc: int | None,
+    ) -> InstructionDatasetType:
         config = (
             f"{self.config}-instruction" if self.config is not None else "instruction"
         )
@@ -246,7 +265,7 @@ class RetrievalDatasetLoader:
 def _combine_queries_with_instructions_datasets(
     queries_dataset: QueryDatasetType,
     instruction_dataset: InstructionDatasetType | dict[str, str],
-    num_proc: int,
+    num_proc: int | None,
 ) -> Dataset:
     if isinstance(instruction_dataset, Dataset):
         instruction_to_query_idx = {

mteb/abstasks/sts.py

@@ -118,7 +118,7 @@ class AbsTaskSTS(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> STSMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/text/bitext_mining.py

@@ -82,7 +82,7 @@ class AbsTaskBitextMining(AbsTask):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
         """Added load for "parallel" datasets"""
@@ -155,7 +155,7 @@ class AbsTaskBitextMining(AbsTask):
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         parallel: bool = False,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> BitextMiningMetrics | dict[str, BitextMiningMetrics]:
         pairs = self._get_pairs(parallel)

mteb/abstasks/text/reranking.py

@@ -34,7 +34,7 @@ class AbsTaskReranking(AbsTaskRetrieval):
         For dataformat and other information, see [AbsTaskRetrieval][mteb.abstasks.retrieval.AbsTaskRetrieval].
     """
 
-    def load_data(self, num_proc: int = 1, **kwargs) -> None:
+    def load_data(self, num_proc: int | None = None, **kwargs) -> None:
         """Load the dataset."""
         if self.data_loaded:
             return

mteb/abstasks/text/summarization.py

@@ -94,7 +94,7 @@ class AbsTaskSummarization(AbsTask):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> SummarizationMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/zeroshot_classification.py

@@ -127,7 +127,7 @@ class AbsTaskZeroShotClassification(AbsTask):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> ZeroShotClassificationMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/benchmarks/benchmark.py

@@ -164,14 +164,142 @@ class MIEBBenchmark(Benchmark):
 class VidoreBenchmark(Benchmark):
     """Wrapper for Vidore3 benchmark."""
 
-    def _create_summary_table(
+    def _create_vidore_summary_table(
         self, benchmark_results: BenchmarkResults
     ) -> pd.DataFrame:
+        """Create summary table from BenchmarkResults.
+
+        Returns a DataFrame with one row per model containing summary statistics
+        and task type averages. Customized for Vidore benchmark.
+
+        Args:
+            benchmark_results: BenchmarkResults object containing model results
+
+        Returns:
+            DataFrame with model summaries, ready for styling in the leaderboard
+        """
+        import mteb
         from mteb.benchmarks._create_table import (
-            _create_summary_table_mean_public_private,
+            _format_max_tokens,
+            _format_n_parameters,
+            _get_means_per_types,
+            _split_on_capital,
+        )
+        from mteb.get_tasks import get_task
+
+        data = benchmark_results.to_dataframe(format="long")
+
+        if data.empty:
+            no_results_frame = pd.DataFrame(
+                {"No results": ["You can try relaxing your criteria"]}
+            )
+            return no_results_frame
+        public_task_name = benchmark_results._filter_tasks(is_public=True).task_names
+        private_task_name = benchmark_results._filter_tasks(is_public=False).task_names
+        # Convert to DataFrame and pivot
+        per_task = data.pivot(index="model_name", columns="task_name", values="score")
+
+        # Remove models with no scores
+        to_remove = per_task.isna().all(axis="columns")
+        if to_remove.all():
+            no_results_frame = pd.DataFrame(
+                {"No results": ["You can try relaxing your criteria"]}
+            )
+            return no_results_frame
+
+        models_to_remove = list(per_task[to_remove].index)
+        per_task = per_task.drop(models_to_remove, axis=0)
+
+        # Calculate means by task type
+        mean_per_type = _get_means_per_types(per_task)
+        mean_per_type = mean_per_type.pivot(
+            index="model_name", columns="task_type", values="score"
+        )
+        mean_per_type.columns = [
+            _split_on_capital(column) for column in mean_per_type.columns
+        ]
+
+        # Calculate overall means
+        public_mean = per_task[public_task_name].mean(skipna=False, axis=1)
+        private_mean = per_task[private_task_name].mean(skipna=False, axis=1)
+
+        # Build joint table
+        joint_table = mean_per_type.copy()
+        joint_table.insert(1, "mean(public)", public_mean)
+        joint_table.insert(2, "mean(private)", private_mean)
+        task_type = get_task(
+            per_task.columns[0]
+        ).metadata.type  # "DocumentUnderstanding"
+        joint_table = joint_table.sort_values(
+            [_split_on_capital(task_type), "mean(public)", "mean(private)"],
+            ascending=False,
+        )
+
+        joint_table = joint_table.reset_index()
+
+        # Add model metadata
+        model_metas = joint_table["model_name"].map(mteb.get_model_meta)
+        joint_table = joint_table[model_metas.notna()]
+        joint_table["model_link"] = model_metas.map(lambda m: m.reference)
+
+        # Insert model metadata columns
+        joint_table.insert(
+            1,
+            "Max Tokens",
+            model_metas.map(lambda m: _format_max_tokens(m.max_tokens)),
+        )
+        joint_table.insert(
+            1,
+            "Embedding Dimensions",
+            model_metas.map(lambda m: int(m.embed_dim) if m.embed_dim else None),
+        )
+        joint_table.insert(
+            1,
+            "Number of Parameters (B)",
+            model_metas.map(lambda m: _format_n_parameters(m.n_parameters)),
+        )
+        joint_table.insert(
+            1,
+            "Memory Usage (MB)",
+            model_metas.map(
+                lambda m: int(m.memory_usage_mb) if m.memory_usage_mb else None
+            ),
+        )
+
+        # Clean up model names (remove HF organization)
+        joint_table["model_name"] = joint_table["model_name"].map(
+            lambda name: name.split("/")[-1]
+        )
+
+        # Add markdown links to model names
+        name_w_link = (
+            "[" + joint_table["model_name"] + "](" + joint_table["model_link"] + ")"
+        )
+        joint_table["model_name"] = joint_table["model_name"].mask(
+            joint_table["model_link"].notna(), name_w_link
+        )
+        joint_table = joint_table.drop(columns=["model_link"])
+
+        # Rename columns
+        rename_dict = {
+            "model_name": "Model",
+            "mean(public)": "Mean (Public)",
+            "mean(private)": "Mean (Private)",
+        }
+
+        joint_table = joint_table.rename(columns=rename_dict)
+
+        # Add Rank column
+        joint_table.insert(
+            0, "Rank (Mean Task)", [i + 1 for i in range(len(joint_table))]
         )
 
-        joint_table = _create_summary_table_mean_public_private(benchmark_results)
+        return joint_table
+
+    def _create_summary_table(
+        self, benchmark_results: BenchmarkResults
+    ) -> pd.DataFrame:
+        joint_table = self._create_vidore_summary_table(benchmark_results)
         # For ViDoRe (V1, V2, V3): all tasks are Document Understanding type, so Document Understanding column = Mean (Task)
         joint_table = joint_table.rename(
             columns={"Document Understanding": "Mean (Task)"}

mteb/evaluate.py

@@ -91,7 +91,7 @@ def _evaluate_task(
     encode_kwargs: EncodeKwargs,
     prediction_folder: Path | None,
     public_only: bool | None,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> TaskResult | TaskError:
     """The core logic to run a model on a given task. See `evaluate` for more details.
 
@@ -282,7 +282,7 @@ def evaluate(
     prediction_folder: Path | str | None = None,
     show_progress_bar: bool = True,
     public_only: bool | None = None,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> ModelResult:
     """This function runs a model on a given task and returns the results.
 

mteb/leaderboard/figures.py

@@ -125,6 +125,7 @@ def _performance_size_plot(df: pd.DataFrame) -> go.Figure:
     min_score, max_score = df["Mean (Task)"].min(), df["Mean (Task)"].max()
     df["sqrt(dim)"] = np.sqrt(df["Embedding Dimensions"])
     df["Max Tokens"] = df["Max Tokens"].apply(lambda x: _process_max_tokens(x))
+    rank_column = "Rank (Borda)" if "Rank (Borda)" in df.columns else "Rank (Mean Task)"
     fig = px.scatter(
         df,
         x="Number of Parameters",
@@ -141,7 +142,7 @@ def _performance_size_plot(df: pd.DataFrame) -> go.Figure:
             "Embedding Dimensions": True,
             "Number of Parameters": True,
             "Mean (Task)": True,
-            "Rank (Borda)": True,
+            rank_column: True,
             "Log(Tokens)": False,
             "sqrt(dim)": False,
             "model_text": False,

mteb/leaderboard/table.py

@@ -156,6 +156,7 @@ def _apply_summary_table_styling(joint_table: pd.DataFrame) -> gr.DataFrame:
     """
     excluded_columns = [
         "Rank (Borda)",
+        "Rank (Mean Task)",
         "Rank",
         "Model",
         "Number of Parameters (B)",
@@ -183,10 +184,17 @@ def _apply_summary_table_styling(joint_table: pd.DataFrame) -> gr.DataFrame:
         joint_table["Zero-shot"] = joint_table["Zero-shot"].apply(_format_zero_shot)
     joint_table[score_columns] = joint_table[score_columns].map(_format_scores)
 
+    if "Rank (Borda)" in joint_table.columns:
+        rank_column = "Rank (Borda)"
+    elif "Rank (Mean Task)" in joint_table.columns:
+        rank_column = "Rank (Mean Task)"
+    else:
+        raise ValueError("No rank column found in the result table.")
+
     joint_table_style = joint_table.style.format(
         {
             **dict.fromkeys(score_columns, "{:.2f}"),
-            "Rank (Borda)": "{:.0f}",
+            rank_column: "{:.0f}",
             "Memory Usage (MB)": "{:.0f}",
             "Embedding Dimensions": "{:.0f}",
             "Max Tokens": "{:.0f}",
@@ -195,7 +203,7 @@ def _apply_summary_table_styling(joint_table: pd.DataFrame) -> gr.DataFrame:
         na_rep="",
     )
     joint_table_style = joint_table_style.highlight_min(
-        "Rank (Borda)", props="font-weight: bold"
+        rank_column, props="font-weight: bold"
     ).highlight_max(subset=score_columns, props="font-weight: bold")
 
     # Apply background gradients for each selected column

mteb/models/cache_wrappers/cache_backend_protocol.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 if TYPE_CHECKING:
     from pathlib import Path
 
-    import numpy as np
+    from mteb.types import Array
 
 
 @runtime_checkable
@@ -26,7 +26,7 @@ class CacheBackendProtocol(Protocol):
             **kwargs: Additional backend-specific arguments.
         """
 
-    def add(self, item: list[dict[str, Any]], vectors: np.ndarray) -> None:
+    def add(self, item: list[dict[str, Any]], vectors: Array) -> 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: dict[str, Any]) -> np.ndarray | None:
+    def get_vector(self, item: dict[str, Any]) -> Array | None:
         """Retrieve the cached vector for the given item.
 
         Args:

mteb/models/cache_wrappers/cache_backends/faiss_cache.py

@@ -15,7 +15,7 @@ from ._hash_utils import _hash_item
 if TYPE_CHECKING:
     import faiss
 
-    from mteb.types import BatchedInput
+    from mteb.types import Array, BatchedInput
 
 logger = logging.getLogger(__name__)
 
@@ -43,7 +43,7 @@ class FaissCache:
         logger.info(f"Initialized FAISS VectorCacheMap in {self.directory}")
         self.load()
 
-    def add(self, items: list[dict[str, Any]], vectors: np.ndarray) -> None:
+    def add(self, items: list[dict[str, Any]], vectors: Array) -> None:
         """Add vector to FAISS index."""
         import faiss
 
@@ -67,7 +67,7 @@ class FaissCache:
             vectors_array = np.vstack(vectors_to_add).astype(np.float32)
             self.index.add(vectors_array)
 
-    def get_vector(self, item: BatchedInput) -> np.ndarray | None:
+    def get_vector(self, item: dict[str, Any]) -> Array | None:
         """Retrieve vector from index by hash."""
         if self.index is None:
             return None

mteb/models/cache_wrappers/cache_backends/numpy_cache.py

@@ -1,13 +1,18 @@
+from __future__ import annotations
+
 import json
 import logging
 import warnings
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 from ._hash_utils import _hash_item
 
+if TYPE_CHECKING:
+    from mteb.types import Array
+
 logger = logging.getLogger(__name__)
 
 
@@ -27,7 +32,7 @@ class NumpyCache:
         logger.info(f"Initialized VectorCacheMap in directory: {self.directory}")
         self._initialize_vectors_file()
 
-    def add(self, items: list[dict[str, Any]], vectors: np.ndarray) -> None:
+    def add(self, items: list[dict[str, Any]], vectors: Array) -> None:
         """Add a vector to the cache."""
         try:
             if self.vector_dim is None:
@@ -178,7 +183,7 @@ class NumpyCache:
             logger.error(f"Error loading VectorCacheMap: {str(e)}")
             raise
 
-    def get_vector(self, item: dict[str, Any]) -> np.ndarray | None:
+    def get_vector(self, item: dict[str, Any]) -> Array | None:
         """Retrieve vector from index by hash."""
         if self.vectors is None:
             return None

mteb/models/cache_wrappers/cache_wrapper.py

@@ -98,7 +98,7 @@ class CachedEmbeddingWrapper:
             uncached_items: list[dict[str, Any]] = []
             uncached_indices: list[int] = []
             all_items: Dataset = inputs.dataset
-            cached_vectors: dict[int, np.ndarray] = {}
+            cached_vectors: dict[int, Array] = {}
 
             for i, item in enumerate(all_items):
                 vector = cache.get_vector(item)
@@ -108,7 +108,7 @@ class CachedEmbeddingWrapper:
                     uncached_items.append(item)
                     uncached_indices.append(i)
 
-            newly_encoded: dict[int, np.ndarray] = {}
+            newly_encoded: dict[int, Array] = {}
             if uncached_items:
                 logger.info(f"Encoding {len(uncached_items)} new items")
                 # Build a simple DataLoader with only uncached items

mteb/models/model_implementations/bedrock_models.py

@@ -86,7 +86,7 @@ class BedrockModel(AbsEncoder):
 
     def _encode_amazon(
         self, sentences: list[str], show_progress_bar: bool = False
-    ) -> np.ndarray:
+    ) -> Array:
         from botocore.exceptions import ValidationError
 
         all_embeddings = []
@@ -125,7 +125,7 @@ class BedrockModel(AbsEncoder):
         sentences: list[str],
         cohere_task_type: str,
         show_progress_bar: bool = False,
-    ) -> np.ndarray:
+    ) -> Array:
         batches = [
             sentences[i : i + self._max_batch_size]
             for i in range(0, len(sentences), self._max_batch_size)
@@ -149,7 +149,7 @@ class BedrockModel(AbsEncoder):
 
         return np.array(all_embeddings)
 
-    def _embed_amazon(self, sentence: str) -> np.ndarray:
+    def _embed_amazon(self, sentence: str) -> Array:
         response = self._client.invoke_model(
             body=json.dumps({"inputText": sentence}),
             modelId=self._model_id,
@@ -158,7 +158,7 @@ class BedrockModel(AbsEncoder):
         )
         return self._to_numpy(response)
 
-    def _to_numpy(self, embedding_response) -> np.ndarray:
+    def _to_numpy(self, embedding_response) -> Array:
         response = json.loads(embedding_response.get("body").read())
         key = "embedding" if self._provider == "amazon" else "embeddings"
         return np.array(response[key])

mteb/models/model_implementations/bm25.py

@@ -54,7 +54,7 @@ def bm25_loader(model_name, **kwargs) -> SearchProtocol:
             hf_split: str,
             hf_subset: str,
             encode_kwargs: EncodeKwargs,
-            num_proc: int = 1,
+            num_proc: int | None = None,
         ) -> None:
             logger.info("Encoding Corpus...")
             corpus_texts = [
@@ -81,7 +81,7 @@ def bm25_loader(model_name, **kwargs) -> SearchProtocol:
             top_k: int,
             encode_kwargs: EncodeKwargs,
             top_ranked: TopRankedDocumentsType | None = None,
-            num_proc: int = 1,
+            num_proc: int | None = None,
         ) -> RetrievalOutputType:
             logger.info("Encoding Queries...")
             query_ids = list(queries["id"])

mteb/models/model_implementations/mcinext_models.py

@@ -13,7 +13,7 @@ from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta
 
 if TYPE_CHECKING:
-    from mteb.types import PromptType
+    from mteb.types import Array, PromptType
 logger = logging.getLogger(__name__)
 
 HAKIM_CITATION = """@article{sarmadi2025hakim,
@@ -302,7 +302,7 @@ class HakimModelWrapper(AbsEncoder):
         prompt_type: PromptType | None = None,
         batch_size: int = 32,
         **kwargs: Any,
-    ) -> np.ndarray:
+    ) -> Array:
         """Encodes sentences using the API.
 
         Returns:

mteb/models/model_implementations/openai_models.py

@@ -11,6 +11,7 @@ from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
 
 if TYPE_CHECKING:
+    from numpy.typing import NDArray
     from torch.utils.data import DataLoader
 
     from mteb.abstasks.task_metadata import TaskMetadata
@@ -166,7 +167,7 @@ class OpenAIModel(AbsEncoder):
             all_embeddings[mask] = no_empty_embeddings
         return all_embeddings
 
-    def _to_numpy(self, embedding_response) -> np.ndarray:
+    def _to_numpy(self, embedding_response) -> NDArray[np.floating]:
         return np.array([e.embedding for e in embedding_response.data])
 
 

mteb/models/model_implementations/pylate_models.py

@@ -53,7 +53,7 @@ class PylateSearchEncoder:
         hf_split: str,
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
-        num_proc: int,
+        num_proc: int | None,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -89,7 +89,7 @@ class PylateSearchEncoder:
         top_k: int,
         encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
-        num_proc: int,
+        num_proc: int | None,
     ) -> RetrievalOutputType:
         queries_dataloader = create_dataloader(
             queries,
@@ -150,7 +150,7 @@ class PylateSearchEncoder:
         hf_split: str,
         top_k: int,
         encode_kwargs: EncodeKwargs,
-        num_proc: int,
+        num_proc: int | None,
     ) -> dict[str, list[tuple[float, str]]]:
         from pylate import indexes, retrieve
 
@@ -216,7 +216,7 @@ class PylateSearchEncoder:
         hf_subset: str,
         hf_split: str,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> dict[str, list[tuple[float, str]]]:
         """Rerank with PyLate's rank.rerank using per-query candidates.
 

mteb/models/model_implementations/random_baseline.py

@@ -13,6 +13,7 @@ from mteb.similarity_functions import (
 )
 
 if TYPE_CHECKING:
+    from numpy.typing import NDArray
     from PIL import Image
     from torch.utils.data import DataLoader
 
@@ -20,7 +21,7 @@ if TYPE_CHECKING:
     from mteb.types._encoder_io import Array, BatchedInput, PromptType
 
 
-def _string_to_vector(text: str | None, size: int) -> np.ndarray:
+def _string_to_vector(text: str | None, size: int) -> NDArray[np.floating]:
     """Generate a deterministic random vector based on a string.
 
     Args:
@@ -39,7 +40,7 @@ def _string_to_vector(text: str | None, size: int) -> np.ndarray:
     return rng.random(size, dtype=np.float32)
 
 
-def _image_to_vector(image: Image.Image, size: int) -> np.ndarray:
+def _image_to_vector(image: Image.Image, size: int) -> NDArray[np.floating]:
     """Generate a deterministic random vector based on image content.
 
     Args:
@@ -80,7 +81,7 @@ _common_mock_metadata = dict(
 
 def _batch_to_embeddings(
     inputs: DataLoader[BatchedInput], embedding_dim: int
-) -> np.ndarray:
+) -> NDArray[np.floating]:
     """Convert batched text/image inputs into embeddings.
 
     Args:

mteb/models/model_implementations/seed_models.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import logging
 import time
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
@@ -14,6 +16,9 @@ from mteb.types import PromptType
 from .bge_models import bge_chinese_training_data
 from .nvidia_models import nvidia_training_datasets
 
+if TYPE_CHECKING:
+    from mteb.types import Array
+
 logger = logging.getLogger(__name__)
 
 
@@ -110,7 +115,7 @@ class SeedTextEmbeddingModel(AbsEncoder):
         prompt_type: PromptType | None = None,
         retries: int = 5,
         **kwargs: Any,
-    ) -> np.ndarray:
+    ) -> Array:
         trimmed_sentences = []
         for sentence in sentences:
             encoded_sentence = self._encoding.encode(sentence)