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

mteb/models/model_implementations/vista_models.py

@@ -1,14 +1,19 @@
-from typing import Any, Literal
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
 
 import torch
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import requires_image_dependencies
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
 
 VISTA_CITATION = """@article{zhou2024vista,
   title={VISTA: Visualized Text Embedding For Universal Multi-Modal Retrieval},

mteb/models/model_implementations/vlm2vec_models.py

@@ -1,8 +1,9 @@
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import (
@@ -10,10 +11,14 @@ from mteb._requires_package import (
     requires_package,
     suggest_package,
 )
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, PromptType
 
 logger = logging.getLogger(__name__)
 

mteb/models/model_implementations/voyage_models.py

@@ -1,16 +1,22 @@
+from __future__ import annotations
+
 import time
 from functools import wraps
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
+
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
 
 VOYAGE_TRAINING_DATA = set(
     # Self-reported (message from VoyageAI member)

mteb/models/model_implementations/voyage_v.py

@@ -4,17 +4,19 @@ import logging
 from typing import TYPE_CHECKING, Any, Literal
 
 import torch
-from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from mteb._requires_package import requires_image_dependencies, requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
 
 if TYPE_CHECKING:
     from PIL import Image
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +29,8 @@ def _downsample_image(
     Returns:
         The downsampled image.
     """
+    from PIL.Image import Resampling
+
     width, height = image.size
     pixels = width * height
 
@@ -42,15 +46,15 @@ def _downsample_image(
         logger.info(
             f"Downsampling image from {width}x{height} to {new_width}x{new_height}"
         )
-        return image.resize(new_size, Image.LANCZOS)
+        return image.resize(new_size, Resampling.LANCZOS)
     if width > height:
         if width > 10000:
             logger.error("Processing extremely wide images.")
-            return image.resize((10000, height), Image.LANCZOS)
+            return image.resize((10000, height), Resampling.LANCZOS)
     else:
         if height > 10000:
             logger.error("Processing extremely high images.")
-            return image.resize((width, 10000), Image.LANCZOS)
+            return image.resize((width, 10000), Resampling.LANCZOS)
     return image
 
 

mteb/models/model_implementations/yuan_models_en.py

@@ -1,6 +1,6 @@
 from mteb.models.instruct_wrapper import InstructSentenceTransformerModel
 from mteb.models.model_meta import ModelMeta
-from mteb.models.models_protocols import PromptType
+from mteb.types import PromptType
 
 
 def instruction_template(

mteb/models/model_meta.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import json
 import logging
 import warnings
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from dataclasses import field
 from enum import Enum
 from functools import partial
@@ -11,9 +11,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 from huggingface_hub import (
-    GitCommitInfo,
     ModelCard,
-    ModelCardData,
     get_safetensors_metadata,
     hf_hub_download,
     list_repo_commits,
@@ -30,17 +28,24 @@ from huggingface_hub.errors import (
 )
 from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 from transformers import AutoConfig
-from typing_extensions import Self
 
 from mteb._helpful_enum import HelpfulStrEnum
 from mteb.languages import check_language_code
-from mteb.models.models_protocols import EncoderProtocol, MTEBModels
+from mteb.models.models_protocols import MTEBModels
 from mteb.types import ISOLanguageScript, Licenses, Modalities, StrDate, StrURL
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from huggingface_hub import (
+        GitCommitInfo,
+        ModelCardData,
+    )
     from sentence_transformers import CrossEncoder, SentenceTransformer
+    from typing_extensions import Self
 
     from mteb.abstasks import AbsTask
+    from mteb.models.models_protocols import EncoderProtocol
 
 
 logger = logging.getLogger(__name__)
@@ -479,7 +484,7 @@ class ModelMeta(BaseModel):
         if isinstance(tasks[0], str):
             benchmark_datasets = set(tasks)
         else:
-            tasks = cast(Sequence["AbsTask"], tasks)
+            tasks = cast("Sequence[AbsTask]", tasks)
             benchmark_datasets = set()
             for task in tasks:
                 benchmark_datasets.add(task.metadata.name)
@@ -534,7 +539,7 @@ class ModelMeta(BaseModel):
         if isinstance(tasks[0], str):
             benchmark_datasets = set(tasks)
         else:
-            tasks = cast(Sequence["AbsTask"], tasks)
+            tasks = cast("Sequence[AbsTask]", tasks)
             benchmark_datasets = {task.metadata.name for task in tasks}
         overlap = training_datasets & benchmark_datasets
         perc_overlap = 100 * (len(overlap) / len(benchmark_datasets))

mteb/models/models_protocols.py

@@ -1,22 +1,23 @@
-from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+from __future__ import annotations
 
-from torch.utils.data import DataLoader
-from typing_extensions import Unpack
-
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.types import (
-    Array,
-    BatchedInput,
-    CorpusDatasetType,
-    EncodeKwargs,
-    PromptType,
-    QueryDatasetType,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
-)
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+    from typing_extensions import Unpack
+
+    from mteb.abstasks.task_metadata import TaskMetadata
     from mteb.models.model_meta import ModelMeta
+    from mteb.types import (
+        Array,
+        BatchedInput,
+        CorpusDatasetType,
+        EncodeKwargs,
+        PromptType,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
 
 
 @runtime_checkable
@@ -72,7 +73,7 @@ class SearchProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -177,7 +178,7 @@ class EncoderProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 
@@ -236,7 +237,7 @@ class CrossEncoderProtocol(Protocol):
         ...
 
     @property
-    def mteb_model_meta(self) -> "ModelMeta":
+    def mteb_model_meta(self) -> ModelMeta:
         """Metadata of the model"""
         ...
 

mteb/models/search_encoder_index/search_backend_protocol.py

@@ -1,7 +1,11 @@
-from collections.abc import Callable
-from typing import Protocol
+from __future__ import annotations
 
-from mteb.types import Array, TopRankedDocumentsType
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from mteb.types import Array, TopRankedDocumentsType
 
 
 class IndexEncoderSearchProtocol(Protocol):

mteb/models/search_encoder_index/search_indexes/faiss_search_index.py

@@ -1,14 +1,23 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
 
 from mteb._requires_package import requires_package
 from mteb.models.model_meta import ScoringFunction
-from mteb.models.models_protocols import EncoderProtocol
-from mteb.types import Array, TopRankedDocumentsType
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import faiss
+
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import Array, TopRankedDocumentsType
+
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +42,6 @@ class FaissSearchIndex:
             install_instruction="pip install mteb[faiss-cpu]",
         )
 
-        import faiss
         from faiss import IndexFlatIP, IndexFlatL2
 
         # https://github.com/facebookresearch/faiss/wiki/Faiss-indexes

mteb/models/search_wrappers.py

@@ -1,28 +1,35 @@
+from __future__ import annotations
+
 import heapq
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 from datasets import Dataset
-from torch.utils.data import DataLoader
 
 from mteb._create_dataloaders import (
     create_dataloader,
 )
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.types import (
-    Array,
-    BatchedInput,
-    CorpusDatasetType,
-    EncodeKwargs,
     PromptType,
-    QueryDatasetType,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
 )
 
-from .models_protocols import CrossEncoderProtocol, EncoderProtocol
-from .search_encoder_index.search_backend_protocol import IndexEncoderSearchProtocol
+if TYPE_CHECKING:
+    from torch.utils.data import DataLoader
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import (
+        Array,
+        BatchedInput,
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
+    from .models_protocols import CrossEncoderProtocol, EncoderProtocol
+    from .search_encoder_index.search_backend_protocol import IndexEncoderSearchProtocol
 
 logger = logging.getLogger(__name__)
 

mteb/models/sentence_transformer_wrapper.py

@@ -7,19 +7,20 @@ from typing import TYPE_CHECKING, Any
 import numpy as np
 import torch
 from packaging.version import Version
-from torch.utils.data import DataLoader
-from typing_extensions import Unpack
 
 from mteb._log_once import LogOnce
 from mteb.models import ModelMeta
-from mteb.types import Array, BatchedInput, EncodeKwargs, PromptType
+from mteb.types import PromptType
 
 from .abs_encoder import AbsEncoder
 
 if TYPE_CHECKING:
     from sentence_transformers import CrossEncoder, SentenceTransformer
+    from torch.utils.data import DataLoader
+    from typing_extensions import Unpack
 
     from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput, EncodeKwargs
 
 logger = logging.getLogger(__name__)
 

mteb/models/vllm_wrapper.py

@@ -4,23 +4,25 @@ import atexit
 import gc
 import logging
 import os
-from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import torch
-from torch.utils.data import DataLoader
 
 from mteb._requires_package import requires_package
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import ModelMeta
 from mteb.models.abs_encoder import AbsEncoder
-from mteb.types import Array, BatchedInput, PromptType
+from mteb.types import PromptType
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from torch.utils.data import DataLoader
     from vllm.config import PoolerConfig  # type: ignore[import-not-found]
-else:
-    PoolerConfig = dict[str, Any]
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import Array, BatchedInput
+
 
 logger = logging.getLogger(__name__)
 

mteb/results/benchmark_results.py

@@ -4,34 +4,39 @@ import functools
 import json
 import logging
 import warnings
-from collections.abc import Callable, Iterable, Iterator
 from pathlib import Path
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import pandas as pd
 from packaging.version import InvalidVersion, Version
 from pydantic import BaseModel, ConfigDict
-from typing_extensions import Self
 
-from mteb.abstasks.abstask import AbsTask
-from mteb.abstasks.task_metadata import (
-    TaskDomain,
-    TaskType,
-)
 from mteb.benchmarks.benchmark import Benchmark
 from mteb.models import ModelMeta
 from mteb.models.get_model_meta import get_model_metas
-from mteb.types import (
-    ISOLanguage,
-    ISOLanguageScript,
-    Modalities,
-    Score,
-    ScoresDict,
-    SplitName,
-)
 
 from .model_result import ModelResult, _aggregate_and_pivot
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Iterator
+
+    from typing_extensions import Self
+
+    from mteb.abstasks.abstask import AbsTask
+    from mteb.abstasks.task_metadata import (
+        TaskDomain,
+        TaskType,
+    )
+    from mteb.types import (
+        ISOLanguage,
+        ISOLanguageScript,
+        Modalities,
+        Score,
+        ScoresDict,
+        SplitName,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -144,7 +149,7 @@ class BenchmarkResults(BaseModel):
                     raise ValueError("name in ModelMeta is None. It must be a string.")
                 name_rev[name.name] = name.revision
             else:
-                name_ = cast(str, name)
+                name_ = cast("str", name)
                 name_rev[name_] = revision
 
         for model_res in self.model_results:

mteb/results/model_result.py

@@ -2,30 +2,36 @@ from __future__ import annotations
 
 import logging
 import warnings
-from collections.abc import Callable, Iterable
-from typing import Any, Literal, cast
+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 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__)
 
 
@@ -83,7 +89,7 @@ class ModelResult(BaseModel):
     model_revision: str | None
     task_results: list[TaskResult]
     default_modalities: list[Modalities] = Field(
-        default_factory=lambda: [cast(Modalities, "text")], alias="modalities"
+        default_factory=lambda: [cast("Modalities", "text")], alias="modalities"
     )
     model_config = (
         ConfigDict(  # to free up the name model_* which is otherwise protected
@@ -202,8 +208,8 @@ 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)
+        aggregation = cast("Callable[[list[Score]], Any]", aggregation)
+        getter = cast("Callable[[ScoresDict], Score]", getter)
 
         if format == "wide":
             scores = {}

mteb/results/task_result.py

@@ -4,34 +4,40 @@ import json
 import logging
 import warnings
 from collections import defaultdict
-from collections.abc import Callable, Iterable, Mapping
 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.abstasks.task_metadata import TaskDomain
 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__)
 
 
@@ -610,7 +616,10 @@ class TaskResult(BaseModel):
         new_res = {**self.to_dict(), "scores": new_scores}
         return TaskResult.from_validated(**new_res)
 
-    def validate_and_filter_scores(self, task: AbsTask | None = None) -> TaskResult:
+    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.
@@ -658,11 +667,33 @@ class TaskResult(BaseModel):
                 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):
             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)

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():

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(),