Whisper sample script

make_subtitle.py

from pathlib import Path
import whisper
from whisper.utils import get_writer


BASE = Path(__file__).parent


def main(folder: Path, ext: str = "wav", lang: str = "Japanese", model_name: str = "large-v2"):
    model = whisper.load_model(model_name)
    targets = folder.glob(f"*.{ext}")
    writer_args = {
        "highlight_words": False,
        "max_line_count": None,
        "max_line_width": None
    }
    for target in targets:
        result = model.transcribe(str(target), verbose=True, language=lang)
        writer_srt = get_writer("srt", folder)
        writer_srt(result, str(target), writer_args)
        writer_txt = get_writer("txt", folder)
        writer_txt(result, str(target), writer_args)


if __name__ == "__main__":
    """
    $ python make_subtitle.py FOLDERNAME EXT
    $ python make_subtitle.py FOLDERNAME mp3
    """
    import sys
    args = sys.argv
    if len(args) != 3:
        print("No args, need to 2 args: folder-name, ext")
    else:
        folder_name = args[1]
        ext = args[2]
        fullpass_folder = BASE / folder_name
        main(fullpass_folder, ext)