mteb 2.6.7__py3-none-any.whl → 2.6.9__py3-none-any.whl

mteb/benchmarks/benchmarks/benchmarks.py

@@ -2728,3 +2728,27 @@ JMTEB_LITE_V1 = Benchmark(
 """,
     contacts=["lsz05"],
 )
+
+KOVIDORE_V2 = Benchmark(
+    name="KoViDoRe(v2)",
+    display_name="KoViDoRe v2",
+    tasks=get_tasks(
+        tasks=[
+            "KoVidore2CybersecurityRetrieval",
+            "KoVidore2EconomicRetrieval",
+            "KoVidore2EnergyRetrieval",
+            "KoVidore2HrRetrieval",
+        ]
+    ),
+    description="KoViDoRe v2 sets a new industry gold standard for multi-modal, enterprise document visual retrieval evaluation. It addresses a critical challenge in production RAG systems: retrieving accurate information from complex, visually-rich documents.",
+    reference="https://github.com/whybe-choi/kovidore-data-generator",
+    citation=r"""
+@misc{choi2026kovidorev2,
+  author = {Yongbin Choi},
+  note = {A benchmark for evaluating Korean vision document retrieval with multi-page reasoning queries in practical domains},
+  title = {KoViDoRe v2: a comprehensive evaluation of vision document retrieval for enterprise use-cases},
+  url = {https://github.com/whybe-choi/kovidore-data-generator},
+  year = {2026},
+}
+""",
+)

mteb/cli/build_cli.py

@@ -13,6 +13,7 @@ from mteb.cache import ResultCache
 from mteb.cli._display_tasks import _display_benchmarks, _display_tasks
 from mteb.cli.generate_model_card import generate_model_card
 from mteb.evaluate import OverwriteStrategy
+from mteb.types._encoder_io import EncodeKwargs
 
 logger = logging.getLogger(__name__)
 
@@ -64,7 +65,7 @@ def run(args: argparse.Namespace) -> None:
             eval_splits=args.eval_splits,
         )
 
-    encode_kwargs = {}
+    encode_kwargs: EncodeKwargs = {}
     if args.batch_size is not None:
         encode_kwargs["batch_size"] = args.batch_size
 

mteb/deprecated_evaluator.py

@@ -28,7 +28,7 @@ from mteb.models import (
     SentenceTransformerEncoderWrapper,
 )
 from mteb.results import TaskResult
-from mteb.types import ScoresDict
+from mteb.types import EncodeKwargs, ScoresDict
 
 if sys.version_info >= (3, 13):
     from warnings import deprecated
@@ -174,7 +174,7 @@ class MTEB:
         split: str,
         subsets_to_run: list[str] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         **kwargs: Any,
     ):
         tick = time()
@@ -263,7 +263,7 @@ class MTEB:
         overwrite_results: bool = False,
         raise_error: bool = True,
         co2_tracker: bool = False,
-        encode_kwargs: dict[str, Any] | None = None,
+        encode_kwargs: EncodeKwargs | None = None,
         **kwargs,
     ) -> list[TaskResult]:
         """Run the evaluation pipeline on the selected tasks.

mteb/descriptive_stats/Image/DocumentUnderstanding/KoVidore2CybersecurityRetrieval.json

@@ -0,0 +1,32 @@
+{
+    "test": {
+        "num_samples": 1299,
+        "number_of_characters": 9254,
+        "documents_text_statistics": null,
+        "documents_image_statistics": {
+            "min_image_width": 2245,
+            "average_image_width": 2370.324347826087,
+            "max_image_width": 3508,
+            "min_image_height": 2481,
+            "average_image_height": 3289.8060869565215,
+            "max_image_height": 3580,
+            "unique_images": 1132
+        },
+        "queries_text_statistics": {
+            "total_text_length": 9254,
+            "min_text_length": 15,
+            "average_text_length": 62.10738255033557,
+            "max_text_length": 108,
+            "unique_texts": 149
+        },
+        "queries_image_statistics": null,
+        "relevant_docs_statistics": {
+            "num_relevant_docs": 409,
+            "min_relevant_docs_per_query": 1,
+            "average_relevant_docs_per_query": 2.7449664429530203,
+            "max_relevant_docs_per_query": 7,
+            "unique_relevant_docs": 316
+        },
+        "top_ranked_statistics": null
+    }
+}

mteb/descriptive_stats/Image/DocumentUnderstanding/KoVidore2EconomicRetrieval.json

@@ -0,0 +1,32 @@
+{
+    "test": {
+        "num_samples": 1640,
+        "number_of_characters": 8331,
+        "documents_text_statistics": null,
+        "documents_image_statistics": {
+            "min_image_width": 2313,
+            "average_image_width": 2347.5321597833445,
+            "max_image_width": 2481,
+            "min_image_height": 3138,
+            "average_image_height": 3214.301963439404,
+            "max_image_height": 3508,
+            "unique_images": 1442
+        },
+        "queries_text_statistics": {
+            "total_text_length": 8331,
+            "min_text_length": 23,
+            "average_text_length": 51.11042944785276,
+            "max_text_length": 110,
+            "unique_texts": 163
+        },
+        "queries_image_statistics": null,
+        "relevant_docs_statistics": {
+            "num_relevant_docs": 413,
+            "min_relevant_docs_per_query": 1,
+            "average_relevant_docs_per_query": 2.5337423312883436,
+            "max_relevant_docs_per_query": 6,
+            "unique_relevant_docs": 349
+        },
+        "top_ranked_statistics": null
+    }
+}

mteb/descriptive_stats/Image/DocumentUnderstanding/KoVidore2EnergyRetrieval.json

@@ -0,0 +1,32 @@
+{
+    "test": {
+        "num_samples": 2166,
+        "number_of_characters": 9764,
+        "documents_text_statistics": null,
+        "documents_image_statistics": {
+            "min_image_width": 2221,
+            "average_image_width": 2339.4957350727545,
+            "max_image_width": 2480,
+            "min_image_height": 3036,
+            "average_image_height": 3242.8138484696437,
+            "max_image_height": 3508,
+            "unique_images": 1974
+        },
+        "queries_text_statistics": {
+            "total_text_length": 9764,
+            "min_text_length": 22,
+            "average_text_length": 56.4393063583815,
+            "max_text_length": 103,
+            "unique_texts": 173
+        },
+        "queries_image_statistics": null,
+        "relevant_docs_statistics": {
+            "num_relevant_docs": 525,
+            "min_relevant_docs_per_query": 1,
+            "average_relevant_docs_per_query": 3.0346820809248554,
+            "max_relevant_docs_per_query": 7,
+            "unique_relevant_docs": 442
+        },
+        "top_ranked_statistics": null
+    }
+}

mteb/descriptive_stats/Image/DocumentUnderstanding/KoVidore2HrRetrieval.json

@@ -0,0 +1,32 @@
+{
+    "test": {
+        "num_samples": 2330,
+        "number_of_characters": 13131,
+        "documents_text_statistics": null,
+        "documents_image_statistics": {
+            "min_image_width": 1949,
+            "average_image_width": 2430.1152204836417,
+            "max_image_width": 3505,
+            "min_image_height": 2480,
+            "average_image_height": 3350.3921289710765,
+            "max_image_height": 3626,
+            "unique_images": 2096
+        },
+        "queries_text_statistics": {
+            "total_text_length": 13131,
+            "min_text_length": 21,
+            "average_text_length": 59.41628959276018,
+            "max_text_length": 112,
+            "unique_texts": 221
+        },
+        "queries_image_statistics": null,
+        "relevant_docs_statistics": {
+            "num_relevant_docs": 726,
+            "min_relevant_docs_per_query": 1,
+            "average_relevant_docs_per_query": 3.2850678733031673,
+            "max_relevant_docs_per_query": 7,
+            "unique_relevant_docs": 575
+        },
+        "top_ranked_statistics": null
+    }
+}

mteb/evaluate.py

@@ -5,7 +5,7 @@ import warnings
 from collections.abc import Iterable
 from pathlib import Path
 from time import time
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, cast
 
 from datasets.exceptions import DatasetNotFoundError
 from tqdm.auto import tqdm
@@ -27,6 +27,7 @@ from mteb.models.sentence_transformer_wrapper import (
 from mteb.results import ModelResult, TaskResult
 from mteb.results.task_result import TaskError
 from mteb.types import HFSubset, PromptType, SplitName
+from mteb.types._encoder_io import EncodeKwargs
 from mteb.types._metadata import ModelName, Revision
 
 if TYPE_CHECKING:
@@ -85,9 +86,10 @@ def _evaluate_task(
     *,
     splits: dict[SplitName, list[HFSubset]],
     co2_tracker: bool | None,
-    encode_kwargs: dict[str, Any],
+    encode_kwargs: EncodeKwargs,
     prediction_folder: Path | None,
     public_only: bool | None,
+    num_proc: int = 1,
 ) -> TaskResult | TaskError:
     """The core logic to run a model on a given task. See `evaluate` for more details.
 
@@ -270,7 +272,7 @@ def evaluate(
     *,
     co2_tracker: bool | None = None,
     raise_error: bool = True,
-    encode_kwargs: dict[str, Any] | None = None,
+    encode_kwargs: EncodeKwargs | None = None,
     cache: ResultCache | None = ResultCache(),
     overwrite_strategy: str | OverwriteStrategy = "only-missing",
     prediction_folder: Path | str | None = None,

mteb/models/abs_encoder.py

@@ -5,6 +5,7 @@ from collections.abc import Callable, Sequence
 from typing import Any, Literal, cast, get_args, overload
 
 from torch.utils.data import DataLoader
+from typing_extensions import Unpack
 
 import mteb
 from mteb.abstasks.task_metadata import TaskMetadata, TaskType
@@ -19,6 +20,7 @@ from mteb.similarity_functions import (
 from mteb.types import (
     Array,
     BatchedInput,
+    EncodeKwargs,
     PromptType,
 )
 
@@ -370,7 +372,7 @@ class AbsEncoder(ABC):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Encodes the given sentences using the encoder.
 

mteb/models/instruct_wrapper.py

@@ -92,7 +92,7 @@ def instruct_wrapper(
             logger.info(
                 f"Using instruction: '{instruction}' for task: '{task_metadata.name}'"
             )
-            embeddings = super().encode(  # type: ignore[safe-super]
+            embeddings = super().encode(  # type: ignore[safe-super,call-arg]
                 _inputs,  # type: ignore[arg-type]
                 instruction=instruction,
                 *args,

mteb/models/model_implementations/bm25.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Any
 
 from mteb._create_dataloaders import _create_text_queries_dataloader
 from mteb._requires_package import requires_package
@@ -8,6 +7,7 @@ from mteb.models.model_meta import ModelMeta
 from mteb.models.models_protocols import SearchProtocol
 from mteb.types import (
     CorpusDatasetType,
+    EncodeKwargs,
     InstructionDatasetType,
     QueryDatasetType,
     RetrievalOutputType,
@@ -49,7 +49,7 @@ def bm25_loader(model_name, **kwargs) -> SearchProtocol:
             task_metadata: TaskMetadata,
             hf_split: str,
             hf_subset: str,
-            encode_kwargs: dict[str, Any],
+            encode_kwargs: EncodeKwargs,
         ) -> None:
             logger.info("Encoding Corpus...")
             corpus_texts = [
@@ -74,7 +74,7 @@ def bm25_loader(model_name, **kwargs) -> SearchProtocol:
             hf_split: str,
             hf_subset: str,
             top_k: int,
-            encode_kwargs: dict[str, Any],
+            encode_kwargs: EncodeKwargs,
             instructions: InstructionDatasetType | None = None,
             top_ranked: TopRankedDocumentsType | None = None,
         ) -> RetrievalOutputType:

mteb/models/model_implementations/mxbai_models.py

@@ -2,7 +2,10 @@ from mteb.models.model_meta import (
     ModelMeta,
     ScoringFunction,
 )
-from mteb.models.sentence_transformer_wrapper import sentence_transformers_loader
+from mteb.models.sentence_transformer_wrapper import (
+    CrossEncoderWrapper,
+    sentence_transformers_loader,
+)
 
 mixedbread_training_data = {
     # from correspondence:
@@ -122,3 +125,117 @@ mxbai_embed_xsmall_v1 = ModelMeta(
   url={https://www.mixedbread.ai/blog/mxbai-embed-xsmall-v1},
 }""",
 )
+
+mxbai_rerank_xsmall_v1 = ModelMeta(
+    loader=CrossEncoderWrapper,
+    name="mixedbread-ai/mxbai-rerank-xsmall-v1",
+    revision="b5c6e9da73abc3711f593f705371cdbe9e0fe422",
+    release_date="2024-02-29",
+    languages=["eng-Latn"],
+    n_parameters=70830337,
+    memory_usage_mb=135.0,
+    max_tokens=512,
+    embed_dim=None,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=[
+        "PyTorch",
+        "Sentence Transformers",
+        "Transformers",
+        "ONNX",
+        "safetensors",
+    ],
+    reference="https://huggingface.co/mixedbread-ai/mxbai-rerank-xsmall-v1",
+    similarity_fn_name=None,
+    use_instructions=None,
+    training_datasets=None,
+    adapted_from=None,
+    superseded_by=None,
+    modalities=["text"],
+    model_type=["cross-encoder"],
+    citation="""@online{rerank2024mxbai,
+  title={Boost Your Search With The Crispy Mixedbread Rerank Models},
+  author={Aamir Shakir and Darius Koenig and Julius Lipp and Sean Lee},
+  year={2024},
+  url={https://www.mixedbread.ai/blog/mxbai-rerank-v1},
+}""",
+    contacts=None,
+)
+
+mxbai_rerank_base_v1 = ModelMeta(
+    loader=CrossEncoderWrapper,
+    name="mixedbread-ai/mxbai-rerank-base-v1",
+    revision="800f24c113213a187e65bde9db00c15a2bb12738",
+    release_date="2024-02-29",
+    languages=["eng-Latn"],
+    n_parameters=184422913,
+    memory_usage_mb=352.0,
+    max_tokens=512,
+    embed_dim=None,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=[
+        "PyTorch",
+        "Sentence Transformers",
+        "Transformers",
+        "ONNX",
+        "safetensors",
+    ],
+    reference="https://huggingface.co/mixedbread-ai/mxbai-rerank-base-v1",
+    similarity_fn_name=None,
+    use_instructions=None,
+    training_datasets=None,
+    adapted_from=None,
+    superseded_by=None,
+    modalities=["text"],
+    model_type=["cross-encoder"],
+    citation="""@online{rerank2024mxbai,
+  title={Boost Your Search With The Crispy Mixedbread Rerank Models},
+  author={Aamir Shakir and Darius Koenig and Julius Lipp and Sean Lee},
+  year={2024},
+  url={https://www.mixedbread.ai/blog/mxbai-rerank-v1},
+}""",
+    contacts=None,
+)
+
+mxbai_rerank_large_v1 = ModelMeta(
+    loader=CrossEncoderWrapper,
+    name="mixedbread-ai/mxbai-rerank-large-v1",
+    revision="98f655841d5caf0b16eaff79c2b4ca109d920d17",
+    release_date="2024-02-29",
+    languages=["eng-Latn"],
+    n_parameters=435062785,
+    memory_usage_mb=830.0,
+    max_tokens=512,
+    embed_dim=None,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=[
+        "PyTorch",
+        "Sentence Transformers",
+        "Transformers",
+        "ONNX",
+        "safetensors",
+    ],
+    reference="https://huggingface.co/mixedbread-ai/mxbai-rerank-large-v1",
+    similarity_fn_name=None,
+    use_instructions=None,
+    training_datasets=None,
+    adapted_from=None,
+    superseded_by=None,
+    modalities=["text"],
+    model_type=["cross-encoder"],
+    citation="""@online{rerank2024mxbai,
+  title={Boost Your Search With The Crispy Mixedbread Rerank Models},
+  author={Aamir Shakir and Darius Koenig and Julius Lipp and Sean Lee},
+  year={2024},
+  url={https://www.mixedbread.ai/blog/mxbai-rerank-v1},
+}""",
+    contacts=None,
+)

mteb/models/model_implementations/nvidia_models.py

@@ -19,14 +19,24 @@ from mteb.types import Array, BatchedInput, PromptType
 
 logger = logging.getLogger(__name__)
 
-NV_RETRIEVER_CITATION = """@misc{moreira2025nvretrieverimprovingtextembedding,
-      title={NV-Retriever: Improving text embedding models with effective hard-negative mining},
-      author={Gabriel de Souza P. Moreira and Radek Osmulski and Mengyao Xu and Ronay Ak and Benedikt Schifferer and Even Oldridge},
+NV_RETRIEVER_CITATION = """@misc{lee2025nvembedimprovedtechniquestraining,
+      title={NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models}, 
+      author={Chankyu Lee and Rajarshi Roy and Mengyao Xu and Jonathan Raiman and Mohammad Shoeybi and Bryan Catanzaro and Wei Ping},
       year={2025},
-      eprint={2407.15831},
+      eprint={2405.17428},
       archivePrefix={arXiv},
-      primaryClass={cs.IR},
-      url={https://arxiv.org/abs/2407.15831}
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2405.17428}, 
+}"""
+
+LlamaEmbedNemotron_CITATION = """@misc{babakhin2025llamaembednemotron8buniversaltextembedding,
+      title={Llama-Embed-Nemotron-8B: A Universal Text Embedding Model for Multilingual and Cross-Lingual Tasks}, 
+      author={Yauhen Babakhin and Radek Osmulski and Ronay Ak and Gabriel Moreira and Mengyao Xu and Benedikt Schifferer and Bo Liu and Even Oldridge},
+      year={2025},
+      eprint={2511.07025},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2511.07025}, 
 }"""
 
 
@@ -614,8 +624,8 @@ llama_embed_nemotron_8b = ModelMeta(
     framework=["PyTorch", "Sentence Transformers", "safetensors", "Transformers"],
     use_instructions=True,
     training_datasets=llama_embed_nemotron_training_datasets,
-    public_training_code=None,  # Will be released later
-    public_training_data=None,  # Will be released later
+    public_training_code="https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/biencoder/llama_embed_nemotron_8b",
+    public_training_data="https://huggingface.co/datasets/nvidia/embed-nemotron-dataset-v1",
     contacts=["ybabakhin"],
-    citation=NV_RETRIEVER_CITATION,
+    citation=LlamaEmbedNemotron_CITATION,
 )

mteb/models/model_implementations/octen_models.py

@@ -163,6 +163,36 @@ _PREDEFINED_PROMPTS = {
     "German1Retrieval": "Given a query, retrieve relevant passages",
 }
 
+Octen_Embedding_0B6 = ModelMeta(
+    loader=InstructSentenceTransformerModel,
+    loader_kwargs=dict(
+        instruction_template=instruction_template,
+        apply_instruction_to_passages=True,
+        prompts_dict=_PREDEFINED_PROMPTS,
+        max_seq_length=18480,
+        model_kwargs={"torch_dtype": "bfloat16"},
+    ),
+    name="bflhc/Octen-Embedding-0.6B",
+    languages=multilingual_langs,
+    open_weights=True,
+    revision="1a00a4e837bd788f6f8d91bc43201a5e52cf8ef8",
+    release_date="2026-01-10",
+    n_parameters=595776512,
+    memory_usage_mb=1136,
+    embed_dim=1024,
+    max_tokens=32768,
+    license="apache-2.0",
+    reference="https://huggingface.co/bflhc/Octen-Embedding-0.6B",
+    similarity_fn_name="cosine",
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
+    use_instructions=True,
+    public_training_code=None,
+    public_training_data=None,
+    training_datasets=training_data,
+    citation=OCTEN_CITATION,
+    adapted_from="Qwen/Qwen3-Embedding-0.6B",
+)
+
 Octen_Embedding_4B = ModelMeta(
     loader=InstructSentenceTransformerModel,
     loader_kwargs=dict(

mteb/models/model_implementations/pylate_models.py

@@ -19,6 +19,7 @@ from mteb.types import (
     Array,
     BatchedInput,
     CorpusDatasetType,
+    EncodeKwargs,
     PromptType,
     QueryDatasetType,
     RetrievalOutputType,
@@ -45,7 +46,7 @@ class PylateSearchEncoder:
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -78,7 +79,7 @@ class PylateSearchEncoder:
         hf_split: str,
         hf_subset: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
     ) -> RetrievalOutputType:
         queries_dataloader = create_dataloader(
@@ -136,7 +137,7 @@ class PylateSearchEncoder:
         hf_subset: str,
         hf_split: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[tuple[float, str]]]:
         from pylate import indexes, retrieve
 
@@ -200,7 +201,7 @@ class PylateSearchEncoder:
         task_metadata: TaskMetadata,
         hf_subset: str,
         hf_split: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[tuple[float, str]]]:
         """Rerank with PyLate's rank.rerank using per-query candidates.
 

mteb/models/models_protocols.py

@@ -1,12 +1,14 @@
 from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 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,
@@ -28,7 +30,7 @@ class SearchProtocol(Protocol):
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -49,7 +51,7 @@ class SearchProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
     ) -> RetrievalOutputType:
         """Search the corpus using the given queries.
@@ -108,7 +110,7 @@ class EncoderProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Encodes the given sentences using the encoder.
 
@@ -214,7 +216,7 @@ class CrossEncoderProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Predicts relevance scores for pairs of inputs. Note that, unlike the encoder, the cross-encoder can compare across inputs.
 

mteb/models/search_wrappers.py

@@ -14,6 +14,7 @@ from mteb.types import (
     Array,
     BatchedInput,
     CorpusDatasetType,
+    EncodeKwargs,
     PromptType,
     QueryDatasetType,
     RetrievalOutputType,
@@ -50,7 +51,7 @@ class SearchEncoderWrapper:
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -88,7 +89,7 @@ class SearchEncoderWrapper:
         hf_split: str,
         hf_subset: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
     ) -> RetrievalOutputType:
         """Search the corpus for the given queries.
@@ -215,7 +216,7 @@ class SearchEncoderWrapper:
         hf_subset: str,
         hf_split: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[tuple[float, str]]]:
         logger.info("Encoding Corpus in batches (this might take a while)...")
         if self.task_corpus is None:
@@ -318,7 +319,7 @@ class SearchEncoderWrapper:
         task_metadata: TaskMetadata,
         hf_subset: str,
         hf_split: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[tuple[float, str]]]:
         """Rerank documents based on pre-ranked documents.
 
@@ -470,7 +471,7 @@ class SearchCrossEncoderWrapper:
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -491,7 +492,7 @@ class SearchCrossEncoderWrapper:
         hf_split: str,
         hf_subset: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
     ) -> RetrievalOutputType:
         """Search the corpus using the given queries.