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

mteb/evaluate.py

@@ -1,23 +1,23 @@
 from __future__ import annotations
 
 import logging
+import warnings
 from collections.abc import Iterable
-from copy import deepcopy
 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
 
 from mteb._helpful_enum import HelpfulStrEnum
 from mteb.abstasks import AbsTaskRetrieval
 from mteb.abstasks.abstask import AbsTask
 from mteb.abstasks.aggregated_task import AbsTaskAggregate
+from mteb.benchmarks.benchmark import Benchmark
 from mteb.cache import ResultCache
 from mteb.models.model_meta import ModelMeta
 from mteb.models.models_protocols import (
-    CrossEncoderProtocol,
-    EncoderProtocol,
     MTEBModels,
 )
 from mteb.models.sentence_transformer_wrapper import (
@@ -25,7 +25,9 @@ from mteb.models.sentence_transformer_wrapper import (
     SentenceTransformerEncoderWrapper,
 )
 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:
@@ -51,62 +53,31 @@ class OverwriteStrategy(HelpfulStrEnum):
     ONLY_CACHE = "only-cache"
 
 
-_empty_model_meta = ModelMeta(
-    loader=None,
-    name=None,
-    revision=None,
-    release_date=None,
-    languages=None,
-    framework=[],
-    similarity_fn_name=None,
-    n_parameters=None,
-    memory_usage_mb=None,
-    max_tokens=None,
-    embed_dim=None,
-    license=None,
-    open_weights=None,
-    public_training_code=None,
-    public_training_data=None,
-    use_instructions=None,
-    training_datasets=None,
-    modalities=[],
-)
-
-
-def _create_empty_model_meta() -> ModelMeta:
-    logger.warning("Model metadata is missing. Using empty metadata.")
-    meta = deepcopy(_empty_model_meta)
-    meta.revision = "no_revision_available"
-    meta.name = "no_model_name_available"
-    return meta
-
-
 def _sanitize_model(
     model: ModelMeta | MTEBModels | SentenceTransformer | CrossEncoder,
 ) -> tuple[MTEBModels | ModelMeta, ModelMeta, ModelName, Revision]:
     from sentence_transformers import CrossEncoder, SentenceTransformer
 
+    wrapped_model: MTEBModels | ModelMeta
     if isinstance(model, SentenceTransformer):
-        _mdl = SentenceTransformerEncoderWrapper(model)
-        meta = _mdl.mteb_model_meta
-        _mdl = cast(EncoderProtocol, _mdl)
-        model = _mdl
+        wrapped_model = SentenceTransformerEncoderWrapper(model)
+        meta = wrapped_model.mteb_model_meta
     elif isinstance(model, CrossEncoder):
-        _mdl = CrossEncoderWrapper(model)
-        _mdl = cast(CrossEncoderProtocol, _mdl)
-        meta = _mdl.mteb_model_meta
-        model = _mdl
+        wrapped_model = CrossEncoderWrapper(model)
+        meta = wrapped_model.mteb_model_meta
     elif hasattr(model, "mteb_model_meta"):
-        meta = model.mteb_model_meta  # type: ignore[attr-defined]
+        meta = getattr(model, "mteb_model_meta")
         if not isinstance(meta, ModelMeta):
-            meta = _create_empty_model_meta()
+            meta = ModelMeta._from_hub(None)
+        wrapped_model = cast(MTEBModels | ModelMeta, model)
     else:
-        meta = _create_empty_model_meta() if not isinstance(model, ModelMeta) else model
+        meta = ModelMeta._from_hub(None) if not isinstance(model, ModelMeta) else model
+        wrapped_model = meta
 
     model_name = cast(str, meta.name)
     model_revision = cast(str, meta.revision)
 
-    return model, meta, model_name, model_revision
+    return wrapped_model, meta, model_name, model_revision
 
 
 def _evaluate_task(
@@ -115,9 +86,11 @@ def _evaluate_task(
     *,
     splits: dict[SplitName, list[HFSubset]],
     co2_tracker: bool | None,
-    encode_kwargs: dict[str, Any],
+    encode_kwargs: EncodeKwargs,
     prediction_folder: Path | None,
-) -> TaskResult:
+    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.
 
     Returns:
@@ -149,8 +122,10 @@ def _evaluate_task(
                 encode_kwargs=encode_kwargs,
                 co2_tracker=False,
                 prediction_folder=prediction_folder,
+                public_only=public_only,
             )
-        result.kg_co2_emissions = tracker.final_emissions
+        if isinstance(result, TaskResult):
+            result.kg_co2_emissions = tracker.final_emissions
         return result
 
     task_results = {}
@@ -159,9 +134,24 @@ def _evaluate_task(
 
     data_loaded = task.data_loaded
     if not data_loaded:
-        task.load_data()
+        try:
+            task.load_data()
+        except DatasetNotFoundError as e:
+            if not task.metadata.is_public and public_only is None:
+                msg = (
+                    f"Dataset for private task '{task.metadata.name}' not found. "
+                    "Make sure you have access to the dataset and that you have set up the authentication correctly. To disable this warning set `public_only=False`"
+                )
+                logger.warning(msg)
+                warnings.warn(msg)
+                return TaskError(
+                    task_name=task.metadata.name,
+                    exception=str(e),
+                )
+            if public_only is False:
+                raise e
 
-    evaluation_time = 0
+    evaluation_time = 0.0
 
     for split, hf_subsets in splits.items():
         tick = time()
@@ -208,12 +198,18 @@ def _check_model_modalities(
         return
 
     model_modalities = set(model.modalities)
+    check_tasks: Iterable[AbsTask] = []
     if isinstance(tasks, AbsTask):
-        tasks = [tasks]
+        check_tasks = [tasks]
+    elif isinstance(tasks, Benchmark):
+        benchmark = cast(Benchmark, tasks)
+        check_tasks = benchmark.tasks
+    else:
+        check_tasks = cast(Iterable[AbsTask], tasks)
 
     warnings, errors = [], []
 
-    for task in tasks:
+    for task in check_tasks:
         # only retrieval tasks have different modalities for query and document and can be run with partial overlaps
         if isinstance(task, AbsTaskRetrieval):
             query_mods = set(task.metadata.get_modalities(PromptType.query))
@@ -256,17 +252,32 @@ def _check_model_modalities(
         logger.warning(msg)
 
 
+def _requires_merge(task: AbsTask, existing_results: TaskResult) -> bool:
+    """Check if the existing results require merging with new results."""
+    # If the task has multiple eval splits and existing results cover only a subset, we need to merge
+    required_evals = dict.fromkeys(task.eval_splits, task.hf_subsets)
+    for split, subsets in required_evals.items():
+        res = existing_results.scores.get(split, None)
+        if res is None:
+            return True
+        hf_subsets = [r["hf_subset"] for r in res]
+        if not set(subsets).issubset(set(hf_subsets)):
+            return True
+    return False
+
+
 def evaluate(
     model: ModelMeta | MTEBModels | SentenceTransformer | CrossEncoder,
     tasks: AbsTask | Iterable[AbsTask],
     *,
     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,
     show_progress_bar: bool = True,
+    public_only: bool | None = None,
 ) -> ModelResult:
     """This function runs a model on a given task and returns the results.
 
@@ -290,6 +301,7 @@ def evaluate(
         prediction_folder: Optional folder in which to save model predictions for the task. Predictions of the tasks will be sabed in `prediction_folder/{task_name}_predictions.json`
         show_progress_bar: Whether to show a progress bar when running the evaluation. Default is True. Setting this to False will also set the
             `encode_kwargs['show_progress_bar']` to False if encode_kwargs is unspecified.
+        public_only: Run only public tasks. If None, it will attempt to run the private task.
 
     Returns:
         The results of the evaluation.
@@ -330,10 +342,10 @@ def evaluate(
 
     # AbsTaskAggregate is a special case where we have to run multiple tasks and combine the results
     if isinstance(tasks, AbsTaskAggregate):
-        task = cast(AbsTaskAggregate, tasks)
+        aggregated_task = cast(AbsTaskAggregate, tasks)
         results = evaluate(
             model,
-            task.metadata.tasks,
+            aggregated_task.metadata.tasks,
             co2_tracker=co2_tracker,
             raise_error=raise_error,
             encode_kwargs=encode_kwargs,
@@ -341,18 +353,21 @@ def evaluate(
             overwrite_strategy=overwrite_strategy,
             prediction_folder=prediction_folder,
             show_progress_bar=show_progress_bar,
+            public_only=public_only,
         )
-        result = task.combine_task_results(results.task_results)
+        combined_results = aggregated_task.combine_task_results(results.task_results)
         return ModelResult(
             model_name=results.model_name,
             model_revision=results.model_revision,
-            task_results=[result],
+            task_results=[combined_results],
         )
 
     if isinstance(tasks, AbsTask):
         task = tasks
     else:
-        results = []
+        tasks = cast(Iterable[AbsTask], tasks)
+        evaluate_results = []
+        exceptions = []
         tasks_tqdm = tqdm(
             tasks,
             desc="Evaluating tasks",
@@ -370,31 +385,40 @@ def evaluate(
                 overwrite_strategy=overwrite_strategy,
                 prediction_folder=prediction_folder,
                 show_progress_bar=False,
+                public_only=public_only,
             )
-            results.extend(_res.task_results)
+            evaluate_results.extend(_res.task_results)
+            if _res.exceptions:
+                exceptions.extend(_res.exceptions)
         return ModelResult(
             model_name=_res.model_name,
             model_revision=_res.model_revision,
-            task_results=results,
+            task_results=evaluate_results,
+            exceptions=exceptions,
         )
 
     overwrite_strategy = OverwriteStrategy.from_str(overwrite_strategy)
 
-    existing_results = None
+    existing_results: TaskResult | None = None
     if cache and overwrite_strategy != OverwriteStrategy.ALWAYS:
-        results = cache.load_task_result(task.metadata.name, meta)
-        if results:
-            existing_results = results
+        cache_results = cache.load_task_result(task.metadata.name, meta)
+        if cache_results:
+            existing_results = cache_results
 
     if (
         existing_results
-        and overwrite_strategy == "only-missing"
-        and overwrite_strategy == OverwriteStrategy.ONLY_MISSING
-        and existing_results.is_mergeable(task)
+        and overwrite_strategy
+        not in (OverwriteStrategy.ALWAYS, OverwriteStrategy.NEVER)
+        and (
+            not _requires_merge(task, existing_results)
+            or existing_results.is_mergeable(task)
+        )
     ):
         missing_eval = existing_results.get_missing_evaluations(task)
     else:
         missing_eval = dict.fromkeys(task.eval_splits, task.hf_subsets)
+        # Will be fully recomputed so we set it to None to avoid merging:
+        existing_results = None
 
     if (
         existing_results
@@ -415,12 +439,13 @@ def evaluate(
         OverwriteStrategy.ONLY_CACHE,
     ]:
         raise ValueError(
-            f"overwrite_strategy is set to '{overwrite_strategy.value}' and the results file exists. However there are the following missing splits (and subsets): {missing_eval}. To rerun these set overwrite_strategy to 'only-missing'."
+            f"overwrite_strategy is set to '{overwrite_strategy.value}' and the results file exists for task {task.metadata.name}. "
+            + f"However there are the following missing splits (and subsets): {missing_eval}. To rerun these set overwrite_strategy to 'only-missing'."
         )
 
     if existing_results:
         logger.info(
-            f"Found existing results for {task.metadata.name}, only running missing splits: {list(missing_eval.keys())}"
+            f"Found existing results for {task.metadata.name}, only running missing splits (subsets): {missing_eval}"
         )
 
     if isinstance(model, ModelMeta):
@@ -439,16 +464,13 @@ def evaluate(
                 co2_tracker=co2_tracker,
                 encode_kwargs=encode_kwargs,
                 prediction_folder=prediction_folder,
+                public_only=public_only,
             )
         except Exception as e:
             logger.error(
                 f"Error while running task {task.metadata.name} on splits {list(missing_eval.keys())}: {e}"
             )
-            return ModelResult(
-                model_name=model_name,
-                model_revision=model_revision,
-                task_results=[],
-            )
+            result = TaskError(task_name=task.metadata.name, exception=str(e))
     else:
         result = _evaluate_task(
             model=model,
@@ -457,9 +479,18 @@ def evaluate(
             co2_tracker=False,
             encode_kwargs=encode_kwargs,
             prediction_folder=prediction_folder,
+            public_only=public_only,
         )
     logger.info(f"✓ Finished evaluation for {task.metadata.name}")
 
+    if isinstance(result, TaskError):
+        return ModelResult(
+            model_name=model_name,
+            model_revision=model_revision,
+            task_results=[],
+            exceptions=[result],
+        )
+
     if existing_results:
         result = result.merge(existing_results)
 

mteb/filter_tasks.py

@@ -1,7 +1,7 @@
 """This script contains functions that are used to get an overview of the MTEB benchmark."""
 
 import logging
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from typing import overload
 
 from mteb.abstasks import (
@@ -34,14 +34,14 @@ def _check_is_valid_language(lang: str) -> None:
 
 @overload
 def filter_tasks(
-    tasks: Sequence[AbsTask],
+    tasks: Iterable[AbsTask],
     *,
-    languages: list[str] | None = None,
-    script: list[str] | None = None,
-    domains: list[TaskDomain] | None = None,
-    task_types: list[TaskType] | None = None,  # type: ignore
-    categories: list[TaskCategory] | None = None,
-    modalities: list[Modalities] | None = None,
+    languages: Sequence[str] | None = None,
+    script: Sequence[str] | None = None,
+    domains: Iterable[TaskDomain] | None = None,
+    task_types: Iterable[TaskType] | None = None,
+    categories: Iterable[TaskCategory] | None = None,
+    modalities: Iterable[Modalities] | None = None,
     exclusive_modality_filter: bool = False,
     exclude_superseded: bool = False,
     exclude_aggregate: bool = False,
@@ -51,14 +51,14 @@ def filter_tasks(
 
 @overload
 def filter_tasks(
-    tasks: Sequence[type[AbsTask]],
+    tasks: Iterable[type[AbsTask]],
     *,
-    languages: list[str] | None = None,
-    script: list[str] | None = None,
-    domains: list[TaskDomain] | None = None,
-    task_types: list[TaskType] | None = None,  # type: ignore
-    categories: list[TaskCategory] | None = None,
-    modalities: list[Modalities] | None = None,
+    languages: Sequence[str] | None = None,
+    script: Sequence[str] | None = None,
+    domains: Iterable[TaskDomain] | None = None,
+    task_types: Iterable[TaskType] | None = None,
+    categories: Iterable[TaskCategory] | None = None,
+    modalities: Iterable[Modalities] | None = None,
     exclusive_modality_filter: bool = False,
     exclude_superseded: bool = False,
     exclude_aggregate: bool = False,
@@ -67,14 +67,14 @@ def filter_tasks(
 
 
 def filter_tasks(
-    tasks: Sequence[AbsTask] | Sequence[type[AbsTask]],
+    tasks: Iterable[AbsTask] | Iterable[type[AbsTask]],
     *,
-    languages: list[str] | None = None,
-    script: list[str] | None = None,
-    domains: list[TaskDomain] | None = None,
-    task_types: list[TaskType] | None = None,  # type: ignore
-    categories: list[TaskCategory] | None = None,
-    modalities: list[Modalities] | None = None,
+    languages: Sequence[str] | None = None,
+    script: Sequence[str] | None = None,
+    domains: Iterable[TaskDomain] | None = None,
+    task_types: Iterable[TaskType] | None = None,
+    categories: Iterable[TaskCategory] | None = None,
+    modalities: Iterable[Modalities] | None = None,
     exclusive_modality_filter: bool = False,
     exclude_superseded: bool = False,
     exclude_aggregate: bool = False,
@@ -92,7 +92,6 @@ def filter_tasks(
         task_types: A string specifying the type of task e.g. "Classification" or "Retrieval". If None, all tasks are included.
         categories: A list of task categories these include "t2t" (text to text), "t2i" (text to image). See TaskMetadata for the full list.
         exclude_superseded: A boolean flag to exclude datasets which are superseded by another.
-        eval_splits: A list of evaluation splits to include. If None, all splits are included.
         modalities: A list of modalities to include. If None, all modalities are included.
         exclusive_modality_filter: If True, only keep tasks where _all_ filter modalities are included in the
             task's modalities and ALL task modalities are in filter modalities (exact match).
@@ -113,12 +112,12 @@ def filter_tasks(
     """
     langs_to_keep = None
     if languages:
-        [_check_is_valid_language(lang) for lang in languages]
+        [_check_is_valid_language(lang) for lang in languages]  # type: ignore[func-returns-value]
         langs_to_keep = set(languages)
 
     script_to_keep = None
     if script:
-        [_check_is_valid_script(s) for s in script]
+        [_check_is_valid_script(s) for s in script]  # type: ignore[func-returns-value]
         script_to_keep = set(script)
 
     domains_to_keep = None
@@ -178,4 +177,4 @@ def filter_tasks(
 
         _tasks.append(t)
 
-    return _tasks
+    return _tasks  # type: ignore[return-value]  # type checker cannot infer the overload return type

mteb/get_tasks.py

@@ -2,8 +2,9 @@
 
 import difflib
 import logging
+import warnings
 from collections import Counter, defaultdict
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from typing import Any
 
 import pandas as pd
@@ -22,12 +23,11 @@ logger = logging.getLogger(__name__)
 def _gather_tasks() -> tuple[type[AbsTask], ...]:
     import mteb.tasks as tasks
 
-    tasks = [
+    return tuple(
         t
         for t in tasks.__dict__.values()
         if isinstance(t, type) and issubclass(t, AbsTask)
-    ]
-    return tuple(tasks)
+    )
 
 
 def _create_name_to_task_mapping(
@@ -43,7 +43,7 @@ def _create_name_to_task_mapping(
     return metadata_names
 
 
-def _create_similar_tasks(tasks: Sequence[type[AbsTask]]) -> dict[str, list[str]]:
+def _create_similar_tasks(tasks: Iterable[type[AbsTask]]) -> dict[str, list[str]]:
     """Create a dictionary of similar tasks.
 
     Returns:
@@ -194,9 +194,8 @@ class MTEBTasks(tuple[AbsTask]):
             string with a LaTeX table.
         """
         if include_citation_in_name and "name" in properties:
-            properties += ["intext_citation"]
-            df = self.to_dataframe(properties)
-            df["name"] = df["name"] + " " + df["intext_citation"]
+            df = self.to_dataframe(tuple(properties) + ("intext_citation",))
+            df["name"] = df["name"] + " " + df["intext_citation"]  # type: ignore[operator]
             df = df.drop(columns=["intext_citation"])
         else:
             df = self.to_dataframe(properties)
@@ -221,17 +220,17 @@ class MTEBTasks(tuple[AbsTask]):
 
 
 def get_tasks(
-    tasks: list[str] | None = None,
+    tasks: Sequence[str] | None = None,
     *,
-    languages: list[str] | None = None,
-    script: list[str] | None = None,
-    domains: list[TaskDomain] | None = None,
-    task_types: list[TaskType] | None = None,  # type: ignore
-    categories: list[TaskCategory] | None = None,
+    languages: Sequence[str] | None = None,
+    script: Sequence[str] | None = None,
+    domains: Sequence[TaskDomain] | None = None,
+    task_types: Sequence[TaskType] | None = None,
+    categories: Sequence[TaskCategory] | None = None,
     exclude_superseded: bool = True,
-    eval_splits: list[str] | None = None,
+    eval_splits: Sequence[str] | None = None,
     exclusive_language_filter: bool = False,
-    modalities: list[Modalities] | None = None,
+    modalities: Sequence[Modalities] | None = None,
     exclusive_modality_filter: bool = False,
     exclude_aggregate: bool = False,
     exclude_private: bool = True,
@@ -287,7 +286,7 @@ def get_tasks(
         ]
         return MTEBTasks(_tasks)
 
-    _tasks = filter_tasks(
+    tasks_: Sequence[type[AbsTask]] = filter_tasks(
         TASK_LIST,
         languages=languages,
         script=script,
@@ -300,12 +299,12 @@ def get_tasks(
         exclude_aggregate=exclude_aggregate,
         exclude_private=exclude_private,
     )
-    _tasks = [
-        cls().filter_languages(languages, script).filter_eval_splits(eval_splits)
-        for cls in _tasks
-    ]
-
-    return MTEBTasks(_tasks)
+    return MTEBTasks(
+        [
+            cls().filter_languages(languages, script).filter_eval_splits(eval_splits)
+            for cls in tasks_
+        ]
+    )
 
 
 _TASK_RENAMES = {"PersianTextTone": "SynPerTextToneClassification"}
@@ -313,10 +312,10 @@ _TASK_RENAMES = {"PersianTextTone": "SynPerTextToneClassification"}
 
 def get_task(
     task_name: str,
-    languages: list[str] | None = None,
-    script: list[str] | None = None,
-    eval_splits: list[str] | None = None,
-    hf_subsets: list[str] | None = None,
+    languages: Sequence[str] | None = None,
+    script: Sequence[str] | None = None,
+    eval_splits: Sequence[str] | None = None,
+    hf_subsets: Sequence[str] | None = None,
     exclusive_language_filter: bool = False,
 ) -> AbsTask:
     """Get a task by name.
@@ -340,9 +339,9 @@ def get_task(
     """
     if task_name in _TASK_RENAMES:
         _task_name = _TASK_RENAMES[task_name]
-        logger.warning(
-            f"The task with the given name '{task_name}' has been renamed to '{_task_name}'. To prevent this warning use the new name."
-        )
+        msg = f"The task with the given name '{task_name}' has been renamed to '{_task_name}'. To prevent this warning use the new name."
+        logger.warning(msg)
+        warnings.warn(msg)
 
     if task_name not in _TASKS_REGISTRY:
         close_matches = difflib.get_close_matches(task_name, _TASKS_REGISTRY.keys())

mteb/languages/language_scripts.py

@@ -1,9 +1,9 @@
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 
 from typing_extensions import Self
 
-from mteb.languages import check_language_code
+from mteb.languages.check_language_code import check_language_code
 
 
 @dataclass
@@ -25,7 +25,9 @@ class LanguageScripts:
 
     @classmethod
     def from_languages_and_scripts(
-        cls, languages: list[str] | None = None, scripts: list[str] | None = None
+        cls,
+        languages: Sequence[str] | None = None,
+        scripts: Sequence[str] | None = None,
     ) -> Self:
         """Create a LanguageScripts object from lists of languages and scripts.