custom handler

handler.py

@@ -0,0 +1,45 @@
+"""Custom handler required to process model in custom way."""
+from typing import Dict, Any, List
+from transformers import pipeline
+import torch
+
+
+class EndpointHandler:
+    """HF class for custom model processing for an inference endpoint."""
+
+    def __init__(self, path: str) -> None:
+        """_summary_
+
+        Args:
+            path (str): Path to the model weights allowing to load the model.
+        """
+        device = 0 if torch.cuda.is_available() else "cpu"
+        self.pipeline = pipeline(
+            "automatic-speech-recognition",
+            model=path,
+            device=device,
+            chunk_length_s=30,
+            )
+
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """Custom processing pipeline.
+
+        Args:
+            request (Dict[str, Any]): _description_
+
+        Returns:
+            List[Dict[str, Any]]: _description_
+        """
+        audio = data.pop("audio")  # bytes
+        lang = data.pop("language")
+        task = data.pop("task")  # One of "translate", "transcribe"
+
+        # Set language and task
+        self.pipeline.model.config.forced_decoder_ids = self.pipeline.tokenizer.get_decoder_prompt_ids(language=lang, task=task)
+        self.pipeline
+
+        # Model inference
+        output = self.pipeline(audio, return_timestamps="word")
+
+        return output

test_handler.py

@@ -0,0 +1,16 @@
+"""Test the custom handler."""
+from handler import EndpointHandler
+import os
+
+handler = EndpointHandler(path="whisper-tiny")
+
+with open("frank.wav", "rb") as file:
+    audio = file.read()
+
+payload = {"audio": audio, "language": "nl", "task": "transcribe"}
+
+output = handler(payload)
+
+print(output["text"])
+print("\n")
+print(output["chunks"])