HFEndpoint handler created

handler.py

@@ -0,0 +1,46 @@
+from typing import Any, Dict, List
+
+import numpy as np
+import torch
+from transformers import WhisperForConditionalGeneration, WhisperProcessor, pipeline
+from transformers.pipelines.audio_utils import ffmpeg_read
+
+SAMPLE_RATE = 16000
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        torch_dtype = torch.float16
+        device = "cuda"
+
+        whisper_processor = WhisperProcessor.from_pretrained(path)
+        whisper_model = WhisperForConditionalGeneration.from_pretrained(
+            path,
+            torch_dtype=torch_dtype,
+        ).to("cuda")
+
+        self.asr_pipeline = pipeline(
+            "automatic-speech-recognition",
+            model=whisper_model,
+            tokenizer=whisper_processor.tokenizer,
+            feature_extractor=whisper_processor.feature_extractor,
+            chunk_length_s=30,
+            torch_dtype=torch_dtype,
+            device=device,
+            model_kwargs={"language": "pt"},
+        )
+
+    def parse_audio(self, audio_bytes):
+        audio_nparray = ffmpeg_read(audio_bytes, SAMPLE_RATE)
+
+        if len(audio_nparray.shape) > 1 and audio_nparray.shape[1] == 2:
+            return np.mean(audio_nparray, axis=1)
+
+        return audio_nparray
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        parameters = data.get("parameters", {})
+
+        audio = self.parse_audio(data["inputs"])
+
+        return self.asr_pipeline(audio, **parameters)