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

mteb/_evaluators/sklearn_evaluator.py

@@ -1,15 +1,15 @@
 import logging
-from typing import Any, Protocol
+from typing import Any, Protocol, cast
 
 import numpy as np
 from datasets import Dataset
 from torch.utils.data import DataLoader
 from typing_extensions import Self
 
-from mteb._create_dataloaders import _create_image_dataloader
+from mteb._create_dataloaders import create_dataloader
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
-from mteb.types import BatchedInput
+from mteb.types import Array, BatchedInput, EncodeKwargs
 
 from .evaluator import Evaluator
 
@@ -17,11 +17,11 @@ logger = logging.getLogger(__name__)
 
 
 class SklearnModelProtocol(Protocol):
-    def fit(self, X: np.ndarray, y: np.ndarray | list[int]) -> None: ...  # noqa: N803
-    def predict(self, X: np.ndarray) -> np.ndarray: ...  # noqa: N803
+    def fit(self, X: Array, y: np.ndarray | list[int]) -> None: ...  # noqa: N803
+    def predict(self, X: Array) -> np.ndarray: ...  # noqa: N803
     def get_params(self) -> dict[str, Any]: ...
-    def set_params(self, **kwargs: dict[str, Any]) -> Self: ...
-    def score(self, X: np.ndarray, y: np.ndarray | list[int]) -> float: ...  # noqa: N803
+    def set_params(self, random_state: int, **kwargs: dict[str, Any]) -> Self: ...
+    def score(self, X: Array, y: np.ndarray | list[int]) -> float: ...  # noqa: N803
 
 
 class SklearnEvaluator(Evaluator):
@@ -50,42 +50,29 @@ class SklearnEvaluator(Evaluator):
         self.evaluator_model = evaluator_model
 
     def create_dataloaders(
-        self, batch_size: int
+        self, encode_kwargs: EncodeKwargs
     ) -> tuple[DataLoader[BatchedInput], DataLoader[BatchedInput]]:
-        if self.task_metadata.modalities == ["image"]:
-            dataloader_train = _create_image_dataloader(
-                self.train_dataset,
-                image_column_name=self.values_column_name,
-                batch_size=batch_size,
-            )
-            dataloader_test = _create_image_dataloader(
-                self.eval_dataset,
-                image_column_name=self.values_column_name,
-                batch_size=batch_size,
-            )
-        elif self.task_metadata.modalities == ["text"]:
-            if self.values_column_name != "text":
-                self.train_dataset = self.train_dataset.rename_column(
-                    self.values_column_name, "text"
-                )
-                self.eval_dataset = self.eval_dataset.rename_column(
-                    self.values_column_name, "text"
-                )
-            dataloader_train = DataLoader(self.train_dataset)
-            dataloader_test = DataLoader(self.eval_dataset)
-        else:
-            raise ValueError(
-                "ClassificationEvaluator only supports image and text modalities."
-            )
+        dataloader_train = create_dataloader(
+            self.train_dataset,
+            self.task_metadata,
+            input_column=self.values_column_name,
+            **encode_kwargs,
+        )
+        dataloader_test = create_dataloader(
+            self.eval_dataset,
+            self.task_metadata,
+            input_column=self.values_column_name,
+            **encode_kwargs,
+        )
         return dataloader_train, dataloader_test
 
     def __call__(  # type: ignore[override]
         self,
         model: EncoderProtocol,
         *,
-        encode_kwargs: dict[str, Any],
-        test_cache: np.ndarray | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
+        encode_kwargs: EncodeKwargs,
+        test_cache: Array | None = None,
+    ) -> tuple[np.ndarray, Array]:
         """Classification evaluation by training a sklearn classifier on the embeddings of the training set and evaluating on the embeddings of the test set.
 
         Args:
@@ -98,7 +85,7 @@ class SklearnEvaluator(Evaluator):
 
         """
         dataloader_train, dataloader_test = self.create_dataloaders(
-            batch_size=encode_kwargs["batch_size"]
+            encode_kwargs=encode_kwargs,
         )
 
         logger.info("Running - Encoding samples...")
@@ -117,6 +104,7 @@ class SklearnEvaluator(Evaluator):
                 hf_subset=self.hf_subset,
                 **encode_kwargs,
             )
+            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,7 +1,5 @@
 import logging
-from typing import Any
 
-import numpy as np
 import torch
 from datasets import Dataset
 from tqdm.auto import tqdm
@@ -10,6 +8,7 @@ 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
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +32,10 @@ class BitextMiningEvaluator(Evaluator):
         self.task_metadata = task_metadata
 
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
     ) -> dict[str, list[dict[str, float]]]:
         pair_elements = {p for pair in self.pairs for p in pair}
         if isinstance(self.sentences, Dataset):
@@ -46,7 +48,10 @@ class BitextMiningEvaluator(Evaluator):
 
         embeddings = {}
         for sub in tqdm(subsets):
-            dataloader = _create_dataloader_from_texts(self.sentences[sub])
+            dataloader = _create_dataloader_from_texts(
+                self.sentences[sub],
+                **encode_kwargs,
+            )
             embeddings[sub] = model.encode(
                 dataloader,
                 task_metadata=self.task_metadata,
@@ -66,11 +71,11 @@ class BitextMiningEvaluator(Evaluator):
 
     def _similarity_search(
         self,
-        query_embeddings: np.ndarray,
-        corpus_embeddings: np.ndarray,
+        query_embeddings: Array,
+        corpus_embeddings: Array,
         model: EncoderProtocol,
         query_chunk_size: int = 100,
-        corpus_chunk_size: int = 500000,
+        corpus_chunk_size: int = 500_000,
     ) -> list[dict[str, float]]:
         """This function performs a cosine similarity search between a list of query embeddings and a list of corpus embeddings.
 
@@ -101,13 +106,15 @@ class BitextMiningEvaluator(Evaluator):
         ):
             query_embeddings = query_embeddings.to(corpus_embeddings.device)
 
-        queries_result_list = [[] for _ in range(len(query_embeddings))]
+        queries_result_list: list[list[dict[str, float]]] = [
+            [] for _ in range(len(query_embeddings))
+        ]
 
         for query_start_idx in range(0, len(query_embeddings), query_chunk_size):
             # Iterate over chunks of the corpus
             for corpus_start_idx in range(0, len(corpus_embeddings), corpus_chunk_size):
                 # Compute cosine similarities
-                similarity_scores = model.similarity(  # type: ignore
+                similarity_scores = model.similarity(
                     query_embeddings[
                         query_start_idx : query_start_idx + query_chunk_size
                     ],
@@ -117,15 +124,17 @@ class BitextMiningEvaluator(Evaluator):
                 )
 
                 # Get top-k scores
-                cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(
-                    torch.tensor(similarity_scores),
-                    1,
-                    dim=1,
-                    largest=True,
-                    sorted=False,
+                cos_scores_top_k_values_tensor, cos_scores_top_k_idx_tensor = (
+                    torch.topk(
+                        torch.tensor(similarity_scores),
+                        1,
+                        dim=1,
+                        largest=True,
+                        sorted=False,
+                    )
                 )
-                cos_scores_top_k_values = cos_scores_top_k_values.cpu().tolist()
-                cos_scores_top_k_idx = cos_scores_top_k_idx.cpu().tolist()
+                cos_scores_top_k_values = cos_scores_top_k_values_tensor.cpu().tolist()
+                cos_scores_top_k_idx = cos_scores_top_k_idx_tensor.cpu().tolist()
 
                 for query_itr in range(len(similarity_scores)):
                     for sub_corpus_id, score in zip(
@@ -138,11 +147,14 @@ class BitextMiningEvaluator(Evaluator):
                             {"corpus_id": corpus_id, "score": score}
                         )
 
+        result_queries_list: list[dict[str, float]] = [
+            {} for _ in range(len(query_embeddings))
+        ]
         # Sort and strip to top_k results
         for idx in range(len(queries_result_list)):
             queries_result_list[idx] = sorted(
                 queries_result_list[idx], key=lambda x: x["score"], reverse=True
             )
-            queries_result_list[idx] = queries_result_list[idx][0]
+            result_queries_list[idx] = queries_result_list[idx][0]
 
-        return queries_result_list
+        return result_queries_list

mteb/_evaluators/text/summarization_evaluator.py

@@ -1,6 +1,6 @@
 import logging
 import sys
-from typing import Any, TypedDict
+from typing import TypedDict
 
 import numpy as np
 import torch
@@ -12,6 +12,7 @@ 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 later than python 3.13 use typing module
 if sys.version_info >= (3, 13):
@@ -94,7 +95,7 @@ class SummarizationEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         *,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> SummarizationDistances:
         # Get the human & machine summaries for the text in one go for all
         human_lens = [len(human_summaries) for human_summaries in self.human_summaries]
@@ -109,7 +110,8 @@ class SummarizationEvaluator(Evaluator):
                     summary
                     for human_summaries in self.human_summaries
                     for summary in human_summaries
-                ]
+                ],
+                **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,
@@ -124,7 +126,8 @@ class SummarizationEvaluator(Evaluator):
                     summary
                     for machine_summaries in self.machine_summaries
                     for summary in machine_summaries
-                ]
+                ],
+                **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,
@@ -133,10 +136,10 @@ class SummarizationEvaluator(Evaluator):
         )
 
         # Split the embeddings into the original human & machine summaries
-        embs_human_summaries_all = np.split(
+        embs_human_summaries_all_split = np.split(
             embs_human_summaries_all, np.cumsum(human_lens)[:-1]
         )
-        embs_machine_summaries_all = np.split(
+        embs_machine_summaries_all_split = np.split(
             embs_machine_summaries_all, np.cumsum(machine_lens)[:-1]
         )
 
@@ -146,7 +149,9 @@ class SummarizationEvaluator(Evaluator):
         all_human_scores = []
 
         for i, (embs_human_summaries, embs_machine_summaries) in tqdm(
-            enumerate(zip(embs_human_summaries_all, embs_machine_summaries_all)),
+            enumerate(
+                zip(embs_human_summaries_all_split, embs_machine_summaries_all_split)
+            ),
             desc="Scoring",
             total=len(self.human_summaries),
         ):
@@ -162,7 +167,7 @@ class SummarizationEvaluator(Evaluator):
                 dot_scores = dot_score(emb_machine_summary, embs_human_summaries)
 
                 _sim_score = [
-                    float(model.similarity(emb_machine_summary, emb_human_summary))  # type: ignore
+                    float(model.similarity(emb_machine_summary, emb_human_summary))
                     for emb_human_summary in embs_human_summaries
                 ]
                 sim_score = torch.tensor(_sim_score)
@@ -214,17 +219,19 @@ class SummarizationEvaluator(Evaluator):
             strict=True,
         ):
             cosine_spearman_scores.append(
-                spearmanr(human_scores, cosine_pred_scores).statistic
+                float(spearmanr(human_scores, cosine_pred_scores).statistic)
             )
             cosine_pearson_scores.append(
-                pearsonr(human_scores, cosine_pred_scores).statistic
+                float(pearsonr(human_scores, cosine_pred_scores).statistic)
             )
             dot_spearman_scores.append(
-                spearmanr(human_scores, dot_pred_scores).statistic
+                float(spearmanr(human_scores, dot_pred_scores).statistic)
             )
-            dot_pearson_scores.append(pearsonr(human_scores, dot_pred_scores).statistic)
-            spearman_scores.append(spearmanr(human_scores, sim_scores).statistic)
-            pearson_scores.append(pearsonr(human_scores, sim_scores).statistic)
+            dot_pearson_scores.append(
+                float(pearsonr(human_scores, dot_pred_scores).statistic)
+            )
+            spearman_scores.append(float(spearmanr(human_scores, sim_scores).statistic))
+            pearson_scores.append(float(pearsonr(human_scores, sim_scores).statistic))
 
         return SummarizationMetrics(
             pearson=float(np.mean(pearson_scores)),
@@ -271,10 +278,10 @@ class DeprecatedSummarizationEvaluator(SummarizationEvaluator):
             pearson_scores.append(pearsonr(human_scores, sim_scores))
 
         return SummarizationMetrics(
-            pearson=float(np.mean(pearson_scores)),
-            spearman=float(np.mean(spearman_scores)),
-            cosine_spearman=float(np.mean(cosine_spearman_scores)),
-            cosine_pearson=float(np.mean(cosine_pearson_scores)),
-            dot_pearson=float(np.mean(dot_pearson_scores)),
-            dot_spearman=float(np.mean(dot_spearman_scores)),
+            pearson=float(np.mean(pearson_scores)),  # type: ignore[arg-type]
+            spearman=float(np.mean(spearman_scores)),  # type: ignore[arg-type]
+            cosine_spearman=float(np.mean(cosine_spearman_scores)),  # type: ignore[arg-type]
+            cosine_pearson=float(np.mean(cosine_pearson_scores)),  # type: ignore[arg-type]
+            dot_pearson=float(np.mean(dot_pearson_scores)),  # type: ignore[arg-type]
+            dot_spearman=float(np.mean(dot_spearman_scores)),  # type: ignore[arg-type]
         )

mteb/_evaluators/zeroshot_classification_evaluator.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Any
 
 from datasets import Dataset
 
@@ -10,7 +9,7 @@ from mteb._create_dataloaders import (
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.similarity_functions import similarity
-from mteb.types import Array
+from mteb.types import Array, EncodeKwargs
 
 from .evaluator import Evaluator
 
@@ -38,18 +37,21 @@ class ZeroShotClassificationEvaluator(Evaluator):
         self.hf_subset = hf_subset
 
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: dict[str, Any]
+        self,
+        model: EncoderProtocol,
+        *,
+        encode_kwargs: EncodeKwargs,
     ) -> Array:
         dataloader = create_dataloader(
             self.dataset,
-            batch_size=encode_kwargs["batch_size"],
             input_column=self.input_column_name,
             task_metadata=self.task_metadata,
+            **encode_kwargs,
         )
 
         logger.info("Running zero-shot classification - Encoding labels...")
         text_label_embeddings = model.encode(
-            _create_dataloader_from_texts(self.candidate_labels),
+            _create_dataloader_from_texts(self.candidate_labels, **encode_kwargs),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,
             hf_split=self.hf_split,

mteb/abstasks/_data_filter/filters.py

@@ -0,0 +1,125 @@
+"""Simplified version of https://gist.github.com/AlexeyVatolin/ea3adc21aa7a767603ff393b22085adc from https://github.com/embeddings-benchmark/mteb/pull/2900"""
+
+import logging
+
+import datasets
+import pandas as pd
+from datasets import Dataset, DatasetDict
+
+from mteb import TaskMetadata
+
+logger = logging.getLogger(__name__)
+
+
+def deduplicate(dataset: Dataset, input_column: str) -> Dataset:
+    """Remove duplicate texts, keeping the first occurrence."""
+    unique_texts = set()
+    indices_to_keep = []
+    for i, text in enumerate(dataset[input_column]):
+        text = text.strip()
+        if text not in unique_texts:
+            unique_texts.add(text)
+            indices_to_keep.append(i)
+
+    logger.info(
+        f"[deduplicate] removed={len(dataset) - len(indices_to_keep)}/{len(dataset)}"
+    )
+    return dataset.select(indices_to_keep)
+
+
+def filter_empty(dataset: Dataset, input_column: str) -> Dataset:
+    """Filter out empty or whitespace-only examples."""
+    before = len(dataset)
+    ds = dataset.filter(lambda x: len(x[input_column].strip()) > 0)
+    logger.info(f"[filter_empty] removed={before - len(ds)}/{before}")
+    return ds
+
+
+def filter_train_leakage(
+    train_dataset: Dataset, test_dataset: Dataset, input_column: str
+) -> Dataset:
+    """Remove test examples that appear in training."""
+    train_texts = set(train_dataset[input_column])
+    before = len(test_dataset)
+    indices = [
+        i
+        for i, text in enumerate(test_dataset[input_column])
+        if text not in train_texts
+    ]
+    logger.info(f"[filter_train_leakage] removed={before - len(indices)}/{before}")
+    return test_dataset.select(indices)
+
+
+def filter_unclear_label(
+    dataset_dict: DatasetDict, input_column: str, label_column: str
+) -> DatasetDict:
+    """Remove examples where the same text appears with multiple different labels."""
+    normalized: dict[str, set[str | tuple[str, ...]]] = {}
+    logger.debug("[filter_controversial] scanning dataset for label conflicts...")
+
+    for split, ds in dataset_dict.items():
+        for text, label in zip(ds[input_column], ds[label_column]):
+            key = text.strip().lower()
+            normalized.setdefault(key, set()).add(
+                label if isinstance(label, (str, int, float)) else tuple(label)  # type: ignore[arg-type]
+            )
+
+    bad_texts = {t for t, labels in normalized.items() if len(labels) > 1}
+    logger.info(f"[filter_controversial] Removing {len(bad_texts)} conflicting texts")
+
+    new_dict = {}
+    for split, ds in dataset_dict.items():
+        before = len(ds)
+        filtered = ds.filter(lambda x: x[input_column].strip().lower() not in bad_texts)
+        logger.debug(
+            f"[filter_controversial:{split}] removed={before - len(filtered)}/{before}"
+        )
+        new_dict[split] = filtered
+
+    return DatasetDict(new_dict)
+
+
+def filter_short(dataset: Dataset, input_column: str, min_words: int = 3) -> Dataset:
+    """Filter out texts with fewer than `min_words`."""
+    before = len(dataset)
+    ds = dataset.filter(lambda x: len(x[input_column].strip().split()) >= min_words)
+    logger.debug(f"[filter_short] removed={before - len(ds)}/{before}")
+    return ds
+
+
+def split_train_test(
+    ds: DatasetDict,
+    metadata: TaskMetadata,
+    train_split: str,
+    label_column: str,
+) -> DatasetDict:
+    if train_split in ds and metadata.eval_splits == train_split:
+        before = len(ds[train_split])
+        logger.info(
+            f"[split_train_test] eval_splits == train_split; performing split on {before} examples"
+        )
+        ds[train_split] = ds[train_split].cast_column(
+            label_column,
+            datasets.ClassLabel(names=list(set(ds[train_split][label_column]))),
+        )
+        label_counts = pd.Series(ds[train_split][label_column]).value_counts()
+        one_sample_labels = set(label_counts[label_counts == 1].index.tolist())
+
+        if one_sample_labels:
+            logger.info(
+                f"[split_train_test] Removing {len(one_sample_labels)} labels with only one instance"
+            )
+            ds[train_split] = ds[train_split].filter(
+                lambda x: x[label_column] not in one_sample_labels
+            )
+
+        splits = ds[train_split].train_test_split(
+            test_size=min(2048, before // 2), seed=42, stratify_by_column=label_column
+        )
+        ds = DatasetDict({train_split: splits[train_split], "test": splits["test"]})
+        metadata.eval_splits = ["test"]
+        logger.info(
+            f"[split_train_test] Train size={len(ds[train_split])}, Test size={len(ds['test'])}"
+        )
+
+    return ds

mteb/abstasks/_data_filter/task_pipelines.py

@@ -0,0 +1,105 @@
+import logging
+
+from datasets import DatasetDict
+
+from mteb import TaskMetadata
+from mteb.abstasks import AbsTaskClassification
+from mteb.abstasks._data_filter.filters import (
+    deduplicate,
+    filter_empty,
+    filter_short,
+    filter_train_leakage,
+    filter_unclear_label,
+    split_train_test,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def clean_dataset(
+    ds: DatasetDict,
+    metadata: TaskMetadata,
+    train_split: str,
+    input_column: str,
+    label_column: str,
+    subset: str | None = None,
+) -> DatasetDict:
+    """Apply the full cleaning pipeline with logging."""
+    logger.info("[clean_dataset] Starting dataset cleaning pipeline...")
+
+    transforms = [
+        ("filter_empty", filter_empty),
+        ("deduplicate", deduplicate),
+    ]
+
+    skip_cjk_codes = {"zho", "jpn", "tha", "mya", "cmn"}
+    logger.info("[clean_dataset] Applying short-text filter")
+    cur_langs = (
+        metadata.eval_langs[subset]
+        if isinstance(metadata.eval_langs, dict) and subset
+        else metadata.eval_langs
+    )
+    apply_short = not any(lang.split("-")[0] in skip_cjk_codes for lang in cur_langs)
+    if apply_short:
+        logger.info("[clean_dataset] Applying short-text filter")
+        transforms.append(("filter_short", filter_short))
+
+    for split in [train_split, *metadata.eval_splits]:
+        if split not in ds:
+            logger.warning(f"[clean_dataset] Split '{split}' missing; skipping.")
+            continue
+
+        for name, fn in transforms:
+            before = len(ds[split])
+            ds[split] = fn(ds[split], input_column=input_column)
+            logger.info(
+                f"[clean_dataset:{split}] {name} removed={before - len(ds[split])}"
+            )
+
+    ds = split_train_test(ds, metadata, train_split, label_column)
+
+    for split in metadata.eval_splits:
+        if split == train_split:
+            continue
+        before = len(ds[split])
+        ds[split] = filter_train_leakage(ds[train_split], ds[split], input_column)
+        logger.info(
+            f"[clean_dataset:{split}] leakage_removed={before - len(ds[split])}"
+        )
+
+    ds = filter_unclear_label(ds, input_column=input_column, label_column=label_column)
+
+    logger.info("[clean_dataset] Cleaning pipeline complete.")
+    return ds
+
+
+def process_classification(
+    task: AbsTaskClassification,
+) -> DatasetDict | dict[str, DatasetDict]:
+    """Process classification task dataset(s) with cleaning pipeline."""
+    if not task.data_loaded:
+        task.load_data()
+    if isinstance(task.dataset, DatasetDict):
+        return clean_dataset(
+            task.dataset,
+            task.metadata,
+            task.train_split,
+            task.input_column_name,
+            task.label_column_name,
+            subset=None,
+        )
+
+    if task.dataset is None:
+        raise ValueError("Task dataset is None.")
+
+    new_ds = {}
+    for subset in task.dataset:
+        new_ds[subset] = clean_dataset(
+            task.dataset[subset],
+            task.metadata,
+            task.train_split,
+            task.input_column_name,
+            task.label_column_name,
+            subset=subset,
+        )
+    return new_ds