Adding SD for superb (speech-classification). (#225)

* Adding SD for superb (speech-classification).

* Style.

* Style.

* Test dependencies.

* Addressing @omar 's comments.

* Forgot pytest..

* `speech-classification` -> `speech-segmentation`.

* Fixing test cache in common + adding simple tests.

Author: Narsil

.github/workflows/python-api-tests.yaml

@@ -25,7 +25,6 @@ jobs:
         working-directory: api-inference-community
         run: |
           pip install --upgrade pip
-          pip install pytest pillow httpx huggingface_hub
-          pip install -e .
+          pip install -e .[test]
       - run: make test
         working-directory: api-inference-community

api-inference-community/api_inference_community/normalizers.py

@@ -0,0 +1,41 @@
+"""
+Helper classes to modify pipeline outputs from tensors to expected pipeline output
+"""
+
+from typing import TYPE_CHECKING, Dict, List, Union
+
+
+Classes = Dict[str, Union[str, float]]
+
+if TYPE_CHECKING:
+    try:
+        import torch
+    except Exception:
+        pass
+
+
+def speaker_diarization_normalize(
+    tensor: "torch.Tensor", sampling_rate: int, classnames: List[str]
+) -> List[Classes]:
+    N = tensor.shape[1]
+    if len(classnames) != N:
+        raise ValueError(
+            f"There is a mismatch between classnames ({len(classnames)}) and number of speakers ({N})"
+        )
+    classes = []
+    for i in range(N):
+        values, counts = tensor[:, i].unique_consecutive(return_counts=True)
+        offset = 0
+        for v, c in zip(values, counts):
+            if v == 1:
+                classes.append(
+                    {
+                        "class": classnames[i],
+                        "start": offset / sampling_rate,
+                        "end": (offset + c.item()) / sampling_rate,
+                    }
+                )
+            offset += c.item()
+
+    classes = sorted(classes, key=lambda x: x["start"])
+    return classes

api-inference-community/api_inference_community/validation.py

@@ -211,12 +211,14 @@ def normalize_payload(
     if task in {
         "automatic-speech-recognition",
         "audio-to-audio",
+        "speech-segmentation",
     }:
         if sampling_rate is None:
             raise EnvironmentError(
                 "We cannot normalize audio file if we don't know the sampling rate"
             )
-        return normalize_payload_audio(bpayload, sampling_rate)
+        outputs = normalize_payload_audio(bpayload, sampling_rate)
+        return outputs
     elif task in {
         "image-classification",
         "image-to-text",

api-inference-community/docker_images/common/app/pipelines/__init__.py

@@ -8,6 +8,7 @@
 from app.pipelines.image_classification import ImageClassificationPipeline
 from app.pipelines.question_answering import QuestionAnsweringPipeline
 from app.pipelines.sentence_similarity import SentenceSimilarityPipeline
+from app.pipelines.speech_segmentation import SpeechSegmentationPipeline
 from app.pipelines.structured_data_classification import (
     StructuredDataClassificationPipeline,
 )

api-inference-community/docker_images/common/app/pipelines/automatic_speech_recognition.py

@@ -21,9 +21,7 @@ def __call__(self, inputs: np.array) -> Dict[str, str]:
         """
         Args:
             inputs (:obj:`np.array`):
-                The raw waveform of audio received. By default at 16KHz.
-                Check `app.validation` if a different sample rate is required
-                or if it depends on the model
+                The raw waveform of audio received. By default at self.sampling_rate, otherwise 16KHz.
         Return:
             A :obj:`dict`:. The object return should be liked {"text": "XXX"} containing
             the detected langage from the input audio

api-inference-community/docker_images/common/app/pipelines/speech_segmentation.py

@@ -0,0 +1,34 @@
+from typing import Dict
+
+import numpy as np
+from app.pipelines import Pipeline
+
+
+class SpeechSegmentationPipeline(Pipeline):
+    def __init__(self, model_id: str):
+        # IMPLEMENT_THIS
+        # Preload all the elements you are going to need at inference.
+        # For instance your model, processors, tokenizer that might be needed.
+        # This function is only called once, so do all the heavy processing I/O here
+        # IMPLEMENT_THIS : Please define a `self.sampling_rate` for this pipeline
+        # to automatically read the input correctly
+        self.sampling_rate = 16000
+        raise NotImplementedError(
+            "Please implement SpeechSegmentationPipeline __init__ function"
+        )
+
+    def __call__(self, inputs: np.array) -> Dict[str, str]:
+        """
+        Args:
+            inputs (:obj:`np.array`):
+                The raw waveform of audio received. By default at self.sampling_rate, otherwise 16KHz.
+        Return:
+            A :obj:`list`:. Each item in the list is like {"class": "XXX", "start": float, "end": float}
+            "class" is the associated class of the audio segment, "start" and "end" are markers expressed in seconds
+            within the audio file.
+        """
+        # IMPLEMENT_THIS
+        # api_inference_community.normalizers.speaker_diarization_normalize could help.
+        raise NotImplementedError(
+            "Please implement SpeechSegmentationPipeline __call__ function"
+        )

api-inference-community/docker_images/common/tests/test_api_audio_to_audio.py

@@ -24,6 +24,12 @@ def setUp(self):
 
         self.app = app
 
+    @classmethod
+    def setUpClass(cls):
+        from app.main import get_pipeline
+
+        get_pipeline.cache_clear()
+
     def tearDown(self):
         if self.old_model_id is not None:
             os.environ["MODEL_ID"] = self.old_model_id

api-inference-community/docker_images/common/tests/test_api_automatic_speech_recognition.py

@@ -22,6 +22,12 @@ def setUp(self):
 
         self.app = app
 
+    @classmethod
+    def setUpClass(cls):
+        from app.main import get_pipeline
+
+        get_pipeline.cache_clear()
+
     def tearDown(self):
         if self.old_model_id is not None:
             os.environ["MODEL_ID"] = self.old_model_id

api-inference-community/docker_images/common/tests/test_api_feature_extraction.py

@@ -22,6 +22,12 @@ def setUp(self):
 
         self.app = app
 
+    @classmethod
+    def setUpClass(cls):
+        from app.main import get_pipeline
+
+        get_pipeline.cache_clear()
+
     def tearDown(self):
         if self.old_model_id is not None:
             os.environ["MODEL_ID"] = self.old_model_id

api-inference-community/docker_images/common/tests/test_api_image_classification.py

@@ -22,6 +22,12 @@ def setUp(self):
 
         self.app = app
 
+    @classmethod
+    def setUpClass(cls):
+        from app.main import get_pipeline
+
+        get_pipeline.cache_clear()
+
     def tearDown(self):
         if self.old_model_id is not None:
             os.environ["MODEL_ID"] = self.old_model_id