mteb 2.7.1__py3-none-any.whl → 2.7.3__py3-none-any.whl

mteb/__init__.py

@@ -3,6 +3,7 @@ from importlib.metadata import version
 from mteb import types
 from mteb.abstasks import AbsTask
 from mteb.abstasks.task_metadata import TaskMetadata
+from mteb.cache import ResultCache
 from mteb.deprecated_evaluator import MTEB
 from mteb.evaluate import evaluate
 from mteb.filter_tasks import filter_tasks
@@ -33,6 +34,7 @@ __all__ = [
     "CrossEncoderProtocol",
     "EncoderProtocol",
     "IndexEncoderSearchProtocol",
+    "ResultCache",
     "SearchProtocol",
     "SentenceTransformerEncoderWrapper",
     "TaskMetadata",

mteb/_create_dataloaders.py

@@ -1,21 +1,28 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Callable
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from datasets import Dataset, Image
 from torch.utils.data import DataLoader, default_collate
 
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.types import (
-    BatchedInput,
-    Conversation,
     ConversationTurn,
     PromptType,
-    QueryDatasetType,
 )
-from mteb.types._encoder_io import CorpusInput, ImageInput, QueryInput, TextInput
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import (
+        BatchedInput,
+        Conversation,
+        QueryDatasetType,
+    )
+    from mteb.types._encoder_io import CorpusInput, ImageInput, QueryInput, TextInput
 
 logger = logging.getLogger(__name__)
 
@@ -128,7 +135,7 @@ def _convert_conv_history_to_query(
     conversation = row["text"]
     # if it's a list of strings, just join them
     if isinstance(conversation, list) and isinstance(conversation[0], str):
-        conversation_ = cast(list[str], conversation)
+        conversation_ = cast("list[str]", conversation)
         conv_str = "; ".join(conversation_)
         current_conversation = [
             ConversationTurn(role="user", content=message) for message in conversation_
@@ -173,7 +180,7 @@ def _convert_conv_history_to_query(
 
     row["text"] = conv_str
     row["conversation"] = current_conversation
-    return cast(dict[str, str | list[ConversationTurn]], row)
+    return cast("dict[str, str | list[ConversationTurn]]", row)
 
 
 def _create_dataloader_for_queries_conversation(

mteb/_evaluators/any_sts_evaluator.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
-from typing import TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
-from datasets import Dataset
 from sklearn.metrics.pairwise import (
     paired_cosine_distances,
     paired_euclidean_distances,
@@ -9,13 +10,17 @@ from sklearn.metrics.pairwise import (
 )
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import compute_pairwise_similarity
-from mteb.types import EncodeKwargs, PromptType
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs, PromptType
+
 logger = logging.getLogger(__name__)
 
 

mteb/_evaluators/clustering_evaluator.py

@@ -1,15 +1,21 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
-from datasets import Dataset
 from sklearn import cluster
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
-from mteb.types import EncodeKwargs
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 

mteb/_evaluators/evaluator.py

@@ -1,10 +1,15 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Mapping
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from mteb.abstasks.abstask import _set_seed
-from mteb.models import EncoderProtocol
-from mteb.types import EncodeKwargs
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs
 
 
 class Evaluator(ABC):

mteb/_evaluators/image/imagetext_pairclassification_evaluator.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -14,13 +13,16 @@ from mteb._create_dataloaders import (
 )
 from mteb._evaluators.evaluator import Evaluator
 from mteb._requires_package import requires_image_dependencies
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models.models_protocols import EncoderProtocol
-from mteb.types import EncodeKwargs
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from PIL.Image import Image
 
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import EncodeKwargs
+
 
 logger = logging.getLogger(__name__)
 

mteb/_evaluators/pair_classification_evaluator.py

@@ -1,8 +1,9 @@
+from __future__ import annotations
+
 import logging
-from typing import Any, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict
 
 import numpy as np
-from datasets import Dataset
 from sklearn.metrics.pairwise import (
     paired_cosine_distances,
     paired_euclidean_distances,
@@ -11,10 +12,14 @@ from sklearn.metrics.pairwise import (
 
 from mteb._create_dataloaders import _create_dataloader_from_texts, create_dataloader
 from mteb._evaluators.evaluator import Evaluator
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import compute_pairwise_similarity
-from mteb.types import EncodeKwargs, PromptType
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs, PromptType
 
 logger = logging.getLogger(__name__)
 

mteb/_evaluators/retrieval_evaluator.py

@@ -1,23 +1,29 @@
-import logging
-from collections.abc import Sequence
+from __future__ import annotations
 
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import SearchProtocol
-from mteb.types import (
-    CorpusDatasetType,
-    EncodeKwargs,
-    QueryDatasetType,
-    RelevantDocumentsType,
-    RetrievalEvaluationResult,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
-)
+import logging
+from typing import TYPE_CHECKING
 
 from .evaluator import Evaluator
 from .retrieval_metrics import (
     calculate_retrieval_scores,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import SearchProtocol
+    from mteb.types import (
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RelevantDocumentsType,
+        RetrievalEvaluationResult,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 

mteb/_evaluators/retrieval_metrics.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 from collections import defaultdict
-from collections.abc import Mapping
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pandas as pd
@@ -9,7 +10,12 @@ import pytrec_eval
 from packaging.version import Version
 from sklearn.metrics import auc
 
-from mteb.types import RelevantDocumentsType, RetrievalEvaluationResult
+from mteb.types import RetrievalEvaluationResult
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from mteb.types import RelevantDocumentsType
 
 logger = logging.getLogger(__name__)
 

mteb/_evaluators/sklearn_evaluator.py

@@ -1,18 +1,22 @@
-import logging
-from typing import Any, Protocol, cast
+from __future__ import annotations
 
-import numpy as np
-from datasets import Dataset
-from torch.utils.data import DataLoader
-from typing_extensions import Self
+import logging
+from typing import TYPE_CHECKING, Any, Protocol, cast
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
-from mteb.types import Array, BatchedInput, EncodeKwargs
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    import numpy as np
+    from datasets import Dataset
+    from torch.utils.data import DataLoader
+    from typing_extensions import Self
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, BatchedInput, EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 
@@ -104,7 +108,7 @@ class SklearnEvaluator(Evaluator):
                 hf_subset=self.hf_subset,
                 **encode_kwargs,
             )
-            test_cache = cast(Array, test_cache)
+            test_cache = cast("Array", test_cache)
 
         logger.info("Running - Fitting classifier...")
         y_train = self.train_dataset[self.label_column_name]

mteb/_evaluators/text/bitext_mining_evaluator.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 import torch
 from datasets import Dataset
@@ -6,9 +9,11 @@ from tqdm.auto import tqdm
 
 from mteb._create_dataloaders import _create_dataloader_from_texts
 from mteb._evaluators.evaluator import Evaluator
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
-from mteb.types import Array, EncodeKwargs
+
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, EncodeKwargs
 
 logger = logging.getLogger(__name__)
 

mteb/_evaluators/text/summarization_evaluator.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import logging
 import sys
-from typing import TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
 import numpy as np
 import torch
@@ -9,10 +11,12 @@ from tqdm.auto import tqdm
 
 from mteb._create_dataloaders import _create_dataloader_from_texts
 from mteb._evaluators.evaluator import Evaluator
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import cos_sim, dot_score
-from mteb.types import EncodeKwargs
+
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs
 
 # if later than python 3.13 use typing module
 if sys.version_info >= (3, 13):

mteb/_evaluators/zeroshot_classification_evaluator.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 from datasets import Dataset
 
@@ -6,13 +9,17 @@ from mteb._create_dataloaders import (
     _create_dataloader_from_texts,
     create_dataloader,
 )
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import similarity
-from mteb.types import Array, EncodeKwargs
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 

mteb/_helpful_enum.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
 from enum import Enum
+from typing import TYPE_CHECKING
 
-from typing_extensions import Self
+if TYPE_CHECKING:
+    from typing_extensions import Self
 
 
 class HelpfulStrEnum(str, Enum):

mteb/abstasks/_data_filter/filters.py

@@ -1,12 +1,18 @@
 """Simplified version of https://gist.github.com/AlexeyVatolin/ea3adc21aa7a767603ff393b22085adc from https://github.com/embeddings-benchmark/mteb/pull/2900"""
 
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 import datasets
 import pandas as pd
-from datasets import Dataset, DatasetDict
+from datasets import DatasetDict
+
+if TYPE_CHECKING:
+    from datasets import Dataset
 
-from mteb import TaskMetadata
+    from mteb import TaskMetadata
 
 logger = logging.getLogger(__name__)
 

mteb/abstasks/_data_filter/task_pipelines.py

@@ -1,9 +1,10 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 from datasets import DatasetDict
 
-from mteb import TaskMetadata
-from mteb.abstasks import AbsTaskClassification
 from mteb.abstasks._data_filter.filters import (
     deduplicate,
     filter_empty,
@@ -13,6 +14,10 @@ from mteb.abstasks._data_filter.filters import (
     split_train_test,
 )
 
+if TYPE_CHECKING:
+    from mteb import TaskMetadata
+    from mteb.abstasks import AbsTaskClassification
+
 logger = logging.getLogger(__name__)
 
 

mteb/abstasks/_statistics_calculation.py

@@ -2,10 +2,8 @@ from __future__ import annotations
 
 import hashlib
 from collections import Counter
-from collections.abc import Mapping
 from typing import TYPE_CHECKING, cast
 
-from mteb.types import TopRankedDocumentsType
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -16,8 +14,12 @@ from mteb.types.statistics import (
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     from PIL import Image
 
+    from mteb.types import TopRankedDocumentsType
+
 
 def calculate_text_statistics(texts: list[str]) -> TextStatistics:
     """Calculate descriptive statistics for a list of texts.
@@ -87,13 +89,13 @@ def calculate_label_statistics(labels: list[int | list[int]]) -> LabelStatistics
 
     if not isinstance(labels[0], list):
         # single label classification
-        single_label = cast(list[int], labels)
+        single_label = cast("list[int]", labels)
         label_len = [1] * len(single_label)
         total_label_len = len(single_label)
         total_labels.extend(single_label)
     elif isinstance(labels[0], list):
         # multilabel classification
-        multilabel_labels = cast(list[list[int]], labels)
+        multilabel_labels = cast("list[list[int]]", labels)
         label_len = [len(l) for l in multilabel_labels]
         total_label_len = sum(label_len)
         for l in multilabel_labels:

mteb/abstasks/abstask.py

@@ -1,30 +1,38 @@
+from __future__ import annotations
+
 import json
 import logging
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 from copy import copy
 from pathlib import Path
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 from datasets import ClassLabel, Dataset, DatasetDict, load_dataset
 from sklearn.preprocessing import MultiLabelBinarizer
 from tqdm.auto import tqdm
-from typing_extensions import Self
 
 from mteb._set_seed import _set_seed
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.languages import LanguageScripts
 from mteb.models import (
     CrossEncoderProtocol,
     EncoderProtocol,
-    MTEBModels,
     SearchProtocol,
 )
-from mteb.types import HFSubset, Modalities, ScoresDict
-from mteb.types._encoder_io import EncodeKwargs
-from mteb.types.statistics import DescriptiveStatistics, SplitDescriptiveStatistics
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from typing_extensions import Self
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import (
+        MTEBModels,
+    )
+    from mteb.types import EncodeKwargs, HFSubset, Modalities, ScoresDict
+    from mteb.types.statistics import DescriptiveStatistics, SplitDescriptiveStatistics
 
 logger = logging.getLogger(__name__)
 
@@ -163,7 +171,7 @@ class AbsTask(ABC):
         if not self.data_loaded:
             self.load_data()
 
-        self.dataset = cast(dict[HFSubset, DatasetDict], self.dataset)
+        self.dataset = cast("dict[HFSubset, DatasetDict]", self.dataset)
 
         scores = {}
         if self.hf_subsets is None:

mteb/abstasks/aggregate_task_metadata.py

@@ -1,28 +1,39 @@
+from __future__ import annotations
+
 import logging
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 from pydantic import ConfigDict, Field, model_validator
-from typing_extensions import Self
 
 from mteb.types import (
-    ISOLanguageScript,
     Languages,
-    Licenses,
-    Modalities,
-    StrDate,
 )
 
 from .abstask import AbsTask
 from .task_metadata import (
-    AnnotatorType,
     MetadataDatasetDict,
-    SampleCreationMethod,
-    TaskDomain,
     TaskMetadata,
-    TaskSubtype,
     TaskType,
 )
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from mteb.types import (
+        ISOLanguageScript,
+        Licenses,
+        Modalities,
+        StrDate,
+    )
+
+    from .task_metadata import (
+        AnnotatorType,
+        SampleCreationMethod,
+        TaskDomain,
+        TaskSubtype,
+    )
+
 logger = logging.getLogger(__name__)
 
 

mteb/abstasks/aggregated_task.py

@@ -1,19 +1,26 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Mapping
-from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
-from datasets import Dataset, DatasetDict
 
-from mteb.models.models_protocols import MTEBModels
 from mteb.results.task_result import TaskResult
-from mteb.types import EncodeKwargs, HFSubset, ScoresDict
-from mteb.types.statistics import DescriptiveStatistics
 
 from .abstask import AbsTask
-from .aggregate_task_metadata import AggregateTaskMetadata
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from pathlib import Path
+
+    from datasets import Dataset, DatasetDict
+
+    from mteb.models.models_protocols import MTEBModels
+    from mteb.types import EncodeKwargs, HFSubset, ScoresDict
+    from mteb.types.statistics import DescriptiveStatistics
+
+    from .aggregate_task_metadata import AggregateTaskMetadata
 
 logger = logging.getLogger(__name__)
 

mteb/abstasks/classification.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 from collections import defaultdict
-from pathlib import Path
-from typing import Any, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict
 
 import numpy as np
 from datasets import Dataset, DatasetDict
@@ -16,12 +17,8 @@ from sklearn.metrics import (
 
 from mteb._evaluators.sklearn_evaluator import SklearnEvaluator, SklearnModelProtocol
 from mteb.models import EncoderProtocol, MTEBModels
-from mteb.types import EncodeKwargs, HFSubset, ScoresDict
 from mteb.types.statistics import (
-    ImageStatistics,
-    LabelStatistics,
     SplitDescriptiveStatistics,
-    TextStatistics,
 )
 
 from ._statistics_calculation import (
@@ -31,6 +28,18 @@ from ._statistics_calculation import (
 )
 from .abstask import AbsTask
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from mteb._evaluators.sklearn_evaluator import SklearnModelProtocol
+    from mteb.models import MTEBModels
+    from mteb.types import EncodeKwargs, HFSubset, ScoresDict
+    from mteb.types.statistics import (
+        ImageStatistics,
+        LabelStatistics,
+        TextStatistics,
+    )
+
 logger = logging.getLogger(__name__)