"""wget https://huggingface.co/kotoba-tech/kotoba-whisper-v2.2/resolve/main/sample_audio/sample_diarization_japanese.mp3"""
from pprint import pprint
from kotoba_whisper import KotobaWhisperPipeline
from transformers.pipelines import PIPELINE_REGISTRY, pipeline
from transformers import WhisperForConditionalGeneration, TFWhisperForConditionalGeneration


test_audio = "sample_diarization_japanese.mp3"
model_alias = "kotoba-tech/kotoba-whisper-v2.2"
PIPELINE_REGISTRY.register_pipeline(
    "kotoba-whisper",
    pipeline_class=KotobaWhisperPipeline,
    pt_model=WhisperForConditionalGeneration,
    tf_model=TFWhisperForConditionalGeneration
)
pipe = pipeline(task="kotoba-whisper", model="kotoba-tech/kotoba-whisper-v2.0", batch_size=16)
# pprint(pipe("fukabori_trimmed_more_more.mp3"))
pprint(pipe(
    "fukabori_trimmed_more_more.mp3",
    add_silence_end=0.5,
    add_silence_start=0.5,
    chunk_length_s=15,
))
pipe.push_to_hub(model_alias)