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

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",

mteb/models/model_implementations/pylate_models.py

@@ -1,30 +1,36 @@
+from __future__ import annotations
+
 import heapq
 import logging
 import shutil
 import tempfile
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.utils.data import DataLoader
 
 from mteb._create_dataloaders import (
     create_dataloader,
 )
 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,
-    CorpusDatasetType,
-    EncodeKwargs,
-    PromptType,
-    QueryDatasetType,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
-)
+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,
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
 
 logger = logging.getLogger(__name__)
 
@@ -47,6 +53,7 @@ class PylateSearchEncoder:
         hf_split: str,
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
+        num_proc: int,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -56,6 +63,7 @@ class PylateSearchEncoder:
             hf_split: Split of current task, allows to know some additional information about current split.
             hf_subset: Subset of current task. Similar to `hf_split` to get more information
             encode_kwargs: Additional arguments to pass to the encoder during indexing.
+            num_proc: Number of processes to use for indexing.
         """
         self.task_corpus = corpus
 
@@ -81,12 +89,14 @@ class PylateSearchEncoder:
         top_k: int,
         encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
+        num_proc: int,
     ) -> RetrievalOutputType:
         queries_dataloader = create_dataloader(
             queries,
             task_metadata,
             prompt_type=PromptType.query,
             batch_size=encode_kwargs.get("batch_size", 32),
+            num_proc=num_proc,
         )
 
         query_embeddings = self.encode(
@@ -110,6 +120,7 @@ class PylateSearchEncoder:
                 hf_subset=hf_subset,
                 hf_split=hf_split,
                 encode_kwargs=encode_kwargs,
+                num_proc=num_proc,
             )
         else:
             result_heaps = self._pylate_full_corpus_search(
@@ -120,6 +131,7 @@ class PylateSearchEncoder:
                 hf_subset=hf_subset,
                 hf_split=hf_split,
                 encode_kwargs=encode_kwargs,
+                num_proc=num_proc,
             )
 
         results = {qid: {} for qid in query_idx_to_id.values()}
@@ -138,6 +150,7 @@ class PylateSearchEncoder:
         hf_split: str,
         top_k: int,
         encode_kwargs: EncodeKwargs,
+        num_proc: int,
     ) -> dict[str, list[tuple[float, str]]]:
         from pylate import indexes, retrieve
 
@@ -164,6 +177,7 @@ class PylateSearchEncoder:
             task_metadata,
             prompt_type=PromptType.document,
             batch_size=encode_kwargs.get("batch_size", 32),
+            num_proc=num_proc,
         )
         documents_embeddings = self.encode(
             documents_loader,
@@ -202,6 +216,7 @@ class PylateSearchEncoder:
         hf_subset: str,
         hf_split: str,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> dict[str, list[tuple[float, str]]]:
         """Rerank with PyLate's rank.rerank using per-query candidates.
 
@@ -224,6 +239,7 @@ class PylateSearchEncoder:
                 task_metadata,
                 prompt_type=PromptType.document,
                 batch_size=encode_kwargs.get("batch_size", 32),
+                num_proc=num_proc,
             ),
             task_metadata=task_metadata,
             hf_split=hf_split,
@@ -346,6 +362,7 @@ colbert_v2 = ModelMeta(
     public_training_data=None,
     release_date="2024-09-21",
     n_parameters=int(110 * 1e6),
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=418,
     max_tokens=180,
     embed_dim=None,
@@ -402,6 +419,7 @@ jina_colbert_v2 = ModelMeta(
     public_training_data=None,
     release_date="2024-08-16",
     n_parameters=int(559 * 1e6),
+    n_embedding_parameters=None,
     memory_usage_mb=1067,
     max_tokens=8192,
     embed_dim=None,
@@ -458,6 +476,7 @@ lightonai__gte_moderncolbert_v1 = ModelMeta(
     public_training_data="https://huggingface.co/datasets/lightonai/ms-marco-en-bge-gemma",
     release_date="2025-04-30",
     n_parameters=int(149 * 1e6),
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     max_tokens=8192,
     embed_dim=None,

mteb/models/model_implementations/qodo_models.py

@@ -36,6 +36,7 @@ Qodo_Embed_1_1_5B = ModelMeta(
     revision="84bbef079b32e8823ec226d4e9e92902706b0eb6",
     release_date="2025-02-19",
     n_parameters=1_780_000_000,
+    n_embedding_parameters=232_928_256,
     memory_usage_mb=6776,
     embed_dim=1536,
     license="https://huggingface.co/Qodo/Qodo-Embed-1-1.5B/blob/main/LICENSE",
@@ -59,6 +60,7 @@ Qodo_Embed_1_7B = ModelMeta(
     revision="f9edd9bf7f687c0e832424058e265120f603cd81",
     release_date="2025-02-24",
     n_parameters=7_613_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=29040,
     embed_dim=3584,
     license="https://huggingface.co/Qodo/Qodo-Embed-1-1.5B/blob/main/LICENSE",

mteb/models/model_implementations/qtack_models.py

@@ -31,6 +31,7 @@ mini_gte = ModelMeta(
     revision="7fbe6f9b4cc42615e0747299f837ad7769025492",
     release_date="2025-01-28",
     n_parameters=int(66.3 * 1e6),
+    n_embedding_parameters=23_440_896,
     memory_usage_mb=253,
     embed_dim=768,
     license="apache-2.0",

mteb/models/model_implementations/qwen3_models.py

@@ -1,6 +1,13 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from mteb.models.instruct_wrapper import InstructSentenceTransformerModel
 from mteb.models.model_meta import ModelMeta
-from mteb.models.models_protocols import EncoderProtocol, PromptType
+from mteb.types import PromptType
+
+if TYPE_CHECKING:
+    from mteb.models.models_protocols import EncoderProtocol
 
 
 def instruction_template(
@@ -140,6 +147,7 @@ Qwen3_Embedding_0B6 = ModelMeta(
     revision="b22da495047858cce924d27d76261e96be6febc0",  # Commit of @tomaarsen
     release_date="2025-06-05",
     n_parameters=595776512,
+    n_embedding_parameters=None,
     memory_usage_mb=1136,
     embed_dim=1024,
     max_tokens=32768,
@@ -163,6 +171,7 @@ Qwen3_Embedding_4B = ModelMeta(
     revision="636cd9bf47d976946cdbb2b0c3ca0cb2f8eea5ff",  # Commit of @tomaarsen
     release_date="2025-06-05",
     n_parameters=4021774336,
+    n_embedding_parameters=None,
     memory_usage_mb=7671,
     embed_dim=2560,
     max_tokens=32768,
@@ -186,6 +195,7 @@ Qwen3_Embedding_8B = ModelMeta(
     revision="4e423935c619ae4df87b646a3ce949610c66241c",  # Commit of @tomaarsen
     release_date="2025-06-05",
     n_parameters=7567295488,
+    n_embedding_parameters=None,
     memory_usage_mb=14433,
     embed_dim=4096,
     max_tokens=32768,

mteb/models/model_implementations/qzhou_models.py

@@ -64,6 +64,7 @@ QZhou_Embedding = ModelMeta(
     revision="f1e6c03ee3882e7b9fa5cec91217715272e433b8",
     release_date="2025-08-24",
     n_parameters=7_070_619_136,
+    n_embedding_parameters=None,
     memory_usage_mb=14436,
     embed_dim=3584,
     license="apache-2.0",
@@ -98,6 +99,7 @@ QZhou_Embedding_Zh = ModelMeta(
     revision="0321ccb126413d1e49c5ce908e802b63d35f18e2",
     release_date="2025-09-28",
     n_parameters=7_575_747_328,
+    n_embedding_parameters=None,
     memory_usage_mb=29431,
     embed_dim=1792,
     license="apache-2.0",

mteb/models/model_implementations/random_baseline.py

@@ -5,18 +5,19 @@ from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import torch
-from torch.utils.data import DataLoader
 
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.model_meta import ModelMeta
 from mteb.similarity_functions import (
     select_pairwise_similarity,
     select_similarity,
 )
-from mteb.types._encoder_io 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._encoder_io import Array, BatchedInput, PromptType
 
 
 def _string_to_vector(text: str | None, size: int) -> np.ndarray:

mteb/models/model_implementations/rasgaard_models.py

@@ -12,6 +12,7 @@ potion_base_8m = ModelMeta(
     revision="387897cfb09992e6d45ea9cd7b28b9fcf119e23a",
     release_date="2025-10-08",
     n_parameters=22893312,
+    n_embedding_parameters=22893312,
     memory_usage_mb=87,
     max_tokens=np.inf,
     embed_dim=256,

mteb/models/model_implementations/reasonir_model.py

@@ -36,12 +36,76 @@ REASONIR_TRAINING_DATA = {
     "DuRetrieval",
     "QuoraRetrieval",
 }
+_prompts_dict = {
+    "BrightBiologyRetrieval": {
+        "query": "Given a Biology post, retrieve relevant passages that help answer the post"
+    },
+    "BrightEarthScienceRetrieval": {
+        "query": "Given a Earth Science post, retrieve relevant passages that help answer the post"
+    },
+    "BrightEconomicsRetrieval": {
+        "query": "Given a Economics post, retrieve relevant passages that help answer the post"
+    },
+    "BrightPsychologyRetrieval": {
+        "query": "Given a Psychology post, retrieve relevant passages that help answer the post"
+    },
+    "BrightRoboticsRetrieval": {
+        "query": "Given a Robotics post, retrieve relevant passages that help answer the post"
+    },
+    "BrightStackoverflowRetrieval": {
+        "query": "Given a Stackoverflow post, retrieve relevant passages that help answer the post"
+    },
+    "BrightSustainableLivingRetrieval": {
+        "query": "Given a Sustainable Living post, retrieve relevant passages that help answer the post"
+    },
+    "BrightPonyRetrieval": {
+        "query": "Given a Pony question, retrieve relevant passages that help answer the question"
+    },
+    "BrightLeetcodeRetrieval": {
+        "query": "Given a coding problem, retrieve relevant examples that help answer the problem",
+    },
+    "BrightAopsRetrieval": {
+        "query": "Given a Math problem, retrieve relevant examples that help answer the problem"
+    },
+    "BrightTheoremQATheoremsRetrieval": {
+        "query": "Given a Math problem, retrieve relevant theorems that help answer the problem",
+    },
+    "BrightTheoremQAQuestionsRetrieval": {
+        "query": "Given a Math problem, retrieve relevant examples that help answer the problem",
+    },
+    "BrightBiologyLongRetrieval": {
+        "query": "Given a Biology post, retrieve relevant documents that help answer the post"
+    },
+    "BrightEarthScienceLongRetrieval": {
+        "query": "Given a Earth Science post, retrieve relevant documents that help answer the post"
+    },
+    "BrightEconomicsLongRetrieval": {
+        "query": "Given a Economics post, retrieve relevant documents that help answer the post"
+    },
+    "BrightPsychologyLongRetrieval": {
+        "query": "Given a Psychology post, retrieve relevant documents that help answer the post"
+    },
+    "BrightRoboticsLongRetrieval": {
+        "query": "Given a Robotics post, retrieve relevant documents that help answer the post"
+    },
+    "BrightStackoverflowLongRetrieval": {
+        "query": "Given a Stackoverflow post, retrieve relevant documents that help answer the post"
+    },
+    "BrightSustainableLivingLongRetrieval": {
+        "query": "Given a Sustainable Living post, retrieve relevant documents that help answer the post"
+    },
+    "BrightPonyLongRetrieval": {
+        "query": "Given a Pony question, retrieve relevant documents that help answer the question"
+    },
+}
+
 
 ReasonIR_8B = ModelMeta(
     loader=InstructSentenceTransformerModel,
     loader_kwargs=dict(
         instruction_template=instruction_template,
         trust_remote_code=True,
+        prompts_dict=_prompts_dict,
     ),
     name="ReasonIR/ReasonIR-8B",
     model_type=["dense"],
@@ -50,6 +114,7 @@ ReasonIR_8B = ModelMeta(
     revision="c3d0690370ff4a8c3d3882d8dfa85c43650034fa",
     release_date="2025-04-29",
     n_parameters=7_500_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=None,
     embed_dim=4096,
     license="cc-by-nc-4.0",

mteb/models/model_implementations/repllama_models.py

@@ -1,22 +1,29 @@
+from __future__ import annotations
+
 import logging
-from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 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.models.models_protocols import EncoderProtocol
-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.abstasks.task_metadata import TaskMetadata
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import Array, BatchedInput
 
 logger = logging.getLogger(__name__)
 
@@ -172,6 +179,7 @@ repllama_llama2_original = ModelMeta(
         "mMARCO-NL",  # translation not trained on
     },
     n_parameters=7_000_000,
+    n_embedding_parameters=131_072_000,
     memory_usage_mb=27,
     max_tokens=4096,
     embed_dim=4096,
@@ -201,6 +209,7 @@ repllama_llama2_reproduced = ModelMeta(
     revision="01c7f73d771dfac7d292323805ebc428287df4f9-ad5c1d0938a1e02954bcafb4d811ba2f34052e71",  # base-peft revision
     release_date="2024-09-15",
     n_parameters=7_000_000,
+    n_embedding_parameters=None,
     memory_usage_mb=27,
     max_tokens=4096,
     embed_dim=4096,

mteb/models/model_implementations/rerankers_custom.py

@@ -1,16 +1,22 @@
+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 mteb._requires_package import requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.model_meta import ModelMeta
-from mteb.types import Array, BatchedInput, PromptType
 
 from .bge_models import bge_m3_training_data
 
+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__)
 
 
@@ -225,6 +231,7 @@ monobert_large = ModelMeta(
     revision="0a97706f3827389da43b83348d5d18c9d53876fa",
     release_date="2020-05-28",
     n_parameters=None,
+    n_embedding_parameters=31_254_528,
     memory_usage_mb=None,
     max_tokens=None,
     embed_dim=None,
@@ -250,6 +257,7 @@ jina_reranker_multilingual = ModelMeta(
     revision="126747772a932960028d9f4dc93bd5d9c4869be4",
     release_date="2024-09-26",
     n_parameters=None,
+    n_embedding_parameters=None,
     memory_usage_mb=531,
     max_tokens=None,
     embed_dim=None,
@@ -313,6 +321,7 @@ bge_reranker_v2_m3 = ModelMeta(
     revision="953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e",
     release_date="2024-06-24",
     n_parameters=None,
+    n_embedding_parameters=256_002_048,
     memory_usage_mb=2166,
     max_tokens=None,
     embed_dim=None,