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

mteb/results/model_result.py

@@ -1,29 +1,37 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 import pandas as pd
 from pydantic import BaseModel, ConfigDict, Field
-from typing_extensions import Self
+from typing_extensions import overload
 
-from mteb.abstasks.abstask import AbsTask
-from mteb.abstasks.task_metadata import (
-    TaskDomain,
-    TaskType,
-)
 from mteb.types import (
-    ISOLanguage,
-    ISOLanguageScript,
     Modalities,
-    Score,
-    ScoresDict,
-    SplitName,
 )
 
 from .task_result import TaskError, TaskResult
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    from mteb.abstasks.abstask import AbsTask
+    from mteb.abstasks.task_metadata import (
+        TaskDomain,
+        TaskType,
+    )
+    from mteb.types import (
+        ISOLanguage,
+        ISOLanguageScript,
+        Score,
+        ScoresDict,
+        SplitName,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -58,7 +66,7 @@ def _aggregate_and_pivot(
             index=index_columns,
             columns=columns,
             values="score",
-            aggfunc=aggregation_fn,
+            aggfunc=aggregation_fn,  # type: ignore[arg-type]
         ).reset_index()
     elif format == "long":
         return (
@@ -81,7 +89,7 @@ class ModelResult(BaseModel):
     model_revision: str | None
     task_results: list[TaskResult]
     default_modalities: list[Modalities] = Field(
-        default_factory=lambda: ["text"], alias="modalities"
+        default_factory=lambda: [cast("Modalities", "text")], alias="modalities"
     )
     model_config = (
         ConfigDict(  # to free up the name model_* which is otherwise protected
@@ -95,16 +103,17 @@ class ModelResult(BaseModel):
         return f"ModelResult(model_name={self.model_name}, model_revision={self.model_revision}, task_results=[...](#{n_entries}))"
 
     @classmethod
-    def from_validated(cls, **data: dict[str, Any]) -> Self:
+    def from_validated(cls, **data: dict[str, Any]) -> ModelResult:
         """Create a ModelResult from validated data.
 
         Args:
             data: The validated data.
         """
-        data["task_results"] = [
-            TaskResult.from_validated(**res) for res in data["task_results"]
+        data["task_results"] = [  # type: ignore[assignment]
+            TaskResult.from_validated(**res)  # type: ignore[arg-type]
+            for res in data["task_results"]
         ]
-        return cls.model_construct(**data)
+        return cls.model_construct(**data)  # type: ignore[arg-type]
 
     def _filter_tasks(
         self,
@@ -114,7 +123,7 @@ class ModelResult(BaseModel):
         task_types: list[TaskType] | None = None,
         modalities: list[Modalities] | None = None,
         is_public: bool | None = None,
-    ) -> Self:
+    ) -> ModelResult:
         new_task_results = []
         for task_result in self.task_results:
             if (task_names is not None) and (task_result.task_name not in task_names):
@@ -142,7 +151,7 @@ class ModelResult(BaseModel):
             task_results=new_task_results,
         )
 
-    def select_tasks(self, tasks: Sequence[AbsTask]) -> Self:
+    def select_tasks(self, tasks: Iterable[AbsTask]) -> ModelResult:
         """Select tasks from the ModelResult based on a list of AbsTask objects.
 
         Args:
@@ -160,6 +169,28 @@ class ModelResult(BaseModel):
             task_results=new_task_results,
         )
 
+    @overload
+    def _get_scores(
+        self,
+        splits: list[SplitName] | None = None,
+        languages: list[ISOLanguage | ISOLanguageScript] | None = None,
+        scripts: list[ISOLanguageScript] | None = None,
+        getter: Callable[[ScoresDict], Score] | None = None,
+        aggregation: Callable[[list[Score]], Any] | None = None,
+        format: Literal["wide"] = "wide",
+    ) -> dict: ...
+
+    @overload
+    def _get_scores(
+        self,
+        splits: list[SplitName] | None = None,
+        languages: list[ISOLanguage | ISOLanguageScript] | None = None,
+        scripts: list[ISOLanguageScript] | None = None,
+        getter: Callable[[ScoresDict], Score] | None = None,
+        aggregation: Callable[[list[Score]], Any] | None = None,
+        format: Literal["long"] = "long",
+    ) -> list: ...
+
     def _get_scores(
         self,
         splits: list[SplitName] | None = None,
@@ -177,21 +208,24 @@ class ModelResult(BaseModel):
             aggregation = aggregation if aggregation is not None else np.mean
         else:
             use_fast = True
+        aggregation = cast("Callable[[list[Score]], Any]", aggregation)
+        getter = cast("Callable[[ScoresDict], Score]", getter)
+
         if format == "wide":
             scores = {}
             for res in self.task_results:
                 try:
                     if use_fast:
                         scores[res.task_name] = res._get_score_fast(
-                            splits=splits,  # type: ignore
-                            languages=languages,  # type: ignore
+                            splits=splits,
+                            languages=languages,
                         )
                     else:
                         scores[res.task_name] = res.get_score(
                             splits=splits,
                             languages=languages,
-                            aggregation=aggregation,  # type: ignore
-                            getter=getter,  # type: ignore
+                            aggregation=aggregation,
+                            getter=getter,
                             scripts=scripts,
                         )
                 except Exception as e:
@@ -206,14 +240,14 @@ class ModelResult(BaseModel):
                     if use_fast:
                         score = task_res._get_score_fast(
                             splits=splits,
-                            languages=languages,  # type: ignore
+                            languages=languages,
                         )
                     else:
                         score = task_res.get_score(
                             splits=splits,
                             languages=languages,
-                            aggregation=aggregation,  # type: ignore
-                            getter=getter,  # type: ignore
+                            aggregation=aggregation,
+                            getter=getter,
                             scripts=scripts,
                         )
                     entry = dict(
@@ -292,7 +326,9 @@ class ModelResult(BaseModel):
         scores_data = self._get_score_for_table()
 
         if not scores_data:
-            logger.warning("No scores data available. Returning empty DataFrame.")
+            msg = "No scores data available. Returning empty DataFrame."
+            logger.warning(msg)
+            warnings.warn(msg)
             return pd.DataFrame()
 
         # Create DataFrame
@@ -315,7 +351,7 @@ class ModelResult(BaseModel):
     def __hash__(self) -> int:
         return id(self)
 
-    def __iter__(self) -> Iterable[TaskResult]:
+    def __iter__(self) -> Iterable[TaskResult]:  # type: ignore[override]
         return iter(self.task_results)
 
     def __getitem__(self, index) -> TaskResult:
@@ -368,13 +404,13 @@ class ModelResult(BaseModel):
         return [task_res.task_name for task_res in self.task_results]
 
     @property
-    def modalities(self) -> list[str]:
+    def modalities(self) -> list[Modalities]:
         """Get all modalities in the task results.
 
         Returns:
             A list of modalities in the task results.
         """
-        mods = []
+        mods: list[Modalities] = []
         for task_res in self.task_results:
             task_modalities = getattr(task_res, "modalities", [])
             mods.extend(task_modalities)

mteb/results/task_result.py

@@ -2,33 +2,42 @@ from __future__ import annotations
 
 import json
 import logging
-from argparse import Namespace
+import warnings
 from collections import defaultdict
-from collections.abc import Callable, Iterable
 from functools import cached_property
 from importlib.metadata import version
-from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 from huggingface_hub import EvalResult
 from packaging.version import Version
 from pydantic import BaseModel, field_validator
-from typing_extensions import Self
 
+from mteb import TaskMetadata
 from mteb._helpful_enum import HelpfulStrEnum
+from mteb.abstasks import AbsTaskClassification
 from mteb.abstasks.abstask import AbsTask
 from mteb.languages import LanguageScripts
 from mteb.models.model_meta import ScoringFunction
 from mteb.types import (
-    HFSubset,
-    ISOLanguage,
-    ISOLanguageScript,
-    Score,
     ScoresDict,
     SplitName,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Mapping
+    from pathlib import Path
+
+    from typing_extensions import Self
+
+    from mteb.abstasks.task_metadata import TaskDomain
+    from mteb.types import (
+        HFSubset,
+        ISOLanguage,
+        ISOLanguageScript,
+        Score,
+    )
+
 logger = logging.getLogger(__name__)
 
 
@@ -39,67 +48,59 @@ class Criteria(HelpfulStrEnum):
     DATASET_REVISION = "dataset_revision"
 
 
-class ScalaNbClassificationDummy:
+class ScalaNbClassificationDummy(AbsTaskClassification):
     """A dummy task for loading historic results from before v1.11.0"""
 
-    metadata = Namespace(  # type: ignore
+    metadata = TaskMetadata(
         name="ScalaNbClassification",
+        description="A dummy",
         main_score="accuracy",
         type="Classification",
-        hf_subsets_to_langscripts={
-            "default": ["nob-Latn"],
-        },
-        dataset={"revision": "revision_not_applicable"},
-        revision="revision_not_applicable",
+        eval_langs=["nob-Latn"],
+        dataset={"path": "not/exists", "revision": "revision_not_applicable"},
     )
 
 
-class ScalaNnClassificationDummy:
+class ScalaNnClassificationDummy(AbsTaskClassification):
     """A dummy task for loading historic results from before v1.11.0"""
 
-    metadata = Namespace(  # type: ignore
+    metadata = TaskMetadata(
         name="ScalaNnClassification",
+        description="A dummy",
         main_score="accuracy",
         type="Classification",
-        hf_subsets_to_langscripts={
-            "default": ["nno-Latn"],
-        },
-        dataset={"revision": "revision_not_applicable"},
-        revision="revision_not_applicable",
+        eval_langs=["nob-Latn"],
+        dataset={"path": "not/exists", "revision": "revision_not_applicable"},
     )
 
 
-class ScalaDaClassificationDummy:
+class ScalaDaClassificationDummy(AbsTaskClassification):
     """A dummy task for loading historic results from before v1.11.0"""
 
-    metadata = Namespace(  # type: ignore
+    metadata = TaskMetadata(
         name="ScalaDaClassification",
+        description="A dummy",
         main_score="accuracy",
         type="Classification",
-        hf_subsets_to_langscripts={
-            "default": ["dan-Latn"],
-        },
-        dataset={"revision": "revision_not_applicable"},
-        revision="revision_not_applicable",
+        eval_langs=["dan-Latn"],
+        dataset={"path": "not/exists", "revision": "revision_not_applicable"},
     )
 
 
-class ScalaSvClassificationDummy:
+class ScalaSvClassificationDummy(AbsTaskClassification):
     """A dummy task for loading historic results from before v1.11.0"""
 
-    metadata = Namespace(  # type: ignore
+    metadata = TaskMetadata(
         name="ScalaSvClassification",
+        description="A dummy",
         main_score="accuracy",
         type="Classification",
-        hf_subsets_to_langscripts={
-            "default": ["swe-Latn"],
-        },
-        dataset={"revision": "revision_not_applicable"},
-        revision="revision_not_applicable",
+        eval_langs=["swe-Latn"],
+        dataset={"path": "not/exists", "revision": "revision_not_applicable"},
     )
 
 
-outdated_tasks = {
+outdated_tasks: dict[str, type[AbsTask]] = {
     "ScalaNbClassification": ScalaNbClassificationDummy,
     "ScalaNnClassification": ScalaNnClassificationDummy,
     "ScalaDaClassification": ScalaDaClassificationDummy,
@@ -166,10 +167,10 @@ class TaskResult(BaseModel):
     def from_task_results(
         cls,
         task: AbsTask | type[AbsTask],
-        scores: dict[SplitName, dict[HFSubset, ScoresDict]],
+        scores: dict[SplitName, Mapping[HFSubset, ScoresDict]],
         evaluation_time: float,
         kg_co2_emissions: float | None = None,
-    ) -> Self:
+    ) -> TaskResult:
         """Create a TaskResult from the task and scores.
 
         Args:
@@ -246,12 +247,12 @@ class TaskResult(BaseModel):
         return get_task(self.task_name)
 
     @property
-    def domains(self) -> list[str]:
+    def domains(self) -> list[TaskDomain]:
         """Get the domains of the task."""
         doms = self.task.metadata.domains
         if doms is None:
             doms = []
-        return doms  # type: ignore
+        return doms
 
     @property
     def task_type(self) -> str:
@@ -307,7 +308,7 @@ class TaskResult(BaseModel):
                     if isinstance(v, dict):
                         self._round_scores(v, n)
                     elif isinstance(v, float):
-                        value[i] = round(v, n)
+                        value[i] = round(v, n)  # type: ignore[call-overload]
 
             elif isinstance(value, float):
                 scores[key] = round(value, n)
@@ -325,7 +326,7 @@ class TaskResult(BaseModel):
             json.dump(json_obj, f, indent=2)
 
     @classmethod
-    def from_disk(cls, path: Path, load_historic_data: bool = True) -> Self:  # type: ignore
+    def from_disk(cls, path: Path, load_historic_data: bool = True) -> TaskResult:
         """Load TaskResult from disk.
 
         Args:
@@ -356,7 +357,7 @@ class TaskResult(BaseModel):
         )  # assume it is before 1.11.0 if the version is not present
 
         try:
-            obj = cls.model_validate(data)
+            obj: TaskResult = cls.model_validate(data)
         except Exception as e:
             if not pre_1_11_load:
                 raise e
@@ -381,6 +382,7 @@ class TaskResult(BaseModel):
         from mteb import get_task
 
         task_name = obj.task_name
+        task: AbsTask | type[AbsTask]
         if task_name in outdated_tasks:
             task = outdated_tasks[task_name]
         else:
@@ -393,11 +395,11 @@ class TaskResult(BaseModel):
                     for key in list(hf_subset_scores.keys()):
                         if isinstance(hf_subset_scores[key], dict):
                             for k, v in hf_subset_scores[key].items():
-                                hf_subset_scores[f"{key}_{k}"] = v
-                            hf_subset_scores.pop(key)
+                                hf_subset_scores[f"{key}_{k}"] = v  # type: ignore[index]
+                            hf_subset_scores.pop(key)  # type: ignore[attr-defined]
 
     @classmethod
-    def _convert_from_before_v1_11_0(cls, data: dict) -> Self:
+    def _convert_from_before_v1_11_0(cls, data: dict) -> TaskResult:
         from mteb.get_tasks import _TASKS_REGISTRY
 
         # in case the task name is not found in the registry, try to find a lower case version
@@ -462,7 +464,9 @@ class TaskResult(BaseModel):
                     if main_score in hf_subset_scores:
                         hf_subset_scores["main_score"] = hf_subset_scores[main_score]
                     else:
-                        logger.warning(f"Main score {main_score} not found in scores")
+                        msg = f"Main score {main_score} not found in scores"
+                        logger.warning(msg)
+                        warnings.warn(msg)
                         hf_subset_scores["main_score"] = None
 
         # specific fixes:
@@ -481,7 +485,7 @@ class TaskResult(BaseModel):
                 scores["test"]["fra-fra"] = scores["test"].pop("fr")
 
         result: TaskResult = TaskResult.from_task_results(
-            task,  # type: ignore
+            task,
             scores,
             evaluation_time,
             kg_co2_emissions=None,
@@ -532,7 +536,7 @@ class TaskResult(BaseModel):
     def _get_score_fast(
         self,
         splits: Iterable[str] | None = None,
-        languages: str | None = None,
+        languages: list[ISOLanguage | ISOLanguageScript] | None = None,
         subsets: Iterable[str] | None = None,
     ) -> float:
         """Sped up version of get_score that will be used if no aggregation, script or getter needs to be specified.
@@ -581,7 +585,7 @@ class TaskResult(BaseModel):
         return val_sum / n_val
 
     @classmethod
-    def from_validated(cls, **data) -> Self:
+    def from_validated(cls, **data) -> TaskResult:
         """Create a TaskResult from validated data.
 
         Returns:
@@ -592,13 +596,13 @@ class TaskResult(BaseModel):
     def __repr__(self) -> str:
         return f"TaskResult(task_name={self.task_name}, scores=...)"
 
-    def only_main_score(self) -> Self:
+    def only_main_score(self) -> TaskResult:
         """Return a new TaskResult object with only the main score.
 
         Returns:
             A new TaskResult object with only the main score.
         """
-        new_scores = {}
+        new_scores: dict[str, list[Score]] = {}
         for split in self.scores:
             new_scores[split] = []
             for subset_scores in self.scores[split]:
@@ -610,10 +614,12 @@ class TaskResult(BaseModel):
                     }
                 )
         new_res = {**self.to_dict(), "scores": new_scores}
-        new_res = TaskResult.from_validated(**new_res)
-        return new_res
+        return TaskResult.from_validated(**new_res)
 
-    def validate_and_filter_scores(self, task: AbsTask | None = None) -> Self:
+    def validate_and_filter_scores(
+        self,
+        task: AbsTask | None = None,
+    ) -> TaskResult:
         """Validate and filter the scores against the task metadata.
 
         This ensures that the scores are correct for the given task, by removing any splits besides those specified in the task metadata.
@@ -635,22 +641,32 @@ class TaskResult(BaseModel):
         splits = task.eval_splits
         hf_subsets = set(task.hf_subsets)  # Convert to set once
 
-        new_scores = {}
+        new_scores: dict[str, list[Score]] = {}
         seen_splits = set()
         for split in self.scores:
             if split not in splits:
                 continue
             seen_subsets = set()
-            # Use list comprehension for better performance
-            new_scores[split] = [
-                _scores
-                for _scores in self.scores[split]
-                if _scores["hf_subset"] in hf_subsets
-            ]
+            if task.is_aggregate:
+                # aggregate tasks only have the default subset, but in metadata can be multiple
+                new_scores[split] = [
+                    _scores
+                    for _scores in self.scores[split]
+                    if _scores["hf_subset"] == "default"
+                ]
+                seen_subsets = {"default"}
+            else:
+                new_scores[split] = [
+                    _scores
+                    for _scores in self.scores[split]
+                    if _scores["hf_subset"] in hf_subsets
+                ]
             for _scores in new_scores[split]:
                 seen_subsets.add(_scores["hf_subset"])
 
-            if seen_subsets != hf_subsets:
+            if seen_subsets != hf_subsets and not (
+                task.is_aggregate and "default" in seen_subsets
+            ):
                 missing_subsets = hf_subsets - seen_subsets
                 if len(missing_subsets) > 2:
                     subset1, subset2 = list(missing_subsets)[:2]
@@ -658,14 +674,36 @@ class TaskResult(BaseModel):
                 else:
                     missing_subsets_str = str(missing_subsets)
 
-                logger.warning(
-                    f"{task.metadata.name}: Missing subsets {missing_subsets_str} for split {split}"
-                )
+                msg = f"{task.metadata.name}: Missing subsets {missing_subsets_str} for split {split}"
+                logger.warning(msg)
+                warnings.warn(msg)
+                for missing_subset in missing_subsets:
+                    new_scores[split].append(
+                        {
+                            "hf_subset": missing_subset,
+                            "main_score": np.nan,
+                            "languages": task.metadata.hf_subsets_to_langscripts.get(
+                                missing_subset, []
+                            ),
+                        }
+                    )
             seen_splits.add(split)
         if seen_splits != set(splits):
-            logger.warning(
-                f"{task.metadata.name}: Missing splits {set(splits) - seen_splits}"
-            )
+            msg = f"{task.metadata.name}: Missing splits {set(splits) - seen_splits}"
+            logger.warning(msg)
+            warnings.warn(msg)
+            for missing_split in set(splits) - seen_splits:
+                new_scores[missing_split] = []
+                for missing_subset in hf_subsets:
+                    new_scores[missing_split].append(
+                        {
+                            "hf_subset": missing_subset,
+                            "main_score": np.nan,
+                            "languages": task.metadata.hf_subsets_to_langscripts.get(
+                                missing_subset, []
+                            ),
+                        }
+                    )
         data = self.model_dump()
         data["scores"] = new_scores
         return type(self).model_construct(**data)
@@ -736,7 +774,7 @@ class TaskResult(BaseModel):
             "mteb_version",
             "dataset_revision",
         ],
-    ) -> Self:
+    ) -> TaskResult:
         """Merges two TaskResult objects.
 
         Args:

mteb/similarity_functions.py

@@ -1,8 +1,14 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import torch
 
-from mteb.models import EncoderProtocol
 from mteb.models.model_meta import ScoringFunction
-from mteb.types import Array
+
+if TYPE_CHECKING:
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array
 
 
 def _use_torch_compile():
@@ -186,7 +192,7 @@ def max_sim(a: Array, b: Array) -> torch.Tensor:
         b,
     )
 
-    return scores.max(axis=-1).values.sum(axis=-1)
+    return scores.max(axis=-1).values.sum(axis=-1)  # type: ignore[call-overload]
 
 
 # https://github.com/lightonai/pylate/blob/2d094a724866d6e15701781528368438081c0157/pylate/scores/scores.py#L67C1-L122C38
@@ -217,7 +223,7 @@ def pairwise_max_sim(
             document_embedding,
         )
 
-        scores.append(query_document_score.max(axis=-1).values.sum())
+        scores.append(query_document_score.max(axis=-1).values.sum())  # type: ignore[call-overload]
 
     return torch.stack(scores, dim=0)
 
@@ -317,11 +323,15 @@ def similarity(text_embeddings: Array, input_embeddings: Array) -> Array:
     Returns:
         Matrix with similarities
     """
-    text_embeddings = _convert_to_tensor(text_embeddings)
-    input_embeddings = _convert_to_tensor(input_embeddings)
+    text_embeddings_tensor = _convert_to_tensor(text_embeddings)
+    input_embeddings_tensor = _convert_to_tensor(input_embeddings)
 
-    text_embeddings = text_embeddings / text_embeddings.norm(dim=-1, keepdim=True)
-    input_embeddings = input_embeddings / input_embeddings.norm(dim=-1, keepdim=True)
-    logits = torch.matmul(input_embeddings, text_embeddings.T)
+    text_embeddings_tensor = text_embeddings_tensor / text_embeddings_tensor.norm(
+        dim=-1, keepdim=True
+    )
+    input_embeddings_tensor = input_embeddings_tensor / input_embeddings_tensor.norm(
+        dim=-1, keepdim=True
+    )
+    logits = torch.matmul(input_embeddings_tensor, text_embeddings_tensor.T)
     probs = (logits * 100).softmax(dim=-1)
     return probs

mteb/tasks/aggregated_tasks/eng/cqadupstack_retrieval.py

@@ -1,5 +1,5 @@
-from mteb.abstasks import AbsTask
-from mteb.abstasks.aggregated_task import AbsTaskAggregate, AggregateTaskMetadata
+from mteb.abstasks.aggregate_task_metadata import AggregateTaskMetadata
+from mteb.abstasks.aggregated_task import AbsTaskAggregate
 from mteb.tasks.retrieval import (
     CQADupstackAndroidRetrieval,
     CQADupstackEnglishRetrieval,
@@ -15,7 +15,7 @@ from mteb.tasks.retrieval import (
     CQADupstackWordpressRetrieval,
 )
 
-task_list_cqa: list[AbsTask] = [
+task_list_cqa = [
     CQADupstackAndroidRetrieval(),
     CQADupstackEnglishRetrieval(),
     CQADupstackGamingRetrieval(),

mteb/tasks/aggregated_tasks/eng/sts17_multilingual_visual_sts_eng.py

@@ -1,10 +1,10 @@
-from mteb.abstasks.abstask import AbsTask
-from mteb.abstasks.aggregated_task import AbsTaskAggregate, AggregateTaskMetadata
+from mteb.abstasks.aggregate_task_metadata import AggregateTaskMetadata
+from mteb.abstasks.aggregated_task import AbsTaskAggregate
 from mteb.tasks.sts.multilingual.sts17_multilingual_visual_sts import (
     STS17MultilingualVisualSTS,
 )
 
-task_list_sts17: list[AbsTask] = [
+task_list_sts17 = [
     STS17MultilingualVisualSTS().filter_languages(
         languages=["eng"], hf_subsets=["en-en"]
     )