init

pipeline/kotoba_whisper.py

@@ -20,7 +20,7 @@ class Punctuator:
 
     ja_punctuations = ["!", "?", "、", "。"]
 
-    def __init__(self, model: str = "pcs_47lang"):
+    def __init__(self, model: str = "1-800-BAD-CODE/xlm-roberta_punctuation_fullstop_truecase"):
         self.punctuation_model = PunctCapSegModelONNX.from_pretrained(model)
 
     def punctuate(self, text: str) -> str:
@@ -123,7 +123,7 @@ class KotobaWhisperPipeline(AutomaticSpeechRecognitionPipeline):
         }
         postprocess_params = {"add_punctuation": add_punctuation, "return_timestamps": True, "return_language": False}
         forward_params = {} if generate_kwargs is None else generate_kwargs
-        forward_params.update({"max_new_tokens": max_new_tokens, "return_timestamps": True})
+        forward_params.update({"max_new_tokens": max_new_tokens, "return_timestamps": True, "language": "ja", "task": "transcribe"})
         return preprocess_params, forward_params, postprocess_params
 
     def preprocess(self,

pipeline/push_pipeline.py

@@ -14,8 +14,6 @@ PIPELINE_REGISTRY.register_pipeline(
     tf_model=TFWhisperForConditionalGeneration
 )
 pipe = pipeline(task="kotoba-whisper", model="kotoba-tech/kotoba-whisper-v2.0", chunk_length_s=15, batch_size=16)
-output = pipe(test_audio, add_punctuation=True)
-pprint(output)
 pipe.push_to_hub(model_alias)
 
 

pipeline/test_pipeline.py

@@ -6,6 +6,5 @@ pipe = pipeline(model="kotoba-tech/kotoba-whisper-v2.2", chunk_length_s=None, ba
 output = pipe("sample_diarization_japanese.mp3")
 pprint(output)
 
-pipe = pipeline(model="kotoba-tech/kotoba-whisper-v2.2", chunk_length_s=None, batch_size=16, trust_remote_code=True, return_unique_speaker=False)
-output = pipe("sample_diarization_japanese.mp3")
+output = pipe("sample_diarization_japanese.mp3", add_punctu)
 pprint(output)