from pprint import pprint
from kotoba_whisper import KotobaWhisperPipeline
from transformers.pipelines import PIPELINE_REGISTRY, pipeline
from transformers import WhisperForConditionalGeneration, TFWhisperForConditionalGeneration


model_alias = "kotoba-tech/kotoba-whisper-v2.2"
PIPELINE_REGISTRY.register_pipeline(
    "kotoba-whisper",
    pipeline_class=KotobaWhisperPipeline,
    pt_model=WhisperForConditionalGeneration,
    tf_model=TFWhisperForConditionalGeneration
)
test_audio = "/Users/asahiu/Desktop/speaker_diariazation_sample_1.wav"
pipe = pipeline(task="kotoba-whisper", model="kotoba-tech/kotoba-whisper-v2.0", chunk_length_s=15, batch_size=16, return_unique_speaker=True)
output = pipe(test_audio)
pprint(output)
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)