Update handler.py

handler.py

@@ -3,6 +3,7 @@ import re
 from itertools import groupby
 from dataclasses import dataclass
 from typing import Optional, Tuple, Union, Dict, List, Any
+from huggingface_hub import hf_hub_download
 import torch
 import torch.nn as nn
 from transformers.modeling_outputs import ModelOutput
@@ -47,8 +48,11 @@ ONSETS = {
 
 class SpeechToJyutpingPipeline(Pipeline):
     def _sanitize_parameters(self, **kwargs):
+        tone_vocab_file = hf_hub_download(
+            repo_id="hon9kon9ize/wav2vec2bert-jyutping", filename="tone_vocab.json"
+        )
         self.tone_tokenizer = Wav2Vec2CTCTokenizer(
-            "tone_vocab.json",
+            tone_vocab_file,
             unk_token="[UNK]",
             pad_token="[PAD]",
             word_delimiter_token="|",
@@ -95,7 +99,6 @@ class SpeechToJyutpingPipeline(Pipeline):
 
         sample_rate = 16000
         symbols = [w for w in transcription.split(" ") if len(w) > 0]
-        duration_sec = model_outputs["duration"] / sample_rate
 
         ids_w_index = [(i, _id.item()) for i, _id in enumerate(predicted_ids[0])]
         # remove entries which are just "padding" (i.e. no characers are recognized)
@@ -151,7 +154,7 @@ class SpeechToJyutpingPipeline(Pipeline):
         transcription = re.sub(
             r"\s+", " ", "".join(transcription).replace("_", " ").strip()
         )
-        tone_probs = torch.stack(tone_probs).cpu().numpy()
+        tone_probs = torch.stack(tone_probs).cpu().tolist()
 
         return {"transcription": transcription, "tone_probs": tone_probs}
 
@@ -388,15 +391,14 @@ class Wav2Vec2BertForCantonese(Wav2Vec2BertPreTrainedModel):
 
 
 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="|"
-        )
+    def __init__(self, path="."):
+        model_path = "hon9kon9ize/wav2vec2bert-jyutping"
+        feature_extractor = SeamlessM4TFeatureExtractor.from_pretrained(model_path)
+        tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(model_path)
 
         self.pipeline = pipeline(
             task="speech-to-jyutping",
-            model=Wav2Vec2BertForCantonese.from_pretrained(path),
+            model=Wav2Vec2BertForCantonese.from_pretrained(model_path),
             feature_extractor=feature_extractor,
             tokenizer=tokenizer,
         )