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


model_alias = "kotoba-tech/kotoba-whisper-v2.1"
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,
    punctuator=True,
    stable_ts=True,
)
pipe.push_to_hub(model_alias)
pipe = pipeline(model=model_alias,
                punctuator=True,
                stable_ts=True,
                chunk_length_s=15,
                batch_size=16,
                trust_remote_code=True)