Allow generic inference for ASR for superb (#185)

Author: osanseviero

api-inference-community/docker_images/superb/Dockerfile

@@ -3,6 +3,10 @@ LABEL maintainer="Omar [email protected]"
 
 # Add any system dependency here
 # RUN apt-get update -y && apt-get install libXXX -y
+RUN apt-get update -y
+RUN apt-get install git -y
+RUN apt-get install libsndfile1 -y
+RUN apt-get install ffmpeg -y
 
 COPY ./requirements.txt /app
 RUN pip install --no-cache-dir -r requirements.txt
@@ -16,6 +20,7 @@ COPY ./prestart.sh /app/
 ARG max_workers=1
 ENV MAX_WORKERS=$max_workers
 ENV HUGGINGFACE_HUB_CACHE=/data
+ENV TORCH_HOME=/data
 
 # Necessary on GPU environment docker.
 # TIMEOUT env variable is used by nvcr.io/nvidia/pytorch:xx for another purpose

api-inference-community/docker_images/superb/app/main.py

@@ -4,7 +4,7 @@
 from typing import Dict, Type
 
 from api_inference_community.routes import pipeline_route, status_ok
-from app.pipelines import Pipeline
+from app.pipelines import AutomaticSpeechRecognitionPipeline, Pipeline
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.middleware.gzip import GZipMiddleware
@@ -32,7 +32,9 @@
 # ALLOWED_TASKS = {"automatic-speech-recognition": AutomaticSpeechRecognitionPipeline}
 # You can check the requirements and expectations of each pipelines in their respective
 # directories. Implement directly within the directories.
-ALLOWED_TASKS: Dict[str, Type[Pipeline]] = {}
+ALLOWED_TASKS: Dict[str, Type[Pipeline]] = {
+    "automatic-speech-recognition": AutomaticSpeechRecognitionPipeline,
+}
 
 
 @functools.lru_cache()

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

@@ -1,12 +1,5 @@
 from app.pipelines.base import Pipeline, PipelineException  # isort:skip
 
-from app.pipelines.audio_to_audio import AudioToAudioPipeline
 from app.pipelines.automatic_speech_recognition import (
     AutomaticSpeechRecognitionPipeline,
 )
-from app.pipelines.feature_extraction import FeatureExtractionPipeline
-from app.pipelines.image_classification import ImageClassificationPipeline
-from app.pipelines.question_answering import QuestionAnsweringPipeline
-from app.pipelines.sentence_similarity import SentenceSimilarityPipeline
-from app.pipelines.text_to_speech import TextToSpeechPipeline
-from app.pipelines.token_classification import TokenClassificationPipeline

api-inference-community/docker_images/superb/app/pipelines/audio_to_audio.py

@@ -1,7 +1,11 @@
+import os
+import subprocess
+import sys
 from typing import Dict
 
 import numpy as np
 from app.pipelines import Pipeline
+from huggingface_hub import snapshot_download
 
 
 class AutomaticSpeechRecognitionPipeline(Pipeline):
@@ -12,10 +16,24 @@ def __init__(self, model_id: str):
         # 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
+        filepath = snapshot_download(model_id)
+        sys.path.append(filepath)
+        if "requirements.txt" in os.listdir(filepath):
+            subprocess.check_call(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "install",
+                    "-r",
+                    os.path.join(filepath, "requirements.txt"),
+                ]
+            )
+
+        from model import PreTrainedModel
+
+        self.model = PreTrainedModel(filepath)
         self.sampling_rate = 16000
-        raise NotImplementedError(
-            "Please implement AutomaticSpeechRecognitionPipeline __init__ function"
-        )
 
     def __call__(self, inputs: np.array) -> Dict[str, str]:
         """
@@ -28,7 +46,4 @@ def __call__(self, inputs: np.array) -> Dict[str, str]:
             A :obj:`dict`:. The object return should be liked {"text": "XXX"} containing
             the detected langage from the input audio
         """
-        # IMPLEMENT_THIS
-        raise NotImplementedError(
-            "Please implement AutomaticSpeechRecognitionPipeline __call__ function"
-        )
+        return self.model(inputs)