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

mteb/abstasks/retrieval_dataset_loaders.py

@@ -136,7 +136,7 @@ class RetrievalDatasetLoader:
                 "_id", "id"
             )
         logger.info("Loaded %d %s Documents.", len(corpus_ds), self.split.upper())
-        logger.info("Doc Example: %s", corpus_ds[0])
+        logger.debug("Doc Example: %s", corpus_ds[0])
         return corpus_ds
 
     def _load_queries(self) -> QueryDatasetType:
@@ -152,7 +152,7 @@ class RetrievalDatasetLoader:
             )
 
         logger.info("Loaded %d %s queries.", len(queries_ds), self.split.upper())
-        logger.info("Query Example: %s", queries_ds[0])
+        logger.debug("Query Example: %s", queries_ds[0])
 
         return queries_ds
 

mteb/abstasks/sts.py

@@ -7,8 +7,8 @@ from scipy.stats import pearsonr, spearmanr
 
 from mteb._evaluators import AnySTSEvaluator
 from mteb._evaluators.any_sts_evaluator import STSEvaluatorScores
-from mteb.models import EncoderProtocol
-from mteb.types import PromptType
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs, PromptType
 from mteb.types.statistics import (
     ImageStatistics,
     ScoreStatistics,
@@ -103,14 +103,17 @@ class AbsTaskSTS(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
         **kwargs: Any,
     ) -> STSMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         normalized_scores = list(map(self._normalize, data_split["score"]))
         data_split = data_split.select_columns(list(self.column_names))
 
@@ -142,7 +145,7 @@ class AbsTaskSTS(AbsTask):
     ) -> STSMetrics:
         def compute_corr(x: list[float], y: list[float]) -> tuple[float, float]:
             """Return (pearson, spearman) correlations between x and y."""
-            return pearsonr(x, y)[0], spearmanr(x, y)[0]
+            return float(pearsonr(x, y)[0]), float(spearmanr(x, y)[0])
 
         cosine_pearson, cosine_spearman = compute_corr(
             normalized_scores, scores["cosine_scores"]

mteb/abstasks/task_metadata.py

@@ -2,9 +2,10 @@ import json
 import logging
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from huggingface_hub import (
+    CardData,
     DatasetCard,
     DatasetCardData,
     constants,
@@ -150,7 +151,7 @@ _TASK_TYPE = (
     "InstructionReranking",
 ) + MIEB_TASK_TYPE
 
-TaskType = Literal[_TASK_TYPE]
+TaskType = Literal[_TASK_TYPE]  # type: ignore[valid-type]
 """The type of the task. E.g. includes "Classification", "Retrieval" and "Clustering"."""
 
 
@@ -192,8 +193,10 @@ AnnotatorType = Literal[
 """The type of the annotators. Is often important for understanding the quality of a dataset."""
 
 
-PromptDict = TypedDict(
-    "PromptDict", {prompt_type.value: str for prompt_type in PromptType}, total=False
+PromptDict = TypedDict(  # type: ignore[misc]
+    "PromptDict",
+    {prompt_type.value: str for prompt_type in PromptType},
+    total=False,
 )
 """A dictionary containing the prompt used for the task.
 
@@ -365,7 +368,7 @@ class TaskMetadata(BaseModel):
         """Return a dictionary mapping huggingface subsets to languages."""
         if isinstance(self.eval_langs, dict):
             return self.eval_langs
-        return {"default": self.eval_langs}  # type: ignore
+        return {"default": cast(list[str], self.eval_langs)}
 
     @property
     def intext_citation(self, include_cite: bool = True) -> str:
@@ -376,9 +379,8 @@ class TaskMetadata(BaseModel):
         if include_cite and cite:
             # check for whitespace in the citation
             if " " in cite:
-                logger.warning(
-                    "Citation contains whitespace. Please ensure that the citation is correctly formatted."
-                )
+                msg = "Citation contains whitespace. Please ensure that the citation is correctly formatted."
+                logger.warning(msg)
             return f"\\cite{{{cite}}}"
         return cite
 
@@ -414,7 +416,7 @@ class TaskMetadata(BaseModel):
         for subset, subset_value in stats.items():
             if subset == "hf_subset_descriptive_stats":
                 continue
-            n_samples[subset] = subset_value["num_samples"]  # type: ignore
+            n_samples[subset] = subset_value["num_samples"]
         return n_samples
 
     @property
@@ -447,7 +449,7 @@ class TaskMetadata(BaseModel):
         Raises:
             ValueError: If the prompt type is not recognized.
         """
-        if prompt_type is None:
+        if prompt_type is None or self.category is None:
             return self.modalities
         query_modalities, doc_modalities = self.category.split("2")
         category_to_modality: dict[str, Modalities] = {
@@ -467,7 +469,7 @@ class TaskMetadata(BaseModel):
 
     def _create_dataset_card_data(
         self,
-        existing_dataset_card_data: DatasetCardData | None = None,
+        existing_dataset_card_data: CardData | None = None,
     ) -> tuple[DatasetCardData, dict[str, Any]]:
         """Create a DatasetCardData object from the task metadata.
 
@@ -483,7 +485,6 @@ class TaskMetadata(BaseModel):
         dataset_type = [
             *self._hf_task_type(),
             *self._hf_task_category(),
-            *self._hf_subtypes(),
         ]
         languages = self._hf_languages()
 
@@ -502,12 +503,13 @@ class TaskMetadata(BaseModel):
 
         tags = ["mteb"] + self.modalities
 
-        descriptive_stats = self.descriptive_stats
-        if descriptive_stats is not None:
-            for split, split_stat in descriptive_stats.items():
+        descriptive_stats = ""
+        if self.descriptive_stats is not None:
+            descriptive_stats_ = self.descriptive_stats
+            for split, split_stat in descriptive_stats_.items():
                 if len(split_stat.get("hf_subset_descriptive_stats", {})) > 10:
                     split_stat.pop("hf_subset_descriptive_stats", {})
-            descriptive_stats = json.dumps(descriptive_stats, indent=4)
+            descriptive_stats = json.dumps(descriptive_stats_, indent=4)
 
         dataset_card_data_params = existing_dataset_card_data.to_dict()
         # override the existing values
@@ -584,10 +586,8 @@ class TaskMetadata(BaseModel):
 
     def _hf_subtypes(self) -> list[str]:
         # to get full list of available task_ids execute
-        # requests.post("https://huggingface.co/api/validate-yaml", json={
-        #   "content": "---\ntask_ids: 'test'\n---",
-        #   "repoType": "dataset"
-        # })
+        # https://huggingface.co/api/datasets-tags-by-type?type=task_ids
+        # ref https://huggingface-openapi.hf.space/#tag/datasets/GET/api/datasets-tags-by-type
         mteb_to_hf_subtype = {
             "Article retrieval": ["document-retrieval"],
             "Conversational retrieval": ["conversational", "utterance-retrieval"],
@@ -609,7 +609,7 @@ class TaskMetadata(BaseModel):
                 "hate-speech-detection",
             ],
             "Thematic clustering": [],
-            "Scientific Reranking": [],
+            "Scientific Reranking": ["text-scoring"],
             "Claim verification": ["fact-checking", "fact-checking-retrieval"],
             "Topic classification": ["topic-classification"],
             "Code retrieval": [],
@@ -617,21 +617,21 @@ class TaskMetadata(BaseModel):
             "Cross-Lingual Semantic Discrimination": [],
             "Textual Entailment": ["natural-language-inference"],
             "Counterfactual Detection": [],
-            "Emotion classification": [],
+            "Emotion classification": ["sentiment-classification"],
             "Reasoning as Retrieval": [],
             "Rendered Texts Understanding": [],
             "Image Text Retrieval": [],
             "Object recognition": [],
             "Scene recognition": [],
             "Caption Pairing": ["image-captioning"],
-            "Emotion recognition": [],
+            "Emotion recognition": ["sentiment-scoring"],
             "Textures recognition": [],
             "Activity recognition": [],
             "Tumor detection": [],
             "Duplicate Detection": [],
             "Rendered semantic textual similarity": [
                 "semantic-similarity-scoring",
-                "rendered semantic textual similarity",
+                "semantic-similarity-classification",
             ],
             "Intent classification": [
                 "intent-classification",
@@ -645,10 +645,8 @@ class TaskMetadata(BaseModel):
 
     def _hf_task_type(self) -> list[str]:
         # to get full list of task_types execute:
-        # requests.post("https://huggingface.co/api/validate-yaml", json={
-        #     "content": "---\ntask_categories: ['test']\n---", "repoType": "dataset"
-        # }).json()
-        # or look at https://huggingface.co/tasks
+        # https://huggingface.co/api/datasets-tags-by-type?type=task_categories
+        # ref https://huggingface-openapi.hf.space/#tag/datasets/GET/api/datasets-tags-by-type
         mteb_task_type_to_datasets = {
             # Text
             "BitextMining": ["translation"],
@@ -667,7 +665,7 @@ class TaskMetadata(BaseModel):
             "Any2AnyRetrieval": ["visual-document-retrieval"],
             "Any2AnyMultilingualRetrieval": ["visual-document-retrieval"],
             "VisionCentricQA": ["visual-question-answering"],
-            "ImageClustering": ["image-clustering"],
+            "ImageClustering": ["image-feature-extraction"],
             "ImageClassification": ["image-classification"],
             "ImageMultilabelClassification": ["image-classification"],
             "DocumentUnderstanding": ["visual-document-retrieval"],
@@ -695,11 +693,11 @@ class TaskMetadata(BaseModel):
 
     def _hf_languages(self) -> list[str]:
         languages: list[str] = []
-        if self.is_multilingual:
-            for val in list(self.eval_langs.values()):
+        if self.is_multilingual and isinstance(self.eval_langs, dict):
+            for val in self.eval_langs.values():
                 languages.extend(val)
         else:
-            languages = self.eval_langs
+            languages = cast(list[str], self.eval_langs)
         # value "python" is not valid. It must be an ISO 639-1, 639-2 or 639-3 code (two/three letters),
         # or a special value like "code", "multilingual".
         readme_langs = []
@@ -711,7 +709,7 @@ class TaskMetadata(BaseModel):
                 readme_langs.append(lang_name)
         return sorted(set(readme_langs))
 
-    def _hf_license(self) -> str:
+    def _hf_license(self) -> str | None:
         dataset_license = self.license
         if dataset_license:
             license_mapping = {

mteb/abstasks/text/bitext_mining.py

@@ -1,7 +1,7 @@
 import logging
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, ClassVar, TypedDict
+from typing import Any, ClassVar, TypedDict, cast
 
 from datasets import Dataset, DatasetDict
 from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
@@ -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,11 +73,14 @@ 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]:
         """Added load for "parallel" datasets"""
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         if not self.data_loaded:
             self.load_data()
 
@@ -87,11 +90,16 @@ class AbsTaskBitextMining(AbsTask):
         if subsets_to_run is not None:
             hf_subsets = [s for s in hf_subsets if s in subsets_to_run]
 
-        scores = {}
+        encoder_model = cast(EncoderProtocol, model)
+
+        if self.dataset is None:
+            raise ValueError("Dataset is not loaded.")
+
+        scores: dict[str, BitextMiningMetrics] = {}
         if self.parallel_subsets:
-            scores = self._evaluate_subset(
-                model,
-                self.dataset[split],  # type: ignore
+            scores = self._evaluate_subset(  # type: ignore[assignment]
+                encoder_model,
+                self.dataset[split],
                 parallel=True,
                 hf_split=split,
                 hf_subset="parallel",
@@ -109,8 +117,8 @@ class AbsTaskBitextMining(AbsTask):
                     data_split = self.dataset[split]
                 else:
                     data_split = self.dataset[hf_subset][split]
-                scores[hf_subset] = self._evaluate_subset(
-                    model,
+                scores[hf_subset] = self._evaluate_subset(  # type: ignore[assignment]
+                    encoder_model,
                     data_split,
                     hf_split=split,
                     hf_subset=hf_subset,
@@ -119,32 +127,32 @@ class AbsTaskBitextMining(AbsTask):
                     **kwargs,
                 )
 
-        return scores
+        return cast(dict[HFSubset, ScoresDict], scores)
 
     def _get_pairs(self, parallel: bool) -> list[tuple[str, str]]:
         pairs = self._DEFAULT_PAIR
         if parallel:
-            pairs = [langpair.split("-") for langpair in self.hf_subsets]
+            pairs = [langpair.split("-") for langpair in self.hf_subsets]  # type: ignore[misc]
         return pairs
 
-    def _evaluate_subset(
+    def _evaluate_subset(  # type: ignore[override]
         self,
         model: EncoderProtocol,
         data_split: Dataset,
         *,
         hf_split: str,
         hf_subset: str,
-        parallel: bool = False,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        parallel: bool = False,
         **kwargs,
-    ) -> ScoresDict:
+    ) -> BitextMiningMetrics | dict[str, BitextMiningMetrics]:
         pairs = self._get_pairs(parallel)
 
         evaluator = BitextMiningEvaluator(
             data_split,
             task_metadata=self.metadata,
-            pair_columns=pairs,  # type: ignore
+            pair_columns=pairs,
             hf_split=hf_split,
             hf_subset=hf_subset,
             **kwargs,
@@ -168,16 +176,16 @@ class AbsTaskBitextMining(AbsTask):
             )
 
         if parallel:
-            metrics = {}
+            parallel_metrics = {}
             for keys, nearest_neighbors in neighbours.items():
-                metrics[keys] = self._compute_metrics(nearest_neighbors, gold)
+                parallel_metrics[keys] = self._compute_metrics(nearest_neighbors, gold)
 
-            for v in metrics.values():
+            for v in parallel_metrics.values():
                 self._add_main_score(v)
-        else:
-            def_pair_str = "-".join(self._DEFAULT_PAIR[0])
-            metrics = self._compute_metrics(neighbours[def_pair_str], gold)
-            self._add_main_score(metrics)
+            return parallel_metrics
+        def_pair_str = "-".join(self._DEFAULT_PAIR[0])
+        metrics = self._compute_metrics(neighbours[def_pair_str], gold)
+        self._add_main_score(metrics)
         return metrics
 
     def _compute_metrics(
@@ -250,8 +258,11 @@ class AbsTaskBitextMining(AbsTask):
         )
 
     def _push_dataset_to_hub(self, repo_name: str) -> None:
+        if self.dataset is None:
+            raise ValueError("Dataset is not loaded.")
+
         if self.metadata.is_multilingual:
-            dataset = defaultdict(dict)
+            dataset: dict[str, dict[str, list[str]]] = defaultdict(dict)
             for config in self.metadata.eval_langs:
                 logger.info(f"Converting {config} of {self.metadata.name}")
 
@@ -266,10 +277,10 @@ class AbsTaskBitextMining(AbsTask):
                     for split in self.dataset[config]:
                         dataset[split][lang_1] = self.dataset[config][split][sent_1]
                         dataset[split][lang_2] = self.dataset[config][split][sent_2]
-            for split in dataset:
-                dataset[split] = Dataset.from_dict(dataset[split])
-            dataset = DatasetDict(dataset)
-            dataset.push_to_hub(repo_name)
+            dataset_dict = DatasetDict(
+                {split: Dataset.from_dict(dataset[split]) for split in dataset}
+            )
+            dataset_dict.push_to_hub(repo_name)
         else:
             sentences = {}
             for split in self.dataset:

mteb/abstasks/text/reranking.py

@@ -16,7 +16,7 @@ else:
 
 logger = logging.getLogger(__name__)
 
-OLD_FORMAT_RERANKING_TASKS = []
+OLD_FORMAT_RERANKING_TASKS: list[str] = []
 
 
 @deprecated(
@@ -100,12 +100,14 @@ class AbsTaskReranking(AbsTaskRetrieval):
         if self.metadata.name not in OLD_FORMAT_RERANKING_TASKS:
             return
 
-        logging.info(
+        logger.info(
             f"Transforming old format to standard format for {self.metadata.name}"
         )
 
         given_dataset = copy(given_dataset)
-        self.dataset = defaultdict(lambda: defaultdict(dict))
+        self.dataset: dict[str, dict[str, RetrievalSplitData]] = defaultdict(
+            lambda: defaultdict(dict)  # type: ignore[arg-type]
+        )
 
         hf_subsets = self.hf_subsets
 
@@ -115,19 +117,19 @@ class AbsTaskReranking(AbsTaskRetrieval):
                 if hf_subset in cur_dataset:
                     cur_dataset = cur_dataset[hf_subset]
             elif "name" in self.metadata.dataset:
-                cur_dataset = datasets.load_dataset(**self.metadata.dataset)  # type: ignore
+                cur_dataset = datasets.load_dataset(**self.metadata.dataset)
                 assert hf_subset == "default", (
                     f"Only default subset is supported for {self.metadata.name} since `name` is given in the metadata."
                 )
             else:
                 cur_dataset = datasets.load_dataset(
                     **self.metadata.dataset, name=hf_subset
-                )  # type: ignore
+                )
 
             for split in cur_dataset:
                 corpus = []
                 queries = []
-                relevant_docs = defaultdict(dict)
+                relevant_docs: dict[str, dict[str, int]] = defaultdict(dict)
                 top_ranked = defaultdict(list)
 
                 # Create an enumerated dataset to pass indices

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
@@ -12,7 +11,8 @@ from mteb.abstasks._statistics_calculation import (
     calculate_text_statistics,
 )
 from mteb.abstasks.abstask import AbsTask
-from mteb.models import EncoderProtocol
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs
 from mteb.types.statistics import (
     ScoreStatistics,
     SplitDescriptiveStatistics,
@@ -77,17 +77,22 @@ class AbsTaskSummarization(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> SummarizationMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         normalized_scores = [
-            (np.array(x) - self.min_score) / (self.max_score - self.min_score)
+            (
+                (np.array(x) - self.min_score) / (self.max_score - self.min_score)
+            ).tolist()
             for x in data_split[self.relevancy_column_name]
         ]
         evaluator = self.evaluator(

mteb/abstasks/zeroshot_classification.py

@@ -1,13 +1,14 @@
 import logging
 from pathlib import Path
-from typing import Any, TypedDict
+from typing import TypedDict
 
 import torch
 from datasets import Dataset
 from sklearn import metrics
 
 from mteb._evaluators import ZeroShotClassificationEvaluator
-from mteb.models import EncoderProtocol
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -111,15 +112,18 @@ class AbsTaskZeroShotClassification(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> ZeroShotClassificationMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         candidate_labels = self.get_candidate_labels()
         data_split = data_split.select_columns(
             [self.input_column_name, self.label_column_name]

mteb/benchmarks/benchmark.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Iterator, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Literal
 
@@ -19,6 +19,7 @@ class Benchmark:
 
     Args:
         name: The name of the benchmark
+        aliases: Alternative names for the benchmark
         tasks: The tasks within the benchmark.
         description: A description of the benchmark, should include its intended goal and potentially a description of its construction
         reference: A link reference, to a source containing additional information typically to a paper, leaderboard or github.
@@ -38,6 +39,7 @@ class Benchmark:
 
     name: str
     tasks: Sequence[AbsTask]
+    aliases: Sequence[str] = field(default_factory=tuple)
     description: str | None = None
     reference: StrURL | None = None
     citation: str | None = None
@@ -47,7 +49,7 @@ class Benchmark:
     display_name: str | None = None
     language_view: list[str] | Literal["all"] = field(default_factory=list)
 
-    def __iter__(self) -> Iterable[AbsTask]:
+    def __iter__(self) -> Iterator[AbsTask]:
         return iter(self.tasks)
 
     def __len__(self) -> int:

mteb/benchmarks/benchmarks/__init__.py

@@ -6,6 +6,7 @@ from mteb.benchmarks.benchmarks.benchmarks import (
     BUILT_MTEB,
     C_MTEB,
     CHEMTEB,
+    CHEMTEB_V1_1,
     CODE_RAG,
     ENCODECHKA,
     FA_MTEB,
@@ -14,6 +15,7 @@ from mteb.benchmarks.benchmarks.benchmarks import (
     JINA_VDR,
     JMTEB_LITE_V1,
     JMTEB_V2,
+    KOVIDORE_V2,
     LONG_EMBED,
     MIEB_ENG,
     MIEB_IMG,
@@ -69,6 +71,7 @@ __all__ = [
     "BRIGHT_LONG",
     "BUILT_MTEB",
     "CHEMTEB",
+    "CHEMTEB_V1_1",
     "CODE_RAG",
     "C_MTEB",
     "ENCODECHKA",
@@ -79,6 +82,7 @@ __all__ = [
     "JINA_VDR",
     "JMTEB_LITE_V1",
     "JMTEB_V2",
+    "KOVIDORE_V2",
     "LONG_EMBED",
     "MIEB_ENG",
     "MIEB_IMG",