TheirStory/custom-whisper-handler

1from typing import Dict2from transformers.pipelines.audio_utils import ffmpeg_read3import whisper4import torch5 6SAMPLE_RATE = 160007 8class EndpointHandler():9    def __init__(self, path=""):10        # load the model11        self.model = whisper.load_model("large")12 13    def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:14        """15        Args:16            data (:obj:):17                includes the deserialized audio file as bytes18        Return:19            A :obj:`dict`:. base64 encoded image20        """21        # process input22        inputs = data.pop("inputs", data)23        audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)24        audio_tensor = torch.from_numpy(audio_nparray)25 26        # run inference pipeline27        result = self.model.transcribe(audio_nparray)28 29        # postprocess the prediction30        return result31