Add ESC50 and zero-shot classification

mteb/abstasks/Audio/AbsTaskAudioZeroshotClassification.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from datasets import Dataset
+
+from ...encoder_interface import Encoder
+from ...evaluation.evaluators import AudioZeroshotClassificationEvaluator
+from ..AbsTask import AbsTask, ScoresDict
+
+logger = logging.getLogger(__name__)
+
+
+class AbsTaskAudioZeroshotClassification(AbsTask):
+    """Abstract class for ZeroshotClassification tasks
+    The similarity between audio and candidate text prompts, such as as an audio wav of a dog barking and candidate text prompts like "Sound of a dog barking" or "Sound of a airplane".
+
+    self.load_data() must generate a huggingface dataset with a split matching self.metadata_dict["eval_splits"], and assign it to self.dataset. It must contain the following columns:
+        image: list of Image.Image
+        labels: list of int
+    """
+
+    audio_column_name: str = "audio"
+    label_column_name: str = "target"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def _add_main_score(self, scores) -> None:
+        scores["main_score"] = scores[self.metadata.main_score]
+
+    def _calculate_metrics_from_split(
+        self, split: str, hf_subset: str | None = None, compute_overall: bool = False
+    ):
+        pass
+
+    def _evaluate_subset(
+        self,
+        model: Encoder,
+        dataset: Dataset,
+        *,
+        encode_kwargs: dict[str, Any] = {},
+        **kwargs,
+    ) -> ScoresDict:
+        candidate_labels = self.get_candidate_labels()
+
+        evaluator = AudioZeroshotClassificationEvaluator(
+            dataset,
+            self.audio_column_name,
+            self.label_column_name,
+            candidate_labels,
+            task_name=self.metadata.name,
+            **kwargs,
+        )
+        metrics = evaluator(model, encode_kwargs=encode_kwargs)
+
+        scores = {
+            "accuracy": metrics["accuracy"],
+            "f1": metrics["f1"],
+            "f1_weighted": metrics["f1_weighted"],
+            "precision": metrics["precision"],
+            "recall": metrics["recall"],
+        }
+        self._add_main_score(scores)
+        return scores
+
+    def get_candidate_labels(self) -> list[str]:
+        """Return the text candidates for zeroshot classification"""
+        raise NotImplementedError("This method should be overridden by subclasses")

mteb/abstasks/TaskMetadata.py

@@ -125,6 +125,7 @@
     "VisualSTS",
     "ZeroShotClassification",
     "AudioMultilabelClassification",
+    "AudioZeroshotClassification",
 ]
 
 

mteb/evaluation/evaluators/Audio/ZeroshotClassificationEvaluator.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import io
+import logging
+import math
+import os
+from typing import Any
+
+import torch
+import torchaudio
+from sklearn import metrics
+from torch.utils.data import DataLoader
+
+from mteb.encoder_interface import Encoder
+
+from ..Evaluator import Evaluator
+
+logger = logging.getLogger(__name__)
+
+
+class AudioDataset(torch.utils.data.Dataset):
+    def __init__(self, hf_dataset, audio_column_name: str = "image", transform=None):
+        self.dataset = hf_dataset
+        self.transform = transform
+        self.audio_column_name = audio_column_name
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        audio = self.dataset[idx][self.audio_column_name]
+        if isinstance(audio, bytes):
+            waveform, sample_rate = torchaudio.load(io.BytesIO(audio))
+        elif isinstance(audio, str):
+            # Assuming audio is a file path
+            waveform, sample_rate = torchaudio.load(audio)
+        else:
+            # Assume audio is already a tensor or in a usable format
+            waveform = audio
+        if self.transform is not None:
+            waveform = self.transform(waveform)
+        return waveform
+
+
+def custom_collate_fn(batch):
+    return batch
+
+
+class AudioZeroshotClassificationEvaluator(Evaluator):
+    def __init__(
+        self,
+        dataset,
+        audio_column_name: str,
+        label_column_name: str,
+        candidate_labels: list[str],
+        task_name: str | None = None,
+        transform=None,
+        batch_size: int = 32,
+        **kwargs,
+    ):
+        """Initialize zero-shot audio classification evaluator.
+
+        Args:
+            dataset: HuggingFace dataset containing audio data
+            audio_column_name: Name of column containing audio data
+            label_column_name: Name of column containing label indices
+            candidate_labels: List of text descriptions for possible classes
+            task_name: Optional name of the task
+            transform: Optional audio transforms
+            batch_size: Batch size for processing
+            **kwargs: Additional keyword arguments
+        """
+        super().__init__(**kwargs)
+        self.dataset = AudioDataset(
+            dataset, audio_column_name=audio_column_name, transform=transform
+        )
+        self.labels = dataset[label_column_name]
+        self.candidate_labels = candidate_labels
+        self.task_name = task_name
+        self.batch_size = batch_size
+
+    def __call__(
+        self, model: Encoder, *, encode_kwargs: dict[str, Any] = {}
+    ) -> dict[str, float]:
+        """Evaluate zero-shot classification performance."""
+        logger.info("Getting text embeddings for candidate labels...")
+
+        text_embeddings = model.get_text_embeddings(self.candidate_labels)
+
+        logger.info("Processing audio data...")
+        dataloader = DataLoader(
+            self.dataset,
+            batch_size=encode_kwargs.get("batch_size", self.batch_size),
+            collate_fn=custom_collate_fn,
+            num_workers=min(math.floor(os.cpu_count() / 2), 16),
+        )
+
+        audio_embeddings = model.get_audio_embeddings(dataloader)
+
+        # Calculate similarity scores
+        similarity = (
+            torch.from_numpy(audio_embeddings) @ torch.from_numpy(text_embeddings).T
+        )
+
+        predictions = similarity.argmax(dim=1).cpu().numpy()
+
+        # Calculate metrics
+        scores = {
+            "accuracy": metrics.accuracy_score(self.labels, predictions),
+            "f1": metrics.f1_score(self.labels, predictions, average="macro"),
+            "f1_weighted": metrics.f1_score(self.labels, predictions, average="macro"),
+            "precision": metrics.precision_score(
+                self.labels, predictions, average="macro"
+            ),
+            "recall": metrics.recall_score(self.labels, predictions, average="macro"),
+        }
+
+        return scores

mteb/evaluation/evaluators/__init__.py

@@ -3,6 +3,7 @@
 from .Audio.Any2AnyRetrievalEvaluator import *
 from .Audio.ClassificationEvaluator import *
 from .Audio.ClusteringEvaluator import *
+from .Audio.ZeroshotClassificationEvaluator import *
 from .BitextMiningEvaluator import *
 from .ClassificationEvaluator import *
 from .ClusteringEvaluator import *