mteb 2.6.6__py3-none-any.whl → 2.6.8__py3-none-any.whl

mteb/_create_dataloaders.py

@@ -23,7 +23,7 @@ logger = logging.getLogger(__name__)
 def _create_dataloader_from_texts(
     text: list[str],
     batch_size: int = 32,
-    **kwargs: dict[str, Any],
+    **kwargs: Any,
 ) -> DataLoader[TextInput]:
     """Create a dataloader from a list of text.
 
@@ -191,7 +191,8 @@ def _create_dataloader_for_queries_conversation(
     """
     return DataLoader(
         queries.map(
-            _convert_conv_history_to_query, desc="Converting conversations to queries"
+            _convert_conv_history_to_query,
+            desc="Converting conversations to queries",
         ),
         collate_fn=_custom_collate_fn,
         batch_size=batch_size,
@@ -361,6 +362,9 @@ def _create_document_dataloader(
         task_metadata: Metadata of the task to determine the document type.
         input_column: The column to use as input. If None, it will use the first column that matches the modality.
         batch_size: Batch size for the dataloader.
+
+    Returns:
+        A dataloader for the documents.
     """
     document_type = task_metadata.get_modalities(PromptType.document)
     if document_type == ["text"]:  # text only
@@ -383,7 +387,7 @@ def create_dataloader(
     prompt_type: PromptType | None = None,
     input_column: str | None = None,
     batch_size: int = 32,
-    **kwargs: dict[str, Any],
+    **kwargs: Any,
 ) -> DataLoader[BatchedInput]:
     """Create a dataloader from a dataset.
 

mteb/_evaluators/any_sts_evaluator.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, TypedDict
+from typing import TypedDict
 
 from datasets import Dataset
 from sklearn.metrics.pairwise import (
@@ -12,7 +12,7 @@ from mteb._create_dataloaders import create_dataloader
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.similarity_functions import compute_pairwise_similarity
-from mteb.types import PromptType
+from mteb.types import EncodeKwargs, PromptType
 
 from .evaluator import Evaluator
 
@@ -57,7 +57,10 @@ class AnySTSEvaluator(Evaluator):
         self.input2_prompt_type = input2_prompt_type
 
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
     ) -> STSEvaluatorScores:
         logger.info("Running semantic similarity - Encoding samples (1/2)")
         embeddings1 = model.encode(

mteb/_evaluators/clustering_evaluator.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Any
 
 from datasets import Dataset
 from sklearn import cluster
@@ -7,6 +6,7 @@ from sklearn import cluster
 from mteb._create_dataloaders import create_dataloader
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
+from mteb.types import EncodeKwargs
 
 from .evaluator import Evaluator
 
@@ -38,7 +38,7 @@ class ClusteringEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> list[int]:
         data_loader = create_dataloader(
             self.dataset,

mteb/_evaluators/evaluator.py

@@ -4,6 +4,7 @@ from typing import Any
 
 from mteb.abstasks.abstask import _set_seed
 from mteb.models import EncoderProtocol
+from mteb.types import EncodeKwargs
 
 
 class Evaluator(ABC):
@@ -18,7 +19,7 @@ class Evaluator(ABC):
 
     @abstractmethod
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self, model: EncoderProtocol, *, encode_kwargs: EncodeKwargs
     ) -> Mapping[str, float] | Iterable[Any]:
         """This is called during training to evaluate the model.
 

mteb/_evaluators/image/imagetext_pairclassification_evaluator.py

@@ -6,16 +6,17 @@ from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
-from datasets import Dataset
 from torch.utils.data import DataLoader
 
 from mteb._create_dataloaders import (
+    _create_dataloader_from_texts,
     _transform_image_to_rgb,
 )
 from mteb._evaluators.evaluator import Evaluator
 from mteb._requires_package import requires_image_dependencies
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.models_protocols import EncoderProtocol
+from mteb.types import EncodeKwargs
 
 if TYPE_CHECKING:
     from PIL.Image import Image
@@ -84,7 +85,10 @@ class ImageTextPairClassificationEvaluator(Evaluator):
         self.hf_subset = hf_subset
 
     def __call__(  # type: ignore[override]
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
     ) -> list[torch.Tensor]:
         images = []
         if isinstance(self.images_column_names, str):
@@ -105,8 +109,8 @@ class ImageTextPairClassificationEvaluator(Evaluator):
                     texts.append(row[col])
 
         text_embeddings = model.encode(
-            DataLoader(
-                Dataset.from_dict({"text": texts}),
+            _create_dataloader_from_texts(
+                texts,
                 **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
@@ -127,7 +131,6 @@ class ImageTextPairClassificationEvaluator(Evaluator):
             DataLoader(
                 CustomImageDataset(images),
                 collate_fn=lambda x: {"image": [item["image"] for item in x]},
-                **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,

mteb/_evaluators/pair_classification_evaluator.py

@@ -14,7 +14,7 @@ from mteb._evaluators.evaluator import Evaluator
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.similarity_functions import compute_pairwise_similarity
-from mteb.types import PromptType
+from mteb.types import EncodeKwargs, PromptType
 
 logger = logging.getLogger(__name__)
 
@@ -85,7 +85,7 @@ class PairClassificationEvaluator(Evaluator):
     def __call__(
         self,
         model: EncoderProtocol,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> PairClassificationDistances:
         logger.info("Running pair classification - Encoding samples (1/2)")
         embeddings1 = model.encode(

mteb/_evaluators/retrieval_evaluator.py

@@ -1,11 +1,11 @@
 import logging
 from collections.abc import Sequence
-from typing import Any
 
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import SearchProtocol
 from mteb.types import (
     CorpusDatasetType,
+    EncodeKwargs,
     QueryDatasetType,
     RelevantDocumentsType,
     RetrievalEvaluationResult,
@@ -48,7 +48,7 @@ class RetrievalEvaluator(Evaluator):
     def __call__(  # type: ignore[override]
         self,
         search_model: SearchProtocol,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> RetrievalOutputType:
         logger.info("Running retrieval task - Indexing corpus...")
         search_model.index(

mteb/_evaluators/sklearn_evaluator.py

@@ -9,7 +9,7 @@ from typing_extensions import Self
 from mteb._create_dataloaders import create_dataloader
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
-from mteb.types import Array, BatchedInput
+from mteb.types import Array, BatchedInput, EncodeKwargs
 
 from .evaluator import Evaluator
 
@@ -50,7 +50,7 @@ class SklearnEvaluator(Evaluator):
         self.evaluator_model = evaluator_model
 
     def create_dataloaders(
-        self, encode_kwargs: dict[str, Any]
+        self, encode_kwargs: EncodeKwargs
     ) -> tuple[DataLoader[BatchedInput], DataLoader[BatchedInput]]:
         dataloader_train = create_dataloader(
             self.train_dataset,
@@ -70,7 +70,7 @@ class SklearnEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         test_cache: Array | None = None,
     ) -> tuple[np.ndarray, Array]:
         """Classification evaluation by training a sklearn classifier on the embeddings of the training set and evaluating on the embeddings of the test set.

mteb/_evaluators/text/bitext_mining_evaluator.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Any
 
 import torch
 from datasets import Dataset
@@ -9,7 +8,7 @@ from mteb._create_dataloaders import _create_dataloader_from_texts
 from mteb._evaluators.evaluator import Evaluator
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
-from mteb.types import Array
+from mteb.types import Array, EncodeKwargs
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +32,10 @@ class BitextMiningEvaluator(Evaluator):
         self.task_metadata = task_metadata
 
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[dict[str, float]]]:
         pair_elements = {p for pair in self.pairs for p in pair}
         if isinstance(self.sentences, Dataset):

mteb/_evaluators/text/summarization_evaluator.py

@@ -1,6 +1,6 @@
 import logging
 import sys
-from typing import Any, TypedDict
+from typing import TypedDict
 
 import numpy as np
 import torch
@@ -12,6 +12,7 @@ from mteb._evaluators.evaluator import Evaluator
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.similarity_functions import cos_sim, dot_score
+from mteb.types import EncodeKwargs
 
 # if later than python 3.13 use typing module
 if sys.version_info >= (3, 13):
@@ -94,7 +95,7 @@ class SummarizationEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> SummarizationDistances:
         # Get the human & machine summaries for the text in one go for all
         human_lens = [len(human_summaries) for human_summaries in self.human_summaries]

mteb/_evaluators/zeroshot_classification_evaluator.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Any
 
 from datasets import Dataset
 
@@ -10,7 +9,7 @@ from mteb._create_dataloaders import (
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.similarity_functions import similarity
-from mteb.types import Array
+from mteb.types import Array, EncodeKwargs
 
 from .evaluator import Evaluator
 
@@ -38,7 +37,10 @@ class ZeroShotClassificationEvaluator(Evaluator):
         self.hf_subset = hf_subset
 
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
     ) -> Array:
         dataloader = create_dataloader(
             self.dataset,

mteb/abstasks/abstask.py

@@ -23,6 +23,7 @@ from mteb.models import (
     SearchProtocol,
 )
 from mteb.types import HFSubset, Modalities, ScoresDict
+from mteb.types._encoder_io import EncodeKwargs
 from mteb.types.statistics import DescriptiveStatistics, SplitDescriptiveStatistics
 
 logger = logging.getLogger(__name__)
@@ -121,7 +122,7 @@ class AbsTask(ABC):
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> Mapping[HFSubset, ScoresDict]:
@@ -201,7 +202,7 @@ class AbsTask(ABC):
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> ScoresDict:

mteb/abstasks/aggregated_task.py

@@ -9,7 +9,7 @@ from datasets import Dataset, DatasetDict
 
 from mteb.models.models_protocols import MTEBModels
 from mteb.results.task_result import TaskResult
-from mteb.types import HFSubset, ScoresDict
+from mteb.types import EncodeKwargs, HFSubset, ScoresDict
 from mteb.types.statistics import DescriptiveStatistics
 
 from .abstask import AbsTask
@@ -127,7 +127,7 @@ class AbsTaskAggregate(AbsTask):
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
@@ -141,7 +141,7 @@ class AbsTaskAggregate(AbsTask):
         self,
         model: MTEBModels,
         data_split: DatasetDict | Dataset,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         **kwargs: Any,
     ) -> ScoresDict:
         raise NotImplementedError(

mteb/abstasks/classification.py

@@ -16,7 +16,7 @@ from sklearn.metrics import (
 
 from mteb._evaluators.sklearn_evaluator import SklearnEvaluator, SklearnModelProtocol
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import HFSubset, ScoresDict
+from mteb.types import EncodeKwargs, HFSubset, ScoresDict
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -125,7 +125,7 @@ class AbsTaskClassification(AbsTask):
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
@@ -184,7 +184,7 @@ class AbsTaskClassification(AbsTask):
         model: MTEBModels,
         data_split: DatasetDict,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,

mteb/abstasks/clustering.py

@@ -12,7 +12,7 @@ from sklearn.metrics.cluster import v_measure_score
 
 from mteb._create_dataloaders import create_dataloader
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import Array, HFSubset, ScoresDict
+from mteb.types import Array, EncodeKwargs, HFSubset, ScoresDict
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -156,7 +156,7 @@ class AbsTaskClustering(AbsTask):
         model: MTEBModels,
         data_split: Dataset,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,

mteb/abstasks/clustering_legacy.py

@@ -9,7 +9,7 @@ from sklearn import metrics
 
 from mteb._evaluators import ClusteringEvaluator
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import ScoresDict
+from mteb.types import EncodeKwargs, ScoresDict
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -83,7 +83,7 @@ class AbsTaskClusteringLegacy(AbsTask):
         model: MTEBModels,
         data_split: Dataset,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,

mteb/abstasks/image/image_text_pair_classification.py

@@ -13,6 +13,7 @@ from mteb.abstasks._statistics_calculation import (
 )
 from mteb.abstasks.abstask import AbsTask
 from mteb.models.models_protocols import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs
 from mteb.types.statistics import (
     ImageStatistics,
     SplitDescriptiveStatistics,
@@ -119,7 +120,7 @@ class AbsTaskImageTextPairClassification(AbsTask):
         model: MTEBModels,
         data_split: Dataset,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,

mteb/abstasks/multilabel_classification.py

@@ -17,7 +17,7 @@ from mteb._create_dataloaders import create_dataloader
 from mteb._evaluators.classification_metrics import hamming_score
 from mteb._evaluators.sklearn_evaluator import SklearnModelProtocol
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import Array
+from mteb.types import Array, EncodeKwargs
 
 from .classification import AbsTaskClassification
 
@@ -83,7 +83,7 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
         model: MTEBModels,
         data_split: DatasetDict,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,

mteb/abstasks/pair_classification.py

@@ -19,7 +19,7 @@ from mteb.abstasks._statistics_calculation import (
 from mteb.abstasks.abstask import AbsTask
 from mteb.models.model_meta import ScoringFunction
 from mteb.models.models_protocols import EncoderProtocol, MTEBModels
-from mteb.types import PromptType
+from mteb.types import EncodeKwargs, PromptType
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -84,7 +84,7 @@ class AbsTaskPairClassification(AbsTask):
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, str],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> dict[str, float]:

mteb/abstasks/retrieval.py

@@ -25,6 +25,7 @@ from mteb.models import (
     SearchProtocol,
 )
 from mteb.types import (
+    EncodeKwargs,
     HFSubset,
     QueryDatasetType,
     RelevantDocumentsType,
@@ -184,17 +185,17 @@ class AbsTaskRetrieval(AbsTask):
             return queries, corpus
 
         if self.metadata.is_multilingual:
-            for subset in self.queries:
-                for split in self.queries[subset]:
-                    queries = self.queries[subset][split]
-                    corpus = self.corpus[subset][split]
+            for subset in self.queries:  # type: ignore[attr-defined]
+                for split in self.queries[subset]:  # type: ignore[attr-defined]
+                    queries = self.queries[subset][split]  # type: ignore[attr-defined]
+                    corpus = self.corpus[subset][split]  # type: ignore[attr-defined]
 
                     (
                         self.dataset[subset][split]["queries"],
                         self.dataset[subset][split]["corpus"],
                     ) = _process_split(queries, corpus)
 
-                    self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[
+                    self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[  # type: ignore[attr-defined]
                         subset
                     ][split]
                     if hasattr(self, "instructions"):
@@ -211,15 +212,15 @@ class AbsTaskRetrieval(AbsTask):
                         ][split]
         else:
             subset = "default"
-            for split in self.queries:
-                queries = self.queries[split]
-                corpus = self.corpus[split]
+            for split in self.queries:  # type: ignore[attr-defined]
+                queries = self.queries[split]  # type: ignore[attr-defined]
+                corpus = self.corpus[split]  # type: ignore[attr-defined]
                 (
                     self.dataset[subset][split]["queries"],
                     self.dataset[subset][split]["corpus"],
                 ) = _process_split(queries, corpus)
 
-                self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[
+                self.dataset[subset][split]["relevant_docs"] = self.relevant_docs[  # type: ignore[attr-defined]
                     split
                 ].copy()
                 if hasattr(self, "instructions"):
@@ -235,9 +236,9 @@ class AbsTaskRetrieval(AbsTask):
                         split
                     ].copy()
 
-        del self.queries
-        del self.corpus
-        del self.relevant_docs
+        del self.queries  # type: ignore[attr-defined]
+        del self.corpus  # type: ignore[attr-defined]
+        del self.relevant_docs  # type: ignore[attr-defined]
         if hasattr(self, "instructions"):
             del self.instructions
         if hasattr(self, "top_ranked"):
@@ -283,7 +284,7 @@ class AbsTaskRetrieval(AbsTask):
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> Mapping[HFSubset, ScoresDict]:
@@ -320,7 +321,7 @@ class AbsTaskRetrieval(AbsTask):
         self,
         model: MTEBModels,
         data_split: RetrievalSplitData,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,

mteb/abstasks/sts.py

@@ -8,7 +8,7 @@ from scipy.stats import pearsonr, spearmanr
 from mteb._evaluators import AnySTSEvaluator
 from mteb._evaluators.any_sts_evaluator import STSEvaluatorScores
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import PromptType
+from mteb.types import EncodeKwargs, PromptType
 from mteb.types.statistics import (
     ImageStatistics,
     ScoreStatistics,
@@ -105,7 +105,7 @@ class AbsTaskSTS(AbsTask):
         self,
         model: MTEBModels,
         data_split: Dataset,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,

mteb/abstasks/text/bitext_mining.py

@@ -10,7 +10,7 @@ from mteb._evaluators import BitextMiningEvaluator
 from mteb.abstasks._statistics_calculation import calculate_text_statistics
 from mteb.abstasks.abstask import AbsTask
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import HFSubset, ScoresDict
+from mteb.types import EncodeKwargs, HFSubset, ScoresDict
 from mteb.types.statistics import SplitDescriptiveStatistics, TextStatistics
 
 logger = logging.getLogger(__name__)
@@ -73,7 +73,7 @@ class AbsTaskBitextMining(AbsTask):
         split: str = "test",
         subsets_to_run: list[HFSubset] | None = None,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
@@ -142,7 +142,7 @@ class AbsTaskBitextMining(AbsTask):
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         parallel: bool = False,
         **kwargs,

mteb/abstasks/text/summarization.py

@@ -1,6 +1,5 @@
 import logging
 from pathlib import Path
-from typing import Any
 
 import numpy as np
 from datasets import Dataset
@@ -13,6 +12,7 @@ from mteb.abstasks._statistics_calculation import (
 )
 from mteb.abstasks.abstask import AbsTask
 from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs
 from mteb.types.statistics import (
     ScoreStatistics,
     SplitDescriptiveStatistics,
@@ -82,7 +82,7 @@ class AbsTaskSummarization(AbsTask):
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> SummarizationMetrics:

mteb/abstasks/zeroshot_classification.py

@@ -1,6 +1,6 @@
 import logging
 from pathlib import Path
-from typing import Any, TypedDict
+from typing import TypedDict
 
 import torch
 from datasets import Dataset
@@ -8,6 +8,7 @@ from sklearn import metrics
 
 from mteb._evaluators import ZeroShotClassificationEvaluator
 from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -116,7 +117,7 @@ class AbsTaskZeroShotClassification(AbsTask):
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> ZeroShotClassificationMetrics:

mteb/benchmarks/benchmarks/__init__.py

@@ -14,6 +14,7 @@ from mteb.benchmarks.benchmarks.benchmarks import (
     JINA_VDR,
     JMTEB_LITE_V1,
     JMTEB_V2,
+    KOVIDORE_V2,
     LONG_EMBED,
     MIEB_ENG,
     MIEB_IMG,
@@ -79,6 +80,7 @@ __all__ = [
     "JINA_VDR",
     "JMTEB_LITE_V1",
     "JMTEB_V2",
+    "KOVIDORE_V2",
     "LONG_EMBED",
     "MIEB_ENG",
     "MIEB_IMG",