mteb 2.5.3__py3-none-any.whl → 2.5.4__py3-none-any.whl

mteb/cli/generate_model_card.py

@@ -1,5 +1,6 @@
 import logging
 import warnings
+from collections.abc import Sequence
 from pathlib import Path
 
 from huggingface_hub import ModelCard, ModelCardData, repo_exists
@@ -13,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 def generate_model_card(
     model_name: str,
-    tasks: list[AbsTask] | None = None,
+    tasks: Sequence[AbsTask] | None = None,
     existing_model_card_id_or_path: str | Path | None = None,
     results_cache: ResultCache = ResultCache(),
     output_path: Path = Path("model_card.md"),
@@ -48,8 +49,8 @@ def generate_model_card(
         for task_result in models_results.task_results:
             eval_results.extend(task_result.get_hf_eval_results())
 
-    existing_model_card_data = (
-        existing_model_card.data if existing_model_card else ModelCardData()
+    existing_model_card_data: ModelCardData = (
+        existing_model_card.data if existing_model_card else ModelCardData()  # type: ignore[assignment]
     )
 
     if existing_model_card_data.eval_results is None:
@@ -89,7 +90,8 @@ def generate_model_card(
             benchmark_results, existing_model_card
         )
 
-    if push_to_hub:
+    if push_to_hub and existing_model_card_id_or_path:
+        existing_model_card_id_or_path = str(existing_model_card_id_or_path)
         if repo_exists(existing_model_card_id_or_path):
             existing_model_card.push_to_hub(existing_model_card_id_or_path, token=token)
         else:

mteb/deprecated_evaluator.py

@@ -6,23 +6,23 @@ import os
 import sys
 import traceback
 import warnings
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from copy import deepcopy
 from datetime import datetime
 from itertools import chain
 from pathlib import Path
 from time import time
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import datasets
 
 import mteb
 from mteb.abstasks import AbsTask
+from mteb.abstasks.aggregated_task import AbsTaskAggregate
 from mteb.abstasks.task_metadata import TaskCategory, TaskType
 from mteb.benchmarks import Benchmark
 from mteb.models import (
     CrossEncoderWrapper,
-    EncoderProtocol,
     ModelMeta,
     MTEBModels,
     SentenceTransformerEncoderWrapper,
@@ -53,7 +53,7 @@ class MTEB:
     )
     def __init__(
         self,
-        tasks: Iterable[AbsTask | Benchmark],
+        tasks: Iterable[AbsTask] | Iterable[Benchmark],
         *,
         err_logs_path: str = "error_logs.txt",
     ) -> None:
@@ -64,15 +64,14 @@ class MTEB:
                 `mteb.get_tasks(["task1","task2"]) or `mteb.get_benchmark("MTEB(eng, classic)").
             err_logs_path: Path to save error logs.
         """
-        from mteb.benchmarks import Benchmark
-
-        self.tasks = list(tasks)
-        if len(self.tasks) > 0 and isinstance(self.tasks[0], Benchmark):
+        if isinstance(next(iter(tasks)), Benchmark):
             self.benchmarks = tasks
-            self.tasks = list(chain.from_iterable(self.tasks))
+            self.tasks = list(chain.from_iterable(cast(Iterable[Benchmark], tasks)))
+        elif isinstance(next(iter(tasks)), AbsTask):
+            self.tasks = list(cast(Iterable[AbsTask], tasks))
 
         self.err_logs_path = Path(err_logs_path)
-        self.last_evaluated_splits = {}
+        self._last_evaluated_splits: dict[str, list[str]] = {}
 
     @property
     def available_tasks(self) -> list[str]:
@@ -85,7 +84,7 @@ class MTEB:
         return sorted({x.metadata.type for x in self.tasks})
 
     @property
-    def available_task_categories(self) -> set[TaskCategory]:
+    def available_task_categories(self) -> set[TaskCategory | None]:
         """Set of available task categories."""
         return {x.metadata.category for x in self.tasks}
 
@@ -232,13 +231,14 @@ class MTEB:
         merged_kg_co2_emissions = None
         if existing_kg_co2_emissions and new_kg_co2_emissions:
             merged_kg_co2_emissions = existing_kg_co2_emissions + new_kg_co2_emissions
+        existing_evaluation_time = existing_results.evaluation_time or 0
+        new_evaluation_time = new_results.evaluation_time or 0
         merged_results = TaskResult(
             dataset_revision=new_results.dataset_revision,
             task_name=new_results.task_name,
             mteb_version=new_results.mteb_version,
             scores=merged_scores,
-            evaluation_time=existing_results.evaluation_time
-            + new_results.evaluation_time,
+            evaluation_time=existing_evaluation_time + new_evaluation_time,
             kg_co2_emissions=merged_kg_co2_emissions,
         )
 
@@ -307,13 +307,16 @@ class MTEB:
         elif verbosity == 3:
             datasets.logging.set_verbosity(logging.DEBUG)
 
-        meta = self.create_model_meta(model)
-        output_path = self._create_output_folder(meta, output_folder)
-
+        mteb_model: MTEBModels
         if isinstance(model, SentenceTransformer):
-            model = SentenceTransformerEncoderWrapper(model)
+            mteb_model = SentenceTransformerEncoderWrapper(model)
         elif isinstance(model, CrossEncoder):
-            model = CrossEncoderWrapper(model)
+            mteb_model = CrossEncoderWrapper(model)
+        else:
+            mteb_model = cast(MTEBModels, model)
+
+        meta = self.create_model_meta(mteb_model)
+        output_path = self._create_output_folder(meta, output_folder)
 
         # Disable co2_tracker for API models
         if "API" in meta.framework:
@@ -334,7 +337,7 @@ class MTEB:
         )  # save them in case we re-use the object (e.g. for reranking)
 
         # To evaluate missing splits, we keep track of the task name and the corresponding splits.
-        self.last_evaluated_splits = {}
+        self._last_evaluated_splits = {}
 
         while len(self.tasks) > 0:
             task = self.tasks[0]
@@ -343,9 +346,10 @@ class MTEB:
             )
 
             if task.is_aggregate:
-                self_ = MTEB(tasks=task.metadata.tasks)
-                task_results = self_.run(
-                    model,
+                aggregated_task = cast(AbsTaskAggregate, task)
+                self_ = MTEB(tasks=aggregated_task.metadata.tasks)
+                aggregated_task_results = self_.run(
+                    mteb_model,
                     verbosity=verbosity - 1,
                     output_folder=output_folder,
                     eval_splits=eval_splits,
@@ -356,12 +360,15 @@ class MTEB:
                     encode_kwargs=encode_kwargs,
                     **kwargs,
                 )
-                new_results = task.combine_task_results(task_results)
+                new_results = aggregated_task.combine_task_results(
+                    aggregated_task_results
+                )
                 evaluation_results.append(new_results)
 
                 if output_path:
-                    save_path = output_path / f"{task.metadata.name}.json"
-                    new_results.to_disk(save_path)
+                    new_results.to_disk(
+                        output_path / f"{aggregated_task.metadata.name}.json"
+                    )
                 del self.tasks[0]
                 continue
 
@@ -383,7 +390,7 @@ class MTEB:
             task_subsets = task.hf_subsets
 
             existing_results = None
-            save_path = None
+            save_path: Path | None = None
             final_splits_to_run = task_eval_splits
             missing_evaluations = self._get_missing_evaluations(
                 existing_results,
@@ -433,7 +440,7 @@ class MTEB:
                     logger.info(
                         f"No splits to evaluate for {task.metadata.name}. Skipping evaluation."
                     )
-                self.last_evaluated_splits[task.metadata.name] = []
+                self._last_evaluated_splits[task.metadata.name] = []
                 del self.tasks[0]
                 continue
 
@@ -441,11 +448,11 @@ class MTEB:
                 task.check_if_dataset_is_superseded()
                 task.load_data()
 
-                task_results = {}
+                task_results: dict[str, dict[str, dict[str, Any]]] = {}
                 evaluation_time = 0
                 kg_co2_emissions: int | None = 0 if co2_tracker else None
 
-                self.last_evaluated_splits[task.metadata.name] = []
+                self._last_evaluated_splits[task.metadata.name] = []
 
                 for split in final_splits_to_run:
                     info = missing_evaluations[split]
@@ -466,7 +473,9 @@ class MTEB:
 
                     if co2_tracker:
                         try:
-                            from codecarbon import EmissionsTracker
+                            from codecarbon import (  # type: ignore[import-untyped]
+                                EmissionsTracker,
+                            )
                         except ImportError:
                             raise ImportError(
                                 "codecarbon is not installed. Please install it using `pip install 'mteb[codecarbon]'` to track CO₂ emissions."
@@ -482,7 +491,7 @@ class MTEB:
                         ) as tracker:
                             results, tick, tock = self._run_eval(
                                 task,
-                                model,
+                                mteb_model,
                                 split,
                                 encode_kwargs=encode_kwargs,
                                 subsets_to_run=subsets_to_run,
@@ -495,7 +504,7 @@ class MTEB:
                     else:
                         results, tick, tock = self._run_eval(
                             task,
-                            model,
+                            mteb_model,
                             split,
                             subsets_to_run=subsets_to_run,
                             encode_kwargs=encode_kwargs,
@@ -511,25 +520,25 @@ class MTEB:
                     if verbosity >= 1:
                         logger.info(f"Scores: {task_results[split]}")
 
-                    self.last_evaluated_splits[task.metadata.name].append(split)
+                    self._last_evaluated_splits[task.metadata.name].append(split)
 
                 # Create new TaskResult
                 new_results = TaskResult.from_task_results(
                     task,
-                    task_results,
+                    task_results,  # type: ignore[arg-type]
                     evaluation_time=evaluation_time,
                     kg_co2_emissions=kg_co2_emissions,
                 )
 
                 # Merge with existing if needed
-                if output_path and save_path.exists():
+                if output_path and save_path and save_path.exists():
                     existing_results = TaskResult.from_disk(save_path)
                 if existing_results:
                     merged_results = self._merge_results(existing_results, new_results)
                 else:
                     merged_results = new_results
 
-                if output_path:
+                if output_path and save_path:
                     merged_results.to_disk(save_path)
 
                 evaluation_results.append(merged_results)
@@ -556,7 +565,7 @@ class MTEB:
     def create_model_meta(model: MTEBModels) -> ModelMeta:
         """Create a ModelMeta object for the given model."""
         if hasattr(model, "mteb_model_meta") and model.mteb_model_meta is not None:
-            meta = model.mteb_model_meta  # type: ignore
+            meta = model.mteb_model_meta
         else:
             meta = MTEB._get_model_meta(model)
 
@@ -582,7 +591,11 @@ class MTEB:
         if output_folder is None:
             return None
 
-        model_revision: str = model_meta.revision  # type: ignore
+        model_revision: str = (
+            model_meta.revision
+            if model_meta.revision is not None
+            else "no_revision_available"
+        )
         model_path_name = model_meta.model_name_as_path()
 
         output_path = Path(output_folder) / model_path_name / model_revision
@@ -604,15 +617,15 @@ class MTEB:
              Tasks with empty lists indicate that results already existed and no splits were evaluated.
         """
         return deepcopy(
-            {task: list(splits) for task, splits in self.last_evaluated_splits.items()}
+            {task: list(splits) for task, splits in self._last_evaluated_splits.items()}
         )
 
     @staticmethod
     def _get_missing_evaluations(
         existing_results: TaskResult | None,
-        task_eval_splits: list[str],
-        task_eval_langs: list[str],
-        eval_subsets: list[str] | None,
+        task_eval_splits: Sequence[str],
+        task_eval_langs: Sequence[str],
+        eval_subsets: Sequence[str] | None,
     ) -> dict[str, dict[str, Any]]:
         """Return a dictionary for each split, indicating if the whole split is missing and which subsets are missing."""
         missing_evaluations = {
@@ -661,7 +674,7 @@ class MTEB:
         return missing_evaluations
 
     @staticmethod
-    def _get_model_meta(model: EncoderProtocol) -> ModelMeta:
+    def _get_model_meta(model: MTEBModels) -> ModelMeta:
         from sentence_transformers import CrossEncoder, SentenceTransformer
 
         if isinstance(model, CrossEncoder):

mteb/evaluate.py

@@ -14,11 +14,10 @@ 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 (
@@ -58,27 +57,26 @@ def _sanitize_model(
 ) -> 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 = ModelMeta.from_hub(None)
+            meta = ModelMeta._from_hub(None)
+        wrapped_model = cast(MTEBModels | ModelMeta, model)
     else:
-        meta = ModelMeta.from_hub(None) 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(
@@ -124,7 +122,8 @@ def _evaluate_task(
                 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 = {}
@@ -150,7 +149,7 @@ def _evaluate_task(
             if public_only is False:
                 raise e
 
-    evaluation_time = 0
+    evaluation_time = 0.0
 
     for split, hf_subsets in splits.items():
         tick = time()
@@ -197,12 +196,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))
@@ -335,10 +340,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,
@@ -348,17 +353,18 @@ def evaluate(
             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,
@@ -379,23 +385,23 @@ def evaluate(
                 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

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