Create handler.py

handler.py

@@ -0,0 +1,42 @@
+from typing import Dict, List, Any
+from transformers import pipeline
+from pipeline import SpeechToJyutpingPipeline
+from model import Wav2Vec2BertForCantonese
+from transformers.pipelines import PIPELINE_REGISTRY
+from transformers import Wav2Vec2CTCTokenizer, SeamlessM4TFeatureExtractor, pipeline
+from model import Wav2Vec2BertForCantonese
+
+PIPELINE_REGISTRY.register_pipeline(
+    "speech-to-jyutping",
+    pipeline_class=SpeechToJyutpingPipeline,
+)
+
+
+class EndpointHandler:
+    def __init__(self, path="hon9kon9ize/wav2vec2bert-jyutping"):
+        feature_extractor = SeamlessM4TFeatureExtractor.from_pretrained(path)
+        tokenizer = Wav2Vec2CTCTokenizer(
+            "vocab.json", unk_token="[UNK]", pad_token="[PAD]", word_delimiter_token="|"
+        )
+
+        self.pipeline = pipeline(
+            task="speech-to-jyutping",
+            model=Wav2Vec2BertForCantonese.from_pretrained(path),
+            feature_extractor=feature_extractor,
+            tokenizer=tokenizer,
+        )
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+         data args:
+              inputs (:obj: `str`)
+        Return:
+              A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # get inputs
+        inputs = data.pop("inputs", data)
+
+        # run normal prediction
+        prediction = self.pipeline(inputs)
+
+        return prediction