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

mteb/models/model_implementations/vi_vn_models.py

@@ -16,6 +16,7 @@ greennode_embedding_large_vn_v1 = ModelMeta(
     loader=sentence_transformers_loader,
     open_weights=True,
     n_parameters=568_000_000,
+    n_embedding_parameters=256_002_048,
     memory_usage_mb=2167,
     embed_dim=1024,
     license="cc-by-4.0",
@@ -41,6 +42,7 @@ greennode_embedding_large_vn_mixed_v1 = ModelMeta(
     loader=sentence_transformers_loader,
     open_weights=True,
     n_parameters=568_000_000,
+    n_embedding_parameters=256_002_048,
     memory_usage_mb=2167,
     embed_dim=1024,
     license="cc-by-4.0",
@@ -66,6 +68,7 @@ aiteamvn_vietnamese_embeddings = ModelMeta(
     loader=sentence_transformers_loader,
     open_weights=True,
     n_parameters=568_000_000,
+    n_embedding_parameters=256_002_048,
     memory_usage_mb=2166,
     embed_dim=1024,
     license="cc-by-4.0",
@@ -98,6 +101,7 @@ hiieu_halong_embedding = ModelMeta(
     use_instructions=False,
     open_weights=True,
     n_parameters=278_000_000,
+    n_embedding_parameters=192_001_536,
     memory_usage_mb=1061,
     embed_dim=768,
     license="apache-2.0",
@@ -129,6 +133,7 @@ sup_simcse_vietnamese_phobert_base_ = ModelMeta(
     use_instructions=False,
     open_weights=True,
     n_parameters=135_000_000,
+    n_embedding_parameters=49_152_768,
     memory_usage_mb=517,
     max_tokens=256,
     embed_dim=768,
@@ -167,6 +172,7 @@ bkai_foundation_models_vietnamese_bi_encoder = ModelMeta(
     use_instructions=False,
     open_weights=True,
     n_parameters=135_000_000,
+    n_embedding_parameters=49_152_768,
     memory_usage_mb=515,
     max_tokens=256,
     embed_dim=768,

mteb/models/model_implementations/vista_models.py

@@ -1,14 +1,19 @@
-from typing import Any, Literal
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
 
 import torch
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import requires_image_dependencies
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
 
 VISTA_CITATION = """@article{zhou2024vista,
   title={VISTA: Visualized Text Embedding For Universal Multi-Modal Retrieval},
@@ -253,6 +258,7 @@ visualized_bge_base = ModelMeta(
     release_date="2024-06-06",
     modalities=["image", "text"],
     n_parameters=196_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=1631,
     max_tokens=512,
     embed_dim=768,
@@ -281,6 +287,7 @@ visualized_bge_m3 = ModelMeta(
     release_date="2024-06-06",
     modalities=["image", "text"],
     n_parameters=872_909_505,
+    n_embedding_parameters=None,
     memory_usage_mb=4263,
     max_tokens=8192,
     embed_dim=1024,

mteb/models/model_implementations/vlm2vec_models.py

@@ -1,8 +1,9 @@
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import (
@@ -10,10 +11,14 @@ from mteb._requires_package import (
     requires_package,
     suggest_package,
 )
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
 
 logger = logging.getLogger(__name__)
 
@@ -275,6 +280,7 @@ vlm2vec_lora = ModelMeta(
     release_date="2024-10-08",
     modalities=["image", "text"],
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     max_tokens=131072,
     embed_dim=3072,
@@ -299,6 +305,7 @@ vlm2vec_full = ModelMeta(
     release_date="2024-10-08",
     modalities=["image", "text"],
     n_parameters=4_150_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=7909,
     max_tokens=131072,
     embed_dim=3072,

mteb/models/model_implementations/voyage_models.py

@@ -1,16 +1,22 @@
+from __future__ import annotations
+
 import time
 from functools import wraps
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
 
 VOYAGE_TRAINING_DATA = set(
     # Self-reported (message from VoyageAI member)
@@ -302,6 +308,7 @@ voyage_3_large = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2025/01/07/voyage-3-large/",
@@ -330,6 +337,7 @@ voyage_3_5 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2025/05/20/voyage-3-5/",
@@ -357,6 +365,7 @@ voyage_3_5_int8 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2025/05/20/voyage-3-5/",
@@ -384,6 +393,7 @@ voyage_3_5_binary = ModelMeta(
     embed_dim=1024,  # Same as original after unpacking from bits
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2025/05/20/voyage-3-5/",
@@ -411,6 +421,7 @@ voyage_large_2_instruct = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/05/05/voyage-large-2-instruct-instruction-tuned-and-rank-1-on-mteb/",
@@ -437,6 +448,7 @@ voyage_finance_2 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/06/03/domain-specific-embeddings-finance-edition-voyage-finance-2/",
@@ -463,6 +475,7 @@ voyage_law_2 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/04/15/domain-specific-embeddings-and-retrieval-legal-edition-voyage-law-2/",
@@ -489,6 +502,7 @@ voyage_code_2 = ModelMeta(
     embed_dim=1536,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/01/23/voyage-code-2-elevate-your-code-retrieval/",
@@ -515,6 +529,7 @@ voyage_code_3 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/12/04/voyage-code-3/",
@@ -542,6 +557,7 @@ voyage_large_2 = ModelMeta(
     embed_dim=1536,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2023/10/29/voyage-embeddings/",
@@ -568,6 +584,7 @@ voyage_2 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2023/10/29/voyage-embeddings/",
@@ -593,6 +610,7 @@ voyage_multilingual_2 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/06/10/voyage-multilingual-2-multilingual-embedding-model/",
@@ -619,6 +637,7 @@ voyage_3 = ModelMeta(
     embed_dim=1024,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/09/18/voyage-3/",
@@ -645,6 +664,7 @@ voyage_3_lite = ModelMeta(
     embed_dim=512,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://blog.voyageai.com/2024/09/18/voyage-3/",
@@ -673,6 +693,7 @@ voyage_3_exp = ModelMeta(
     open_weights=False,
     # from their card https://huggingface.co/voyageai/voyage-3-m-exp#model-information
     n_parameters=int(6918 * 1e6),
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://huggingface.co/voyageai/voyage-3-m-exp",

mteb/models/model_implementations/voyage_v.py

@@ -4,17 +4,19 @@ import logging
 from typing import TYPE_CHECKING, Any, Literal
 
 import torch
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import requires_image_dependencies, requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
 
 if TYPE_CHECKING:
     from PIL import Image
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +29,8 @@ def _downsample_image(
     Returns:
         The downsampled image.
     """
+    from PIL.Image import Resampling
+
     width, height = image.size
     pixels = width * height
 
@@ -42,15 +46,15 @@ def _downsample_image(
         logger.info(
             f"Downsampling image from {width}x{height} to {new_width}x{new_height}"
         )
-        return image.resize(new_size, Image.LANCZOS)
+        return image.resize(new_size, Resampling.LANCZOS)
     if width > height:
         if width > 10000:
             logger.error("Processing extremely wide images.")
-            return image.resize((10000, height), Image.LANCZOS)
+            return image.resize((10000, height), Resampling.LANCZOS)
     else:
         if height > 10000:
             logger.error("Processing extremely high images.")
-            return image.resize((width, 10000), Image.LANCZOS)
+            return image.resize((width, 10000), Resampling.LANCZOS)
     return image
 
 
@@ -211,6 +215,7 @@ voyage_v = ModelMeta(
     revision="1",
     release_date="2024-11-10",
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     max_tokens=32768,
     embed_dim=1024,

mteb/models/model_implementations/xyz_models.py

@@ -31,6 +31,7 @@ xyz_embedding = ModelMeta(
     revision="4004120220b99baea764a1d3508427248ac3bccf",
     release_date="2024-09-13",
     n_parameters=326000000,
+    n_embedding_parameters=21_635_072,
     memory_usage_mb=1242,
     max_tokens=512,
     embed_dim=768,

mteb/models/model_implementations/youtu_models.py

@@ -121,6 +121,7 @@ Youtu_Embedding_V1 = ModelMeta(
     release_date="2025-09-28",
     open_weights=True,
     n_parameters=2672957440,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     embed_dim=2048,
     license="apache-2.0",

mteb/models/model_implementations/yuan_models.py

@@ -20,6 +20,7 @@ yuan_embedding_2_zh = ModelMeta(
     revision="b5ebcace6f4fc6e5a4d1852557eb2dc2d1040cee",
     release_date="2025-11-24",
     n_parameters=326000000,
+    n_embedding_parameters=21_635_072,
     memory_usage_mb=1242,
     embed_dim=1792,
     license="apache-2.0",

mteb/models/model_implementations/yuan_models_en.py

@@ -1,6 +1,6 @@
 from mteb.models.instruct_wrapper import InstructSentenceTransformerModel
 from mteb.models.model_meta import ModelMeta
-from mteb.models.models_protocols import PromptType
+from mteb.types import PromptType
 
 
 def instruction_template(
@@ -43,6 +43,7 @@ yuan_embedding_2_en = ModelMeta(
     revision="b2fd15da3bcae3473c8529593825c15068f09fce",
     release_date="2025-11-27",
     n_parameters=595776512,
+    n_embedding_parameters=None,
     memory_usage_mb=2272,
     embed_dim=1024,
     max_tokens=2048,

mteb/models/model_meta.py

@@ -3,17 +3,16 @@ from __future__ import annotations
 import json
 import logging
 import warnings
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from dataclasses import field
 from enum import Enum
 from functools import partial
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, cast
 
+import numpy as np
 from huggingface_hub import (
-    GitCommitInfo,
     ModelCard,
-    ModelCardData,
     get_safetensors_metadata,
     hf_hub_download,
     list_repo_commits,
@@ -29,18 +28,27 @@ from huggingface_hub.errors import (
     SafetensorsParsingError,
 )
 from pydantic import BaseModel, ConfigDict, field_validator, model_validator
+from sentence_transformers.models import Transformer
+from torch import nn
 from transformers import AutoConfig
-from typing_extensions import Self
 
 from mteb._helpful_enum import HelpfulStrEnum
 from mteb.languages import check_language_code
-from mteb.models.models_protocols import EncoderProtocol, MTEBModels
+from mteb.models.models_protocols import MTEBModels
 from mteb.types import ISOLanguageScript, Licenses, Modalities, StrDate, StrURL
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from huggingface_hub import (
+        GitCommitInfo,
+        ModelCardData,
+    )
     from sentence_transformers import CrossEncoder, SentenceTransformer
+    from typing_extensions import Self
 
     from mteb.abstasks import AbsTask
+    from mteb.models.models_protocols import EncoderProtocol
 
 
 logger = logging.getLogger(__name__)
@@ -94,8 +102,9 @@ class ModelMeta(BaseModel):
         loader: The function that loads the model. If None it assumes that the model is not implemented.
         loader_kwargs: The keyword arguments to pass to the loader function.
         name: The name of the model, ideally the name on huggingface. It should be in the format "organization/model_name".
-        n_parameters: The number of parameters in the model, e.g. 7_000_000 for a 7M parameter model. Can be None if the number of parameters is not known (e.g. for proprietary models) or
-            if the loader returns a SentenceTransformer model from which it can be derived.
+        n_parameters: The total number of parameters in the model, e.g. `7_000_000` for a 7M parameter model. Can be none in case the number of parameters is unknown.
+        n_embedding_parameters: The number of parameters used for the embedding layer. Can be None if the number of embedding parameters is not known (e.g. for proprietary models).
+        n_active_parameters_override: The number of active parameters used bu model. Should be used **only** for Mixture of Experts models.
         memory_usage_mb: The memory usage of the model in MB. Can be None if the memory usage is not known (e.g. for proprietary models). To calculate it use the `calculate_memory_usage_mb` method.
         max_tokens: The maximum number of tokens the model can handle. Can be None if the maximum number of tokens is not known (e.g. for proprietary
             models).
@@ -134,6 +143,8 @@ class ModelMeta(BaseModel):
     release_date: StrDate | None
     languages: list[ISOLanguageScript] | None
     n_parameters: int | None
+    n_active_parameters_override: int | None = None
+    n_embedding_parameters: int | None = None
     memory_usage_mb: float | None
     max_tokens: float | None
     embed_dim: int | None
@@ -192,6 +203,16 @@ class ModelMeta(BaseModel):
         """
         return "cross-encoder" in self.model_type
 
+    @property
+    def n_active_parameters(self):
+        """Number of active parameters. Assumed to be `n_parameters - n_embedding_parameters`. Can be overwritten using `n_active_parameters_override` e.g. for MoE models."""
+        if self.n_active_parameters_override is not None:
+            return self.n_active_parameters_override
+
+        if self.n_parameters is not None and self.n_embedding_parameters is not None:
+            return self.n_parameters - self.n_embedding_parameters
+        return None
+
     @field_validator("similarity_fn_name", mode="before")
     @classmethod
     def _validate_similarity_fn_name(cls, value: str) -> ScoringFunction | None:
@@ -384,6 +405,14 @@ class ModelMeta(BaseModel):
             else model.model_card_data.base_model
         )
         meta = cls._from_hub(name, revision, compute_metadata)
+        try:
+            first = model[0]
+
+            if isinstance(first, Transformer):
+                emb = first.auto_model.get_input_embeddings()
+                meta.n_embedding_parameters = int(np.prod(emb.weight.shape))
+        except Exception as e:
+            logger.warning(f"Could not calculate embedding parameters for {name}: {e}")
         meta.revision = model.model_card_data.base_model_revision or meta.revision
         meta.max_tokens = model.max_seq_length
         meta.embed_dim = model.get_sentence_embedding_dimension()
@@ -455,6 +484,15 @@ class ModelMeta(BaseModel):
         from mteb.models import CrossEncoderWrapper
 
         meta = cls._from_hub(model.model.name_or_path, revision, compute_metadata)
+        try:
+            emb = model.model.get_input_embeddings()
+
+            if isinstance(emb, nn.Embedding):
+                meta.n_embedding_parameters = int(np.prod(emb.weight.shape))
+        except Exception as e:
+            logger.warning(
+                f"Could not calculate embedding parameters for {model.model.name_or_path}: {e}"
+            )
         meta.revision = model.config._commit_hash or meta.revision
         meta.loader = CrossEncoderWrapper
         meta.embed_dim = None
@@ -479,7 +517,7 @@ class ModelMeta(BaseModel):
         if isinstance(tasks[0], str):
             benchmark_datasets = set(tasks)
         else:
-            tasks = cast(Sequence["AbsTask"], tasks)
+            tasks = cast("Sequence[AbsTask]", tasks)
             benchmark_datasets = set()
             for task in tasks:
                 benchmark_datasets.add(task.metadata.name)
@@ -534,7 +572,7 @@ class ModelMeta(BaseModel):
         if isinstance(tasks[0], str):
             benchmark_datasets = set(tasks)
         else:
-            tasks = cast(Sequence["AbsTask"], tasks)
+            tasks = cast("Sequence[AbsTask]", tasks)
             benchmark_datasets = {task.metadata.name for task in tasks}
         overlap = training_datasets & benchmark_datasets
         perc_overlap = 100 * (len(overlap) / len(benchmark_datasets))

mteb/models/models_protocols.py

@@ -1,22 +1,23 @@
-from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+from __future__ import annotations
 
-from torch.utils.data import DataLoader
-from typing_extensions import Unpack
-
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.types import (
-    Array,
-    BatchedInput,
-    CorpusDatasetType,
-    EncodeKwargs,
-    PromptType,
-    QueryDatasetType,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
-)
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+    from typing_extensions import Unpack
+
+    from mteb.abstasks.task_metadata import TaskMetadata
     from mteb.models.model_meta import ModelMeta
+    from mteb.types import (
+        Array,
+        BatchedInput,
+        CorpusDatasetType,
+        EncodeKwargs,
+        PromptType,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
 
 
 @runtime_checkable
@@ -72,7 +73,7 @@ class SearchProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -177,7 +178,7 @@ class EncoderProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -236,7 +237,7 @@ class CrossEncoderProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 

mteb/models/search_encoder_index/search_backend_protocol.py

@@ -1,7 +1,11 @@
-from collections.abc import Callable
-from typing import Protocol
+from __future__ import annotations
 
-from mteb.types import Array, TopRankedDocumentsType
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from mteb.types import Array, TopRankedDocumentsType
 
 
 class IndexEncoderSearchProtocol(Protocol):

mteb/models/search_encoder_index/search_indexes/faiss_search_index.py

@@ -1,14 +1,23 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
 
 from mteb._requires_package import requires_package
 from mteb.models.model_meta import ScoringFunction
-from mteb.models.models_protocols import EncoderProtocol
-from mteb.types import Array, TopRankedDocumentsType
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import faiss
+
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import Array, TopRankedDocumentsType
+
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +42,6 @@ class FaissSearchIndex:
             install_instruction="pip install mteb[faiss-cpu]",
         )
 
-        import faiss
         from faiss import IndexFlatIP, IndexFlatL2
 
         # https://github.com/facebookresearch/faiss/wiki/Faiss-indexes

mteb/models/search_wrappers.py

@@ -1,28 +1,35 @@
+from __future__ import annotations
+
 import heapq
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 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.types import (
-    Array,
-    BatchedInput,
-    CorpusDatasetType,
-    EncodeKwargs,
     PromptType,
-    QueryDatasetType,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
 )
 
-from .models_protocols import CrossEncoderProtocol, EncoderProtocol
-from .search_encoder_index.search_backend_protocol import IndexEncoderSearchProtocol
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import (
+        Array,
+        BatchedInput,
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
+    from .models_protocols import CrossEncoderProtocol, EncoderProtocol
+    from .search_encoder_index.search_backend_protocol import IndexEncoderSearchProtocol
 
 logger = logging.getLogger(__name__)