"""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", chunk_length_s=15, batch_size=16)
output = pipe(test_audio)
pprint(output)
pipe.push_to_hub(model_alias)