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

mteb/models/model_implementations/moka_models.py

@@ -97,6 +97,7 @@ m3e_base = ModelMeta(
     revision="764b537a0e50e5c7d64db883f2d2e051cbe3c64c",
     release_date="2023-06-06",  # first commit
     n_parameters=int(102 * 1e6),
+    n_embedding_parameters=16_226_304,
     memory_usage_mb=390,
     embed_dim=768,
     # They don't give a specific license but commercial use is not allowed
@@ -123,6 +124,7 @@ m3e_small = ModelMeta(
     revision="44c696631b2a8c200220aaaad5f987f096e986df",
     release_date="2023-06-02",  # first commit
     n_parameters=None,
+    n_embedding_parameters=10_817_536,
     memory_usage_mb=None,  # Can't be seen on HF page
     embed_dim=512,
     # They don't give a specific license but commercial use is not allowed
@@ -149,6 +151,7 @@ m3e_large = ModelMeta(
     revision="12900375086c37ba5d83d1e417b21dc7d1d1f388",
     release_date="2023-06-21",  # first commit
     n_parameters=None,
+    n_embedding_parameters=21_635_072,
     memory_usage_mb=None,  # Can't be seen on HF page
     embed_dim=768,
     # They don't give a specific license but commercial use is not allowed

mteb/models/model_implementations/nbailab.py

@@ -12,6 +12,7 @@ nb_sbert = ModelMeta(
     revision="b95656350a076aeafd2d23763660f80655408cc6",
     release_date="2022-11-23",
     n_parameters=1_780_000_000,
+    n_embedding_parameters=91_812_096,
     memory_usage_mb=678,
     embed_dim=4096,
     license="apache-2.0",
@@ -34,6 +35,7 @@ nb_bert_large = ModelMeta(
     revision="f9d0fc184adab4dc354d85e1854b7634540d7550",
     release_date="2021-04-29",
     n_parameters=355087360,
+    n_embedding_parameters=51_200_000,
     memory_usage_mb=1359,
     embed_dim=1024,
     license="cc-by-4.0",
@@ -56,6 +58,7 @@ nb_bert_base = ModelMeta(
     revision="9417c3f62a3adc99f17ff92bff446f35d011f994",
     release_date="2021-01-13",
     n_parameters=177853440,
+    n_embedding_parameters=91_812_096,
     memory_usage_mb=681,
     embed_dim=768,
     license="cc-by-4.0",

mteb/models/model_implementations/no_instruct_sentence_models.py

@@ -1,15 +1,22 @@
-from collections.abc import Generator
+from __future__ import annotations
+
 from itertools import islice
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
-from torch.utils.data import DataLoader
 
-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 collections.abc import Generator
+
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
 
 
 # https://docs.python.org/3/library/itertools.html#itertools.batched
@@ -103,6 +110,7 @@ no_instruct_small_v0 = ModelMeta(
     revision="b38747000553d8268915c95a55fc87e707c9aadd",
     release_date="2024-05-01",  # first commit
     n_parameters=33_400_000,
+    n_embedding_parameters=11_720_448,
     memory_usage_mb=127,
     max_tokens=512,
     embed_dim=384,

mteb/models/model_implementations/nomic_models.py

@@ -1,15 +1,21 @@
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
 from packaging.version import Version
-from torch.utils.data import DataLoader
 
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.model_meta import ModelMeta, ScoringFunction
 from mteb.models.sentence_transformer_wrapper import SentenceTransformerEncoderWrapper
-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
 
 logger = logging.getLogger(__name__)
 
@@ -209,6 +215,7 @@ nomic_embed_v1_5 = ModelMeta(
     release_date="2024-02-10",  # first commit
     citation=NOMIC_CITATION,
     n_parameters=137_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=522,
     max_tokens=8192,
     embed_dim=768,
@@ -243,6 +250,7 @@ nomic_embed_v1 = ModelMeta(
     revision="0759316f275aa0cb93a5b830973843ca66babcf5",
     release_date="2024-01-31",  # first commit
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=522,
     max_tokens=8192,
     embed_dim=768,
@@ -278,6 +286,7 @@ nomic_embed_v1_ablated = ModelMeta(
     revision="7d948905c5d5d3874fa55a925d68e49dbf411e5f",
     release_date="2024-01-15",  # first commit
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     max_tokens=8192,
     embed_dim=768,
@@ -306,6 +315,7 @@ nomic_embed_v1_unsupervised = ModelMeta(
     revision="b53d557b15ae63852847c222d336c1609eced93c",
     release_date="2024-01-15",  # first commit
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     max_tokens=8192,
     embed_dim=768,
@@ -334,6 +344,7 @@ nomic_modern_bert_embed = ModelMeta(
     revision="5960f1566fb7cb1adf1eb6e816639cf4646d9b12",
     release_date="2024-12-29",
     n_parameters=149_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=568,
     max_tokens=8192,
     embed_dim=768,
@@ -473,6 +484,7 @@ nomic_embed_text_v2_moe = ModelMeta(
     revision="1066b6599d099fbb93dfcb64f9c37a7c9e503e85",
     release_date="2025-02-07",
     n_parameters=475292928,
+    n_embedding_parameters=None,
     memory_usage_mb=1813,
     max_tokens=512,
     embed_dim=768,

mteb/models/model_implementations/nomic_models_vision.py

@@ -4,17 +4,18 @@ from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
-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
 
 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, PromptType
 
 NOMIC_EMBED_VISION_CITATION = """@article{nussbaum2024nomicembedvision,
       title={Nomic Embed Vision: Expanding the Latent Space},
@@ -174,6 +175,7 @@ nomic_embed_vision_v1_5 = ModelMeta(
     release_date="2024-06-08",
     modalities=["image", "text"],
     n_parameters=92_900_000,
+    n_embedding_parameters=None,
     memory_usage_mb=355,
     max_tokens=2048,
     embed_dim=768,

mteb/models/model_implementations/nvidia_llama_nemoretriever_colemb.py

@@ -1,14 +1,18 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 import torch
 from packaging.version import Version
 from torch.utils.data import DataLoader
 from transformers import __version__ as transformers_version
 
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta
-from mteb.types import Array, BatchedInput, PromptType
+
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
 
 LLAMA_NEMORETRIEVER_CITATION = """@misc{xu2025llamanemoretrievercolembedtopperforming,
       title={Llama Nemoretriever Colembed: Top-Performing Text-Image Retrieval Model},
@@ -158,6 +162,7 @@ llama_nemoretriever_colembed_1b_v1 = ModelMeta(
     release_date="2025-06-27",
     modalities=["image", "text"],
     n_parameters=2_418_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=4610,
     max_tokens=8192,
     embed_dim=2048,
@@ -185,6 +190,7 @@ llama_nemoretriever_colembed_3b_v1 = ModelMeta(
     release_date="2025-06-27",
     modalities=["image", "text"],
     n_parameters=4_407_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=8403,
     max_tokens=8192,
     embed_dim=3072,

mteb/models/model_implementations/nvidia_models.py

@@ -1,11 +1,11 @@
+from __future__ import annotations
+
 import logging
-from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
 from packaging.version import Version
-from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import AutoModel, AutoTokenizer
 from transformers import __version__ as transformers_version
@@ -16,7 +16,15 @@ from mteb.models import CrossEncoderWrapper
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.instruct_wrapper import InstructSentenceTransformerModel
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.utils.data import DataLoader
+
+    from mteb import TaskMetadata
+    from mteb.types import Array, BatchedInput
 
 logger = logging.getLogger(__name__)
 
@@ -196,6 +204,7 @@ NV_embed_v2 = ModelMeta(
     revision="7604d305b621f14095a1aa23d351674c2859553a",
     release_date="2024-09-09",  # initial commit of hf model.
     n_parameters=7_850_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=14975,
     embed_dim=4096,
     license="cc-by-nc-4.0",
@@ -227,6 +236,7 @@ NV_embed_v1 = ModelMeta(
     revision="570834afd5fef5bf3a3c2311a2b6e0a66f6f4f2c",
     release_date="2024-09-13",  # initial commit of hf model.
     n_parameters=7_850_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=14975,
     embed_dim=4096,
     license="cc-by-nc-4.0",
@@ -616,6 +626,7 @@ llama_embed_nemotron_8b = ModelMeta(
     revision="84a375593d27d3528beb4e104822515659e093b4",
     release_date="2025-10-23",
     n_parameters=7_504_924_672,
+    n_embedding_parameters=None,
     memory_usage_mb=28629,
     embed_dim=4096,
     license="https://huggingface.co/nvidia/llama-embed-nemotron-8b/blob/main/LICENSE",

mteb/models/model_implementations/octen_models.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(
@@ -208,6 +208,7 @@ Octen_Embedding_4B = ModelMeta(
     revision="6e188e3b072c3e3678b235ad84e6e97bcbb71e8f",
     release_date="2025-12-30",
     n_parameters=4021774336,
+    n_embedding_parameters=None,
     memory_usage_mb=7671,
     embed_dim=2560,
     max_tokens=32768,
@@ -238,6 +239,7 @@ Octen_Embedding_8B = ModelMeta(
     revision="f7db178d5a82fb841f606a6a67c423cead2fdbba",
     release_date="2025-12-23",
     n_parameters=7567295488,
+    n_embedding_parameters=None,
     memory_usage_mb=14433,
     embed_dim=4096,
     max_tokens=32768,

mteb/models/model_implementations/openai_models.py

@@ -1,15 +1,20 @@
+from __future__ import annotations
+
 import logging
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 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
+
+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__)
 
@@ -180,6 +185,7 @@ text_embedding_3_small = ModelMeta(
     embed_dim=1536,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://openai.com/index/new-embedding-models-and-api-updates/",
@@ -208,6 +214,7 @@ text_embedding_3_large = ModelMeta(
     framework=["API"],
     use_instructions=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     public_training_code=None,
     public_training_data=None,  # assumed
@@ -233,6 +240,7 @@ text_embedding_ada_002 = ModelMeta(
     framework=["API"],
     use_instructions=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     public_training_code=None,
     public_training_data=None,  # assumed
@@ -257,6 +265,7 @@ text_embedding_3_small_512 = ModelMeta(
     embed_dim=512,
     open_weights=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     license=None,
     reference="https://openai.com/index/new-embedding-models-and-api-updates/",
@@ -287,6 +296,7 @@ text_embedding_3_large_512 = ModelMeta(
     framework=["API"],
     use_instructions=False,
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     public_training_code=None,
     public_training_data=None,  # assumed

mteb/models/model_implementations/openclip_models.py

@@ -1,14 +1,19 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 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
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
 
 OPENCLIP_CITATION = """@inproceedings{cherti2023reproducible,
     title={Reproducible scaling laws for contrastive language-image learning},
@@ -128,6 +133,7 @@ CLIP_ViT_L_14_DataComp_XL_s13B_b90K = ModelMeta(
     release_date="2023-04-26",
     modalities=["image", "text"],
     n_parameters=428_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=1633,
     max_tokens=77,
     embed_dim=768,
@@ -154,6 +160,7 @@ CLIP_ViT_B_32_DataComp_XL_s13B_b90K = ModelMeta(
     release_date="2023-04-26",
     modalities=["image", "text"],
     n_parameters=151_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=576,
     max_tokens=77,
     embed_dim=512,
@@ -180,6 +187,7 @@ CLIP_ViT_B_16_DataComp_XL_s13B_b90K = ModelMeta(
     release_date="2023-04-26",
     modalities=["image", "text"],
     n_parameters=150_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=572,
     max_tokens=77,
     embed_dim=512,
@@ -206,6 +214,7 @@ CLIP_ViT_bigG_14_laion2B_39B_b160k = ModelMeta(
     release_date="2023-01-23",
     modalities=["image", "text"],
     n_parameters=2_540_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=9689,
     max_tokens=77,
     embed_dim=1280,
@@ -232,6 +241,7 @@ CLIP_ViT_g_14_laion2B_s34B_b88K = ModelMeta(
     release_date="2023-03-06",
     modalities=["image", "text"],
     n_parameters=1_367_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=5215,
     max_tokens=77,
     embed_dim=1024,
@@ -258,6 +268,7 @@ CLIP_ViT_H_14_laion2B_s32B_b79K = ModelMeta(
     release_date="2022-09-15",
     modalities=["image", "text"],
     n_parameters=986_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=3762,
     max_tokens=77,
     embed_dim=1024,
@@ -284,6 +295,7 @@ CLIP_ViT_L_14_laion2B_s32B_b82K = ModelMeta(
     release_date="2022-09-15",
     modalities=["image", "text"],
     n_parameters=428_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=1631,
     max_tokens=77,
     embed_dim=768,
@@ -310,6 +322,7 @@ CLIP_ViT_B_32_laion2B_s34B_b79K = ModelMeta(
     release_date="2022-09-15",
     modalities=["image", "text"],
     n_parameters=151_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=577,
     max_tokens=77,
     embed_dim=512,

mteb/models/model_implementations/opensearch_neural_sparse_models.py

@@ -1,12 +1,18 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.utils.data import DataLoader
 
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta
-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
 
 v2_training_data = {
     "MSMARCO",
@@ -134,6 +140,7 @@ opensearch_neural_sparse_encoding_doc_v3_gte = ModelMeta(
     revision="a8abaa916125ee512a7a8f4d706d07eb0128a8e6",
     release_date="2025-06-18",
     n_parameters=137_394_234,
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=549,
     embed_dim=30522,
     license="apache-2.0",
@@ -160,6 +167,7 @@ opensearch_neural_sparse_encoding_doc_v3_distill = ModelMeta(
     revision="babf71f3c48695e2e53a978208e8aba48335e3c0",
     release_date="2025-03-28",
     n_parameters=66_985_530,
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=267,
     embed_dim=30522,
     license="apache-2.0",
@@ -182,6 +190,7 @@ opensearch_neural_sparse_encoding_doc_v2_distill = ModelMeta(
     revision="8921a26c78b8559d6604eb1f5c0b74c079bee38f",
     release_date="2024-07-17",
     n_parameters=66_985_530,
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=267,
     embed_dim=30522,
     license="apache-2.0",
@@ -205,6 +214,7 @@ opensearch_neural_sparse_encoding_doc_v2_mini = ModelMeta(
     revision="4af867a426867dfdd744097531046f4289a32fdd",
     release_date="2024-07-18",
     n_parameters=22_744_506,
+    n_embedding_parameters=11_720_448,
     memory_usage_mb=86,
     embed_dim=30522,
     license="apache-2.0",
@@ -227,6 +237,7 @@ opensearch_neural_sparse_encoding_doc_v1 = ModelMeta(
     revision="98cdcbd72867c547f72f2b7b7bed9cdf9f09922d",
     release_date="2024-03-07",
     n_parameters=132_955_194,
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=507,
     embed_dim=30522,
     license="apache-2.0",

mteb/models/model_implementations/ops_moa_models.py

@@ -1,8 +1,13 @@
-import numpy as np
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta
 
+if TYPE_CHECKING:
+    from mteb.types import Array
+
 
 class OPSWrapper(AbsEncoder):
     def __init__(self, model_name: str, revision: str):
@@ -15,7 +20,7 @@ class OPSWrapper(AbsEncoder):
         )
         self.output_dim = 1536
 
-    def encode(self, sentences: list[str], **kwargs) -> np.ndarray:
+    def encode(self, sentences: list[str], **kwargs) -> Array:
         embeddings = self.model.encode(sentences, **kwargs)
         return embeddings[:, : self.output_dim]
 
@@ -28,6 +33,7 @@ ops_moa_conan_embedding = ModelMeta(
     languages=["zho-Hans"],
     loader=OPSWrapper,
     n_parameters=int(343 * 1e6),
+    n_embedding_parameters=21_635_072,
     memory_usage_mb=1308,
     max_tokens=512,
     embed_dim=1536,
@@ -60,6 +66,7 @@ ops_moa_yuan_embedding = ModelMeta(
     languages=["zho-Hans"],
     loader=OPSWrapper,
     n_parameters=int(343 * 1e6),
+    n_embedding_parameters=21_635_072,
     memory_usage_mb=1242,
     max_tokens=512,
     embed_dim=1536,

mteb/models/model_implementations/ordalietech_solon_embeddings_mini_beta_1_1.py

@@ -4,6 +4,7 @@ solon_embeddings_1_1 = ModelMeta(
     name="OrdalieTech/Solon-embeddings-mini-beta-1.1",
     languages=["fra-Latn"],
     n_parameters=210_000_000,
+    n_embedding_parameters=None,
     public_training_code=None,
     memory_usage_mb=808.0,
     open_weights=True,

mteb/models/model_implementations/pawan_models.py

@@ -20,6 +20,7 @@ pawan_embd_68m = ModelMeta(
     revision="32f295145802bdbd65699ad65fd27d2a5b69a909",
     release_date="2025-12-08",
     n_parameters=68_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=260,
     embed_dim=768,
     license="apache-2.0",

mteb/models/model_implementations/piccolo_models.py

@@ -12,6 +12,7 @@ piccolo_base_zh = ModelMeta(
     revision="47c0a63b8f667c3482e05b2fd45577bb19252196",
     release_date="2023-09-04",  # first commit
     n_parameters=None,
+    n_embedding_parameters=16_226_304,
     memory_usage_mb=None,  # can't see on model card
     embed_dim=768,
     license="mit",
@@ -37,6 +38,7 @@ piccolo_large_zh_v2 = ModelMeta(
     revision="05948c1d889355936bdf9db7d30df57dd78d25a3",
     release_date="2024-04-22",  # first commit
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=None,  # we don't know because they removed the model
     embed_dim=1024,
     license="not specified",

mteb/models/model_implementations/promptriever_models.py

@@ -1,15 +1,21 @@
+from __future__ import annotations
+
 import logging
-from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.utils.data import DataLoader
 
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.models.models_protocols import EncoderProtocol
-from mteb.types import Array, BatchedInput, PromptType
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import Array, BatchedInput, PromptType
 
 from .repllama_models import RepLLaMAModel, model_prompts
 
@@ -81,6 +87,7 @@ promptriever_llama2 = ModelMeta(
     revision="01c7f73d771dfac7d292323805ebc428287df4f9-30b14e3813c0fa45facfd01a594580c3fe5ecf23",  # base-peft revision
     release_date="2024-09-15",
     n_parameters=7_000_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=26703,
     max_tokens=4096,
     embed_dim=4096,
@@ -117,6 +124,7 @@ promptriever_llama3 = ModelMeta(
     },
     release_date="2024-09-15",
     n_parameters=8_000_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=30518,
     max_tokens=8192,
     embed_dim=4096,
@@ -146,6 +154,7 @@ promptriever_llama3_instruct = ModelMeta(
     revision="5206a32e0bd3067aef1ce90f5528ade7d866253f-8b677258615625122c2eb7329292b8c402612c21",  # base-peft revision
     release_date="2024-09-15",
     n_parameters=8_000_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=30518,
     max_tokens=8192,
     embed_dim=4096,
@@ -179,6 +188,7 @@ promptriever_mistral_v1 = ModelMeta(
     revision="7231864981174d9bee8c7687c24c8344414eae6b-876d63e49b6115ecb6839893a56298fadee7e8f5",  # base-peft revision
     release_date="2024-09-15",
     n_parameters=7_000_000_000,
+    n_embedding_parameters=131_072_000,
     memory_usage_mb=26703,
     training_datasets={
         # "samaya-ai/msmarco-w-instructions",