[v2] Refactor descriptive stats

mteb/abstasks/AbsTaskAnyClassification.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from collections import Counter, defaultdict
+from collections import defaultdict
 from typing import Any
 
 import numpy as np
@@ -20,6 +20,11 @@
 )
 
 from ..evaluation.evaluators.ClassificationEvaluator import ClassificationEvaluator
+from ._statistics_calculation import (
+    calculate_image_statistics,
+    calculate_label_statistics,
+    calculate_text_statistics,
+)
 from .AbsTask import AbsTask
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
@@ -31,7 +36,6 @@ class ClassificationDescriptiveStatistics(DescriptiveStatistics):
 
     Attributes:
         num_samples: number of samples in the dataset.
-        number_of_characters: Total number of symbols in the dataset.
         number_texts_intersect_with_train: Number of texts in the train split
 
         text_statistics: Statistics for text
@@ -40,7 +44,6 @@ class ClassificationDescriptiveStatistics(DescriptiveStatistics):
     """
 
     num_samples: int
-    number_of_characters: int | None
     number_texts_intersect_with_train: int | None
 
     text_statistics: TextStatistics | None
@@ -226,79 +229,24 @@ def _calculate_metrics_from_split(
             if split != self.train_split:
                 train_text = self.dataset[self.train_split][self.input_column_name]
 
-        total_text_len = 0
-        text_len = None
-        img_widths, img_heights = None, None
+        image_statistics = None
+        text_statistics = None
         num_texts_in_train = None
 
         if "image" in self.metadata.modalities:
-            img_widths, img_heights = [], []
-            for img in inputs:
-                width, height = img.size  # type: ignore
-                img_heights.append(height)
-                img_widths.append(width)
+            image_statistics = calculate_image_statistics(inputs)
         if "text" in self.metadata.modalities:
-            text_len = [len(t) for t in inputs]
-            total_text_len = sum(text_len)
+            text_statistics = calculate_text_statistics(inputs)
             num_texts_in_train = (
                 len(set(inputs) & set(train_text))
                 if split != self.train_split
                 else None
             )
 
-        if isinstance(label[0], int):
-            label_len = [1] * len(label)
-            total_label_len = len(label)
-            total_labels = label
-        else:
-            # multilabel classification
-            label_len = [len(l) for l in label]
-            total_label_len = sum(label_len)
-            total_labels = []
-            for l in label:
-                total_labels.extend(l if len(l) > 0 else [None])
-
-        label_count = Counter(total_labels)
-
-        if text_len:
-            text_statistics = TextStatistics(
-                min_text_length=min(text_len),
-                average_text_length=total_text_len / len(inputs),
-                max_text_length=max(text_len),
-                unique_texts=len(set(inputs)),
-            )
-        else:
-            text_statistics = None
-
-        if img_widths:
-            image_statistics = ImageStatistics(
-                min_image_width=min(img_widths),
-                average_image_width=sum(img_widths) / len(img_widths),
-                max_image_width=max(img_widths),
-                min_image_height=min(img_heights),
-                average_image_height=sum(img_heights) / len(img_heights),
-                max_image_height=max(img_heights),
-            )
-        else:
-            image_statistics = None
-
-        label_statistics = LabelStatistics(
-            min_labels_per_text=min(label_len),
-            average_label_per_text=total_label_len / len(label),
-            max_labels_per_text=max(label_len),
-            unique_labels=len(label_count),
-            labels={
-                str(label): {
-                    "count": value,
-                }
-                for label, value in label_count.items()
-            },
-        )
+        label_statistics = calculate_label_statistics(label)
 
         return ClassificationDescriptiveStatistics(
             num_samples=len(inputs),
-            # text
-            number_of_characters=total_text_len,
             number_texts_intersect_with_train=num_texts_in_train
             if num_texts_in_train
             else None,

mteb/abstasks/AbsTaskAnySTS.py

@@ -15,6 +15,11 @@
 )
 
 from ..evaluation.evaluators import AnySTSEvaluator
+from ._statistics_calculation import (
+    calculate_image_statistics,
+    calculate_score_statistics,
+    calculate_text_statistics,
+)
 from .AbsTask import AbsTask
 
 logger = logging.getLogger(__name__)
@@ -108,71 +113,32 @@ def _calculate_metrics_from_split(
             score = self.dataset[split]["score"]
 
         if "text" in self.metadata.modalities:
-            text1_statistics = TextStatistics(
-                min_text_length=min(len(s) for s in sentence1),
-                average_text_length=sum(len(s) for s in sentence1) / len(sentence1),
-                max_text_length=max(len(s) for s in sentence1),
-                unique_texts=len(set(sentence1)),
-            )
-            text2_statistics = TextStatistics(
-                min_text_length=min(len(s) for s in sentence2),
-                max_text_length=max(len(s) for s in sentence2),
-                average_text_length=sum(len(s) for s in sentence2) / len(sentence2),
-                unique_texts=len(set(sentence2)),
-            )
-            sentence1_len = [len(s) for s in sentence1]
-            sentence2_len = [len(s) for s in sentence2]
-            number_of_characters = sum(sentence1_len) + sum(sentence2_len)
+            text1_statistics = calculate_text_statistics(sentence1)
+            text2_statistics = calculate_text_statistics(sentence2)
+
             unique_pairs = len(set(zip(sentence1, sentence2)))
         else:
             text1_statistics = None
             text2_statistics = None
-            number_of_characters = None
             unique_pairs = None
 
         if "image" in self.metadata.modalities:
-            img_widths1, img_heights1 = [], []
-            for img in sentence1:
-                width, height = img.size
-                img_heights1.append(height)
-                img_widths1.append(width)
-
-            image1_statistics = ImageStatistics(
-                min_image_width=min(img_widths1),
-                average_image_width=sum(img_widths1) / len(img_widths1),
-                max_image_width=max(img_widths1),
-                min_image_height=min(img_heights1),
-                average_image_height=sum(img_heights1) / len(img_heights1),
-                max_image_height=max(img_widths1),
-            )
-
-            img_widths2, img_heights2 = [], []
-            for img in sentence2:
-                width, height = img.size
-                img_heights2.append(height)
-                img_widths2.append(width)
-
-            image2_statistics = ImageStatistics(
-                min_image_width=min(img_widths2),
-                average_image_width=sum(img_widths2) / len(img_widths2),
-                max_image_width=max(img_widths2),
-                min_image_height=min(img_heights2),
-                average_image_height=sum(img_heights2) / len(img_heights2),
-                max_image_height=max(img_widths2),
-            )
+            image1_statistics = calculate_image_statistics(sentence1)
+            image2_statistics = calculate_image_statistics(sentence2)
         else:
             image1_statistics = None
             image2_statistics = None
 
-        labels_statistics = ScoreStatistics(
-            min_score=min(score),
-            avg_score=sum(score) / len(score),
-            max_score=max(score),
-        )
+        labels_statistics = calculate_score_statistics(score)
 
         return AnySTSDescriptiveStatistics(
             num_samples=len(sentence1),
-            number_of_characters=number_of_characters,
+            number_of_characters=(
+                text1_statistics["total_text_length"]
+                + text2_statistics["total_text_length"]
+                if text1_statistics
+                else None
+            ),
             unique_pairs=unique_pairs,
             text1_statistics=text1_statistics,
             text2_statistics=text2_statistics,

mteb/abstasks/AbsTaskBitextMining.py

@@ -7,9 +7,10 @@
 
 from mteb.encoder_interface import Encoder
 from mteb.types import HFSubset, ScoresDict
-from mteb.types.statistics import DescriptiveStatistics
+from mteb.types.statistics import DescriptiveStatistics, TextStatistics
 
 from ..evaluation.evaluators import BitextMiningEvaluator
+from ._statistics_calculation import calculate_text_statistics
 from .AbsTask import AbsTask
 
 logger = logging.getLogger(__name__)
@@ -23,29 +24,16 @@ class BitextDescriptiveStatistics(DescriptiveStatistics):
         number_of_characters: Total number of symbols in the dataset.
         unique_pairs: Number of duplicate pairs
 
-        min_sentence1_length: Minimum length of sentence1
-        average_sentence1_length: Average length of sentence1
-        max_sentence1_length: Maximum length of sentence1
-        unique_sentence1: Number of duplicates in sentence1
-
-        min_sentence2_length: Minimum length of sentence2
-        average_sentence2_length: Average length of sentence2
-        max_sentence2_length: Maximum length of sentence2
+        sentence1_statistics: Statistics for sentence1
+        sentence2_statistics: Statistics for sentence2
     """
 
     num_samples: int
     number_of_characters: int
     unique_pairs: int
 
-    min_sentence1_length: int
-    average_sentence1_length: float
-    max_sentence1_length: int
-    unique_sentence1: int
-
-    min_sentence2_length: int
-    average_sentence2_length: float
-    max_sentence2_length: int
-    unique_sentence2: int
+    sentence1_statistics: TextStatistics
+    sentence2_statistics: TextStatistics
 
 
 class AbsTaskBitextMining(AbsTask):
@@ -175,26 +163,20 @@ def _calculate_metrics_from_split(
             sent_1, sent_2 = pairs_cols[0]
             sentence1 = self.dataset[split][sent_1]
             sentence2 = self.dataset[split][sent_2]
-        s1_len = [len(s1) for s1 in sentence1]
-        s2_len = [len(s2) for s2 in sentence2]
-        total_s1_len = sum(s1_len)
-        total_s2_len = sum(s2_len)
 
+        text1_statistics = calculate_text_statistics(sentence1)
+        text2_statistics = calculate_text_statistics(sentence2)
         unique_pairs = len(set(zip(sentence1, sentence2)))
-        unique_sentence1 = len(set(sentence1))
-        unique_sentence2 = len(set(sentence2))
+
         return BitextDescriptiveStatistics(
             num_samples=len(sentence1),
-            number_of_characters=total_s1_len + total_s2_len,
+            number_of_characters=(
+                text1_statistics["total_text_length"]
+                + text2_statistics["total_text_length"]
+            ),
             unique_pairs=unique_pairs,
-            min_sentence1_length=min(s1_len),
-            average_sentence1_length=sum(s1_len) / len(sentence1),
-            max_sentence1_length=max(s1_len),
-            unique_sentence1=unique_sentence1,
-            min_sentence2_length=min(s2_len),
-            average_sentence2_length=total_s2_len / len(sentence2),
-            max_sentence2_length=max(s2_len),
-            unique_sentence2=unique_sentence2,
+            sentence1_statistics=text1_statistics,
+            sentence2_statistics=text2_statistics,
         )
 
     def _push_dataset_to_hub(self, repo_name: str) -> None: