Revert "whisperX merged in"

This reverts commit 9c2d684d1a778e559b6b13bb90112712dbc20568.

.gitignore

@@ -1,2 +0,0 @@
-whisperx.egg-info/
-**/__pycache__/

EXAMPLES.md

@@ -1,37 +0,0 @@
-# More Examples
-
-## Other Languages
-
-For non-english ASR, it is best to use the `large` whisper model. Alignment models are automatically picked by the chosen language from the default [lists](https://github.com/m-bain/whisperX/blob/main/whisperx/alignment.py#L18).
-
-Currently support default models tested for {en, fr, de, es, it, ja, zh, nl}
-
-
-If the detected language is not in this list, you need to find a phoneme-based ASR model from [huggingface model hub](https://huggingface.co/models) and test it on your data.
-
-### French
-    whisperx --model large --language fr examples/sample_fr_01.wav
-
-
-https://user-images.githubusercontent.com/36994049/208298804-31c49d6f-6787-444e-a53f-e93c52706752.mov
-
-
-### German
-    whisperx --model large --language de examples/sample_de_01.wav
-
-
-https://user-images.githubusercontent.com/36994049/208298811-e36002ba-3698-4731-97d4-0aebd07e0eb3.mov
-
-
-### Italian
-    whisperx --model large --language de examples/sample_it_01.wav
-
-
-https://user-images.githubusercontent.com/36994049/208298819-6f462b2c-8cae-4c54-b8e1-90855794efc7.mov
-
-
-### Japanese
-    whisperx --model large --language ja examples/sample_ja_01.wav
-
-
-https://user-images.githubusercontent.com/19920981/208731743-311f2360-b73b-4c60-809d-aaf3cd7e06f4.mov

LICENSE

@@ -1,27 +0,0 @@
-Copyright (c) 2022, Max Bain
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-1. Redistributions of source code must retain the above copyright
-   notice, this list of conditions and the following disclaimer.
-2. Redistributions in binary form must reproduce the above copyright
-   notice, this list of conditions and the following disclaimer in the
-   documentation and/or other materials provided with the distribution.
-3. All advertising materials mentioning features or use of this software
-   must display the following acknowledgement:
-   This product includes software developed by Max Bain.
-4. Neither the name of Max Bain nor the
-   names of its contributors may be used to endorse or promote products
-   derived from this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER ''AS IS'' AND ANY
-EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE 
-USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

MANIFEST.in

@@ -1,4 +0,0 @@
-include whisperx/assets/*
-include whisperx/assets/gpt2/*
-include whisperx/assets/multilingual/*
-include whisperx/normalizers/english.json

requirements.txt

@@ -1,10 +0,0 @@
-numpy
-pandas
-torch >=1.9
-torchaudio >=0.10,<1.0
-tqdm
-more-itertools
-transformers>=4.19.0
-ffmpeg-python==0.2.0
-pyannote.audio
-openai-whisper==20230314

setup.py

@@ -1,28 +0,0 @@
-import os
-
-import pkg_resources
-from setuptools import setup, find_packages
-
-setup(
-    name="whisperx",
-    py_modules=["whisperx"],
-    version="2.0",
-    description="Time-Accurate Automatic Speech Recognition using Whisper.",
-    readme="README.md",
-    python_requires=">=3.8",
-    author="Max Bain",
-    url="https://github.com/m-bain/whisperx",
-    license="MIT",
-    packages=find_packages(exclude=["tests*"]),
-    install_requires=[
-        str(r)
-        for r in pkg_resources.parse_requirements(
-            open(os.path.join(os.path.dirname(__file__), "requirements.txt"))
-        )
-    ],
-    entry_points = {
-        'console_scripts': ['whisperx=whisperx.transcribe:cli'],
-    },
-    include_package_data=True,
-    extras_require={'dev': ['pytest']},
-)

whisperx/__init__.py

@@ -1,3 +0,0 @@
-from .transcribe import transcribe, transcribe_with_vad
-from .alignment import load_align_model, align
-from .vad import load_vad_model

whisperx/__main__.py

@@ -1,4 +0,0 @@
-from .transcribe import cli
-
-
-cli()

whisperx/alignment.py

@@ -1,548 +0,0 @@
-""""
-Forced Alignment with Whisper
-C. Max Bain
-"""
-import numpy as np
-import pandas as pd
-from typing import List, Union, Iterator, TYPE_CHECKING
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-import torchaudio
-import torch
-from dataclasses import dataclass
-from whisper.audio import SAMPLE_RATE, load_audio
-from .utils import interpolate_nans
-
-
-LANGUAGES_WITHOUT_SPACES = ["ja", "zh"]
-
-DEFAULT_ALIGN_MODELS_TORCH = {
-    "en": "WAV2VEC2_ASR_BASE_960H",
-    "fr": "VOXPOPULI_ASR_BASE_10K_FR",
-    "de": "VOXPOPULI_ASR_BASE_10K_DE",
-    "es": "VOXPOPULI_ASR_BASE_10K_ES",
-    "it": "VOXPOPULI_ASR_BASE_10K_IT",
-}
-
-DEFAULT_ALIGN_MODELS_HF = {
-    "ja": "jonatasgrosman/wav2vec2-large-xlsr-53-japanese",
-    "zh": "jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn",
-    "nl": "jonatasgrosman/wav2vec2-large-xlsr-53-dutch",
-    "uk": "Yehor/wav2vec2-xls-r-300m-uk-with-small-lm",
-    "pt": "jonatasgrosman/wav2vec2-large-xlsr-53-portuguese",
-    "ar": "jonatasgrosman/wav2vec2-large-xlsr-53-arabic",
-    "ru": "jonatasgrosman/wav2vec2-large-xlsr-53-russian",
-    "pl": "jonatasgrosman/wav2vec2-large-xlsr-53-polish",
-    "hu": "jonatasgrosman/wav2vec2-large-xlsr-53-hungarian",
-    "fi": "jonatasgrosman/wav2vec2-large-xlsr-53-finnish",
-    "fa": "jonatasgrosman/wav2vec2-large-xlsr-53-persian",
-    "el": "jonatasgrosman/wav2vec2-large-xlsr-53-greek",
-    "tr": "mpoyraz/wav2vec2-xls-r-300m-cv7-turkish",
-}
-
-
-def load_align_model(language_code, device, model_name=None):
-    if model_name is None:
-        # use default model
-        if language_code in DEFAULT_ALIGN_MODELS_TORCH:
-            model_name = DEFAULT_ALIGN_MODELS_TORCH[language_code]
-        elif language_code in DEFAULT_ALIGN_MODELS_HF:
-            model_name = DEFAULT_ALIGN_MODELS_HF[language_code]
-        else:
-            print(f"There is no default alignment model set for this language ({language_code}).\
-                Please find a wav2vec2.0 model finetuned on this language in https://huggingface.co/models, then pass the model name in --align_model [MODEL_NAME]")
-            raise ValueError(f"No default align-model for language: {language_code}")
-
-    if model_name in torchaudio.pipelines.__all__:
-        pipeline_type = "torchaudio"
-        bundle = torchaudio.pipelines.__dict__[model_name]
-        align_model = bundle.get_model().to(device)
-        labels = bundle.get_labels()
-        align_dictionary = {c.lower(): i for i, c in enumerate(labels)}
-    else:
-        try:
-            processor = Wav2Vec2Processor.from_pretrained(model_name)
-            align_model = Wav2Vec2ForCTC.from_pretrained(model_name)
-        except Exception as e:
-            print(e)
-            print(f"Error loading model from huggingface, check https://huggingface.co/models for finetuned wav2vec2.0 models")
-            raise ValueError(f'The chosen align_model "{model_name}" could not be found in huggingface (https://huggingface.co/models) or torchaudio (https://pytorch.org/audio/stable/pipelines.html#id14)')
-        pipeline_type = "huggingface"
-        align_model = align_model.to(device)
-        labels = processor.tokenizer.get_vocab()
-        align_dictionary = {char.lower(): code for char,code in processor.tokenizer.get_vocab().items()}
-
-    align_metadata = {"language": language_code, "dictionary": align_dictionary, "type": pipeline_type}
-
-    return align_model, align_metadata
-
-
-def align(
-    transcript: Iterator[dict],
-    model: torch.nn.Module,
-    align_model_metadata: dict,
-    audio: Union[str, np.ndarray, torch.Tensor],
-    device: str,
-    extend_duration: float = 0.0,
-    start_from_previous: bool = True,
-    interpolate_method: str = "nearest",
-):
-    """
-    Force align phoneme recognition predictions to known transcription
-
-    Parameters
-    ----------
-    transcript: Iterator[dict]
-        The Whisper model instance
-    
-    model: torch.nn.Module
-        Alignment model (wav2vec2)
-
-    audio: Union[str, np.ndarray, torch.Tensor]
-        The path to the audio file to open, or the audio waveform
-
-    device: str
-        cuda device
-
-    diarization: pd.DataFrame {'start': List[float], 'end': List[float], 'speaker': List[float]}
-        diarization segments with speaker labels.
-
-    extend_duration: float
-        Amount to pad input segments by. If not using vad--filter then recommended to use 2 seconds
-
-        If the gzip compression ratio is above this value, treat as failed
-
-    interpolate_method: str ["nearest", "linear", "ignore"]
-        Method to assign timestamps to non-aligned words. Words are not able to be aligned when none of the characters occur in the align model dictionary.
-        "nearest" copies timestamp of nearest word within the segment. "linear" is linear interpolation. "drop" removes that word from output.
-
-    Returns
-    -------
-    A dictionary containing the resulting text ("text") and segment-level details ("segments"), and
-    the spoken language ("language"), which is detected when `decode_options["language"]` is None.
-    """
-    if not torch.is_tensor(audio):
-        if isinstance(audio, str):
-            audio = load_audio(audio)
-        audio = torch.from_numpy(audio)
-    if len(audio.shape) == 1:
-        audio = audio.unsqueeze(0)
-
-    MAX_DURATION = audio.shape[1] / SAMPLE_RATE
-
-    model_dictionary = align_model_metadata["dictionary"]
-    model_lang = align_model_metadata["language"]
-    model_type = align_model_metadata["type"]
-
-    aligned_segments = []
-
-    prev_t2 = 0
-
-    char_segments_arr = {
-        "segment-idx": [],
-        "subsegment-idx": [],
-        "word-idx": [],
-        "char": [],
-        "start": [],
-        "end": [],
-        "score": [],
-    }
-
-    for sdx, segment in enumerate(transcript):
-        while True:
-            segment_align_success = False
-
-            # strip spaces at beginning / end, but keep track of the amount.
-            num_leading = len(segment["text"]) - len(segment["text"].lstrip())
-            num_trailing = len(segment["text"]) - len(segment["text"].rstrip())
-            transcription = segment["text"]
-
-            # TODO: convert number tokenizer / symbols to phonetic words for alignment.
-            # e.g. "$300" -> "three hundred dollars"
-            # currently "$300" is ignored since no characters present in the phonetic dictionary
-
-            # split into words
-            if model_lang not in LANGUAGES_WITHOUT_SPACES:
-                per_word = transcription.split(" ")
-            else:
-                per_word = transcription
-
-            # first check that characters in transcription can be aligned (they are contained in align model"s dictionary)
-            clean_char, clean_cdx = [], []
-            for cdx, char in enumerate(transcription):
-                char_ = char.lower()
-                # wav2vec2 models use "|" character to represent spaces
-                if model_lang not in LANGUAGES_WITHOUT_SPACES:
-                    char_ = char_.replace(" ", "|")
-                
-                # ignore whitespace at beginning and end of transcript
-                if cdx < num_leading:
-                    pass
-                elif cdx > len(transcription) - num_trailing - 1:
-                    pass
-                elif char_ in model_dictionary.keys():
-                    clean_char.append(char_)
-                    clean_cdx.append(cdx)
-
-            clean_wdx = []
-            for wdx, wrd in enumerate(per_word):
-                if any([c in model_dictionary.keys() for c in wrd]):
-                    clean_wdx.append(wdx)
-
-            # if no characters are in the dictionary, then we skip this segment...
-            if len(clean_char) == 0:
-                print(f'Failed to align segment ("{segment["text"]}"): no characters in this segment found in model dictionary, resorting to original...')
-                break          
-           
-            transcription_cleaned = "".join(clean_char)
-            tokens = [model_dictionary[c] for c in transcription_cleaned]
-
-            # we only pad if not using VAD filtering
-            if "seg_text" not in segment:
-                # pad according original timestamps
-                t1 = max(segment["start"] - extend_duration, 0)
-                t2 = min(segment["end"] + extend_duration, MAX_DURATION)
-
-            # use prev_t2 as current t1 if it"s later
-            if start_from_previous and t1 < prev_t2:
-                t1 = prev_t2
-
-            # check if timestamp range is still valid
-            if t1 >= MAX_DURATION:
-                print("Failed to align segment: original start time longer than audio duration, skipping...")
-                break
-            if t2 - t1 < 0.02:
-                print("Failed to align segment: duration smaller than 0.02s time precision")
-                break
-
-            f1 = int(t1 * SAMPLE_RATE)
-            f2 = int(t2 * SAMPLE_RATE)
-
-            waveform_segment = audio[:, f1:f2]
-
-            with torch.inference_mode():
-                if model_type == "torchaudio":
-                    emissions, _ = model(waveform_segment.to(device))
-                elif model_type == "huggingface":
-                    emissions = model(waveform_segment.to(device)).logits
-                else:
-                    raise NotImplementedError(f"Align model of type {model_type} not supported.")
-                emissions = torch.log_softmax(emissions, dim=-1)
-
-            emission = emissions[0].cpu().detach()
-
-            trellis = get_trellis(emission, tokens)
-            path = backtrack(trellis, emission, tokens)
-            if path is None:
-                print(f'Failed to align segment ("{segment["text"]}"): backtrack failed, resorting to original...')
-                break
-            char_segments = merge_repeats(path, transcription_cleaned)
-            # word_segments = merge_words(char_segments)
-            
-
-            # sub-segments
-            if "seg-text" not in segment:
-                segment["seg-text"] = [transcription]
-                
-            seg_lens = [0] + [len(x) for x in segment["seg-text"]]
-            seg_lens_cumsum = list(np.cumsum(seg_lens))
-            sub_seg_idx = 0
-
-            wdx = 0
-            duration = t2 - t1
-            ratio = duration * waveform_segment.size(0) / (trellis.size(0) - 1)
-            for cdx, char in enumerate(transcription + " "):
-                is_last = False
-                if cdx == len(transcription):
-                    break
-                elif cdx+1 == len(transcription):
-                    is_last = True
-
-                
-                start, end, score = None, None, None
-                if cdx in clean_cdx:
-                    char_seg = char_segments[clean_cdx.index(cdx)]
-                    start = char_seg.start * ratio + t1
-                    end = char_seg.end * ratio + t1
-                    score = char_seg.score        
-
-                char_segments_arr["char"].append(char)
-                char_segments_arr["start"].append(start)
-                char_segments_arr["end"].append(end)
-                char_segments_arr["score"].append(score)
-                char_segments_arr["word-idx"].append(wdx)
-                char_segments_arr["segment-idx"].append(sdx)
-                char_segments_arr["subsegment-idx"].append(sub_seg_idx)
-
-                # word-level info
-                if model_lang in LANGUAGES_WITHOUT_SPACES:
-                    # character == word
-                    wdx += 1
-                elif is_last or transcription[cdx+1] == " " or cdx == seg_lens_cumsum[sub_seg_idx+1] - 1:
-                    wdx += 1
-
-                if is_last or cdx == seg_lens_cumsum[sub_seg_idx+1] - 1:
-                    wdx = 0
-                    sub_seg_idx += 1
-
-            prev_t2 = segment["end"]
-
-            segment_align_success = True
-            # end while True loop
-            break
-
-        # reset prev_t2 due to drifting issues
-        if not segment_align_success:
-            prev_t2 = 0
-        
-    char_segments_arr = pd.DataFrame(char_segments_arr)
-    not_space = char_segments_arr["char"] != " "
-
-    per_seg_grp = char_segments_arr.groupby(["segment-idx", "subsegment-idx"], as_index = False)
-    char_segments_arr = per_seg_grp.apply(lambda x: x.reset_index(drop = True)).reset_index()
-    per_word_grp = char_segments_arr[not_space].groupby(["segment-idx", "subsegment-idx", "word-idx"])
-    per_subseg_grp = char_segments_arr[not_space].groupby(["segment-idx", "subsegment-idx"])
-    per_seg_grp = char_segments_arr[not_space].groupby(["segment-idx"])
-    char_segments_arr["local-char-idx"] = char_segments_arr.groupby(["segment-idx", "subsegment-idx"]).cumcount()
-    per_word_grp = char_segments_arr[not_space].groupby(["segment-idx", "subsegment-idx", "word-idx"]) # regroup
-
-    word_segments_arr = {}
-
-    # start of word is first char with a timestamp
-    word_segments_arr["start"] = per_word_grp["start"].min().values
-    # end of word is last char with a timestamp
-    word_segments_arr["end"] = per_word_grp["end"].max().values
-    # score of word is mean (excluding nan)
-    word_segments_arr["score"] = per_word_grp["score"].mean().values
-
-    word_segments_arr["segment-text-start"] = per_word_grp["local-char-idx"].min().astype(int).values
-    word_segments_arr["segment-text-end"] = per_word_grp["local-char-idx"].max().astype(int).values+1
-    word_segments_arr = pd.DataFrame(word_segments_arr)
-
-    word_segments_arr[["segment-idx", "subsegment-idx", "word-idx"]] = per_word_grp["local-char-idx"].min().reset_index()[["segment-idx", "subsegment-idx", "word-idx"]].astype(int)
-    segments_arr = {}
-    segments_arr["start"] = per_subseg_grp["start"].min().reset_index()["start"]
-    segments_arr["end"] = per_subseg_grp["end"].max().reset_index()["end"]
-    segments_arr = pd.DataFrame(segments_arr)
-    segments_arr[["segment-idx", "subsegment-idx-start"]] = per_subseg_grp["start"].min().reset_index()[["segment-idx", "subsegment-idx"]]
-    segments_arr["subsegment-idx-end"] = segments_arr["subsegment-idx-start"] + 1
-
-    # interpolate missing words / sub-segments
-    if interpolate_method != "ignore":
-        wrd_subseg_grp = word_segments_arr.groupby(["segment-idx", "subsegment-idx"], group_keys=False)
-        wrd_seg_grp = word_segments_arr.groupby(["segment-idx"], group_keys=False)
-        # we still know which word timestamps are interpolated because their score == nan
-        word_segments_arr["start"] = wrd_subseg_grp['start'].apply(lambda group: interpolate_nans(group, method=interpolate_method))
-        word_segments_arr["end"] = wrd_subseg_grp['end'].apply(lambda group: interpolate_nans(group, method=interpolate_method))
-
-        word_segments_arr["start"] = wrd_seg_grp['start'].apply(lambda group: interpolate_nans(group, method=interpolate_method))
-        word_segments_arr["end"] = wrd_seg_grp['end'].apply(lambda group: interpolate_nans(group, method=interpolate_method))
-
-        sub_seg_grp =  segments_arr.groupby(["segment-idx"], group_keys=False)
-        segments_arr['start'] = sub_seg_grp['start'].apply(lambda group: interpolate_nans(group, method=interpolate_method))
-        segments_arr['end'] = sub_seg_grp['end'].apply(lambda group: interpolate_nans(group, method=interpolate_method))
-
-        # merge words & subsegments which are missing times
-        word_grp = word_segments_arr.groupby(["segment-idx", "subsegment-idx", "end"])
-
-        word_segments_arr["segment-text-start"] = word_grp["segment-text-start"].transform(min)
-        word_segments_arr["segment-text-end"] = word_grp["segment-text-end"].transform(max)
-        word_segments_arr.drop_duplicates(subset=["segment-idx", "subsegment-idx", "end"], inplace=True)
-
-        seg_grp_dup = segments_arr.groupby(["segment-idx", "start", "end"])
-        segments_arr["subsegment-idx-start"] = seg_grp_dup["subsegment-idx-start"].transform(min)
-        segments_arr["subsegment-idx-end"] = seg_grp_dup["subsegment-idx-end"].transform(max)
-        segments_arr.drop_duplicates(subset=["segment-idx", "subsegment-idx-start", "subsegment-idx-end"], inplace=True)
-    else:
-        word_segments_arr.dropna(inplace=True)
-        segments_arr.dropna(inplace=True)
-
-    # if some segments still have missing timestamps (usually because all numerals / symbols), then use original timestamps...
-    segments_arr['start'].fillna(pd.Series([x['start'] for x in transcript]), inplace=True)
-    segments_arr['end'].fillna(pd.Series([x['end'] for x in transcript]), inplace=True)
-    segments_arr['subsegment-idx-start'].fillna(0, inplace=True)
-    segments_arr['subsegment-idx-end'].fillna(1, inplace=True)
-
-
-    aligned_segments = []
-    aligned_segments_word = []
-
-    word_segments_arr.set_index(["segment-idx", "subsegment-idx"], inplace=True)
-    char_segments_arr.set_index(["segment-idx", "subsegment-idx", "word-idx"], inplace=True)
-
-    for sdx, srow in segments_arr.iterrows():
-
-        seg_idx = int(srow["segment-idx"])
-        sub_start = int(srow["subsegment-idx-start"])
-        sub_end = int(srow["subsegment-idx-end"])
-
-        seg = transcript[seg_idx]
-        text = "".join(seg["seg-text"][sub_start:sub_end])
-
-        wseg = word_segments_arr.loc[seg_idx].loc[sub_start:sub_end-1]
-        wseg["start"].fillna(srow["start"], inplace=True)
-        wseg["end"].fillna(srow["end"], inplace=True)
-        wseg["segment-text-start"].fillna(0, inplace=True)
-        wseg["segment-text-end"].fillna(len(text)-1, inplace=True)
-
-        cseg = char_segments_arr.loc[seg_idx].loc[sub_start:sub_end-1]
-        # fixes bug for single segment in transcript
-        cseg['segment-text-start'] = cseg['level_1'] if 'level_1' in cseg else 0
-        cseg['segment-text-end'] = cseg['level_1'] + 1 if 'level_1' in cseg else 1
-        if 'level_1' in cseg: del cseg['level_1']
-        if 'level_0' in cseg: del cseg['level_0']
-        cseg.reset_index(inplace=True)
-        aligned_segments.append(
-            {
-                "start": srow["start"],
-                "end": srow["end"],
-                "text": text,
-                "word-segments": wseg,
-                "char-segments": cseg
-            }
-        )
-
-        def get_raw_text(word_row):
-            return seg["seg-text"][word_row.name][int(word_row["segment-text-start"]):int(word_row["segment-text-end"])+1]
-
-        wdx = 0
-        curr_text = get_raw_text(wseg.iloc[wdx])
-        if len(wseg) > 1:
-            for _, wrow in wseg.iloc[1:].iterrows():
-                if wrow['start'] != wseg.iloc[wdx]['start']:
-                    aligned_segments_word.append(
-                        {
-                            "text": curr_text.strip(),
-                            "start": wseg.iloc[wdx]["start"],
-                            "end": wseg.iloc[wdx]["end"],
-                        }
-                    )
-                    curr_text = ""
-                curr_text += " " + get_raw_text(wrow)
-                wdx += 1
-        aligned_segments_word.append(
-            {
-                "text": curr_text.strip(),
-                "start": wseg.iloc[wdx]["start"],
-                "end": wseg.iloc[wdx]["end"]
-            }
-        )
-
-    
-    return {"segments": aligned_segments, "word_segments": aligned_segments_word}
-
-
-"""
-source: https://pytorch.org/tutorials/intermediate/forced_alignment_with_torchaudio_tutorial.html
-"""
-def get_trellis(emission, tokens, blank_id=0):
-    num_frame = emission.size(0)
-    num_tokens = len(tokens)
-
-    # Trellis has extra diemsions for both time axis and tokens.
-    # The extra dim for tokens represents <SoS> (start-of-sentence)
-    # The extra dim for time axis is for simplification of the code.
-    trellis = torch.empty((num_frame + 1, num_tokens + 1))
-    trellis[0, 0] = 0
-    trellis[1:, 0] = torch.cumsum(emission[:, 0], 0)
-    trellis[0, -num_tokens:] = -float("inf")
-    trellis[-num_tokens:, 0] = float("inf")
-
-    for t in range(num_frame):
-        trellis[t + 1, 1:] = torch.maximum(
-            # Score for staying at the same token
-            trellis[t, 1:] + emission[t, blank_id],
-            # Score for changing to the next token
-            trellis[t, :-1] + emission[t, tokens],
-        )
-    return trellis
-
-@dataclass
-class Point:
-    token_index: int
-    time_index: int
-    score: float
-
-def backtrack(trellis, emission, tokens, blank_id=0):
-    # Note:
-    # j and t are indices for trellis, which has extra dimensions
-    # for time and tokens at the beginning.
-    # When referring to time frame index `T` in trellis,
-    # the corresponding index in emission is `T-1`.
-    # Similarly, when referring to token index `J` in trellis,
-    # the corresponding index in transcript is `J-1`.
-    j = trellis.size(1) - 1
-    t_start = torch.argmax(trellis[:, j]).item()
-
-    path = []
-    for t in range(t_start, 0, -1):
-        # 1. Figure out if the current position was stay or change
-        # Note (again):
-        # `emission[J-1]` is the emission at time frame `J` of trellis dimension.
-        # Score for token staying the same from time frame J-1 to T.
-        stayed = trellis[t - 1, j] + emission[t - 1, blank_id]
-        # Score for token changing from C-1 at T-1 to J at T.
-        changed = trellis[t - 1, j - 1] + emission[t - 1, tokens[j - 1]]
-
-        # 2. Store the path with frame-wise probability.
-        prob = emission[t - 1, tokens[j - 1] if changed > stayed else 0].exp().item()
-        # Return token index and time index in non-trellis coordinate.
-        path.append(Point(j - 1, t - 1, prob))
-
-        # 3. Update the token
-        if changed > stayed:
-            j -= 1
-            if j == 0:
-                break
-    else:
-        # failed
-        return None
-    return path[::-1]
-
-# Merge the labels
-@dataclass
-class Segment:
-    label: str
-    start: int
-    end: int
-    score: float
-
-    def __repr__(self):
-        return f"{self.label}\t({self.score:4.2f}): [{self.start:5d}, {self.end:5d})"
-
-    @property
-    def length(self):
-        return self.end - self.start
-
-def merge_repeats(path, transcript):
-    i1, i2 = 0, 0
-    segments = []
-    while i1 < len(path):
-        while i2 < len(path) and path[i1].token_index == path[i2].token_index:
-            i2 += 1
-        score = sum(path[k].score for k in range(i1, i2)) / (i2 - i1)
-        segments.append(
-            Segment(
-                transcript[path[i1].token_index],
-                path[i1].time_index,
-                path[i2 - 1].time_index + 1,
-                score,
-            )
-        )
-        i1 = i2
-    return segments
-
-def merge_words(segments, separator="|"):
-    words = []
-    i1, i2 = 0, 0
-    while i1 < len(segments):
-        if i2 >= len(segments) or segments[i2].label == separator:
-            if i1 != i2:
-                segs = segments[i1:i2]
-                word = "".join([seg.label for seg in segs])
-                score = sum(seg.score * seg.length for seg in segs) / sum(seg.length for seg in segs)
-                words.append(Segment(word, segments[i1].start, segments[i2 - 1].end, score))
-            i1 = i2 + 1
-            i2 = i1
-        else:
-            i2 += 1
-    return words

whisperx/asr.py

@@ -1,429 +0,0 @@
-import warnings
-from typing import TYPE_CHECKING, Optional, Tuple, Union
-import numpy as np
-import torch
-import tqdm
-import ffmpeg
-from whisper.audio import (
-    FRAMES_PER_SECOND,
-    HOP_LENGTH,
-    N_FRAMES,
-    N_SAMPLES,
-    SAMPLE_RATE,
-    CHUNK_LENGTH,
-    log_mel_spectrogram,
-    pad_or_trim,
-    load_audio
-)
-from whisper.decoding import DecodingOptions, DecodingResult
-from whisper.timing import add_word_timestamps
-from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE, get_tokenizer
-from whisper.utils import (
-    exact_div,
-    format_timestamp,
-    make_safe,
-)
-
-if TYPE_CHECKING:
-    from whisper.model import Whisper
-
-from .vad import merge_chunks
-
-def transcribe(
-    model: "Whisper",
-    audio: Union[str, np.ndarray, torch.Tensor] = None,
-    mel: np.ndarray = None,
-    verbose: Optional[bool] = None,
-    temperature: Union[float, Tuple[float, ...]] = (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
-    compression_ratio_threshold: Optional[float] = 2.4,
-    logprob_threshold: Optional[float] = -1.0,
-    no_speech_threshold: Optional[float] = 0.6,
-    condition_on_previous_text: bool = True,
-    initial_prompt: Optional[str] = None,
-    word_timestamps: bool = False,
-    prepend_punctuations: str = "\"'“¿([{-",
-    append_punctuations: str = "\"'.。,，!！?？:：”)]}、",
-    **decode_options,
-):
-    """
-    Transcribe an audio file using Whisper.
-    We redefine the Whisper transcribe function to allow mel input (for sequential slicing of audio)
-
-    Parameters
-    ----------
-    model: Whisper
-        The Whisper model instance
-
-    audio: Union[str, np.ndarray, torch.Tensor]
-        The path to the audio file to open, or the audio waveform
-
-    mel: np.ndarray
-        Mel spectrogram of audio segment.
-        
-    verbose: bool
-        Whether to display the text being decoded to the console. If True, displays all the details,
-        If False, displays minimal details. If None, does not display anything
-
-    temperature: Union[float, Tuple[float, ...]]
-        Temperature for sampling. It can be a tuple of temperatures, which will be successively used
-        upon failures according to either `compression_ratio_threshold` or `logprob_threshold`.
-
-    compression_ratio_threshold: float
-        If the gzip compression ratio is above this value, treat as failed
-
-    logprob_threshold: float
-        If the average log probability over sampled tokens is below this value, treat as failed
-
-    no_speech_threshold: float
-        If the no_speech probability is higher than this value AND the average log probability
-        over sampled tokens is below `logprob_threshold`, consider the segment as silent
-
-    condition_on_previous_text: bool
-        if True, the previous output of the model is provided as a prompt for the next window;
-        disabling may make the text inconsistent across windows, but the model becomes less prone to
-        getting stuck in a failure loop, such as repetition looping or timestamps going out of sync.
-
-    word_timestamps: bool
-        Extract word-level timestamps using the cross-attention pattern and dynamic time warping,
-        and include the timestamps for each word in each segment.
-
-    prepend_punctuations: str
-        If word_timestamps is True, merge these punctuation symbols with the next word
-
-    append_punctuations: str
-        If word_timestamps is True, merge these punctuation symbols with the previous word
-
-    initial_prompt: Optional[str]
-        Optional text to provide as a prompt for the first window. This can be used to provide, or
-        "prompt-engineer" a context for transcription, e.g. custom vocabularies or proper nouns
-        to make it more likely to predict those word correctly.
-
-    decode_options: dict
-        Keyword arguments to construct `DecodingOptions` instances
-
-    Returns
-    -------
-    A dictionary containing the resulting text ("text") and segment-level details ("segments"), and
-    the spoken language ("language"), which is detected when `decode_options["language"]` is None.
-    """
-    dtype = torch.float16 if decode_options.get("fp16", True) else torch.float32
-    if model.device == torch.device("cpu"):
-        if torch.cuda.is_available():
-            warnings.warn("Performing inference on CPU when CUDA is available")
-        if dtype == torch.float16:
-            warnings.warn("FP16 is not supported on CPU; using FP32 instead")
-            dtype = torch.float32
-
-    if dtype == torch.float32:
-        decode_options["fp16"] = False
-
-    # Pad 30-seconds of silence to the input audio, for slicing
-    if mel is None:
-        if audio is None:
-            raise ValueError("Transcribe needs either audio or mel as input, currently both are none.")
-        mel = log_mel_spectrogram(audio, padding=N_SAMPLES)
-    content_frames = mel.shape[-1] - N_FRAMES
-
-    if decode_options.get("language", None) is None:
-        if not model.is_multilingual:
-            decode_options["language"] = "en"
-        else:
-            if verbose:
-                print(
-                    "Detecting language using up to the first 30 seconds. Use `--language` to specify the language"
-                )
-            mel_segment = pad_or_trim(mel, N_FRAMES).to(model.device).to(dtype)
-            _, probs = model.detect_language(mel_segment)
-            decode_options["language"] = max(probs, key=probs.get)
-            if verbose is not None:
-                print(
-                    f"Detected language: {LANGUAGES[decode_options['language']].title()}"
-                )
-
-    language: str = decode_options["language"]
-    task: str = decode_options.get("task", "transcribe")
-    tokenizer = get_tokenizer(model.is_multilingual, language=language, task=task)
-
-    if word_timestamps and task == "translate":
-        warnings.warn("Word-level timestamps on translations may not be reliable.")
-
-    def decode_with_fallback(segment: torch.Tensor) -> DecodingResult:
-        temperatures = (
-            [temperature] if isinstance(temperature, (int, float)) else temperature
-        )
-        decode_result = None
-
-        for t in temperatures:
-            kwargs = {**decode_options}
-            if t > 0:
-                # disable beam_size and patience when t > 0
-                kwargs.pop("beam_size", None)
-                kwargs.pop("patience", None)
-            else:
-                # disable best_of when t == 0
-                kwargs.pop("best_of", None)
-
-            options = DecodingOptions(**kwargs, temperature=t)
-            decode_result = model.decode(segment, options)
-
-            needs_fallback = False
-            if (
-                compression_ratio_threshold is not None
-                and decode_result.compression_ratio > compression_ratio_threshold
-            ):
-                needs_fallback = True  # too repetitive
-            if (
-                logprob_threshold is not None
-                and decode_result.avg_logprob < logprob_threshold
-            ):
-                needs_fallback = True  # average log probability is too low
-
-            if not needs_fallback:
-                break
-
-        return decode_result
-
-    seek = 0
-    input_stride = exact_div(
-        N_FRAMES, model.dims.n_audio_ctx
-    )  # mel frames per output token: 2
-    time_precision = (
-        input_stride * HOP_LENGTH / SAMPLE_RATE
-    )  # time per output token: 0.02 (seconds)
-    all_tokens = []
-    all_segments = []
-    prompt_reset_since = 0
-
-    if initial_prompt is not None:
-        initial_prompt_tokens = tokenizer.encode(" " + initial_prompt.strip())
-        all_tokens.extend(initial_prompt_tokens)
-    else:
-        initial_prompt_tokens = []
-
-    def new_segment(
-        *, start: float, end: float, tokens: torch.Tensor, result: DecodingResult
-    ):
-        tokens = tokens.tolist()
-        text_tokens = [token for token in tokens if token < tokenizer.eot]
-        return {
-            "seek": seek,
-            "start": start,
-            "end": end,
-            "text": tokenizer.decode(text_tokens),
-            "tokens": tokens,
-            "temperature": result.temperature,
-            "avg_logprob": result.avg_logprob,
-            "compression_ratio": result.compression_ratio,
-            "no_speech_prob": result.no_speech_prob,
-        }
-
-
-    # show the progress bar when verbose is False (if True, transcribed text will be printed)
-    with tqdm.tqdm(
-        total=content_frames, unit="frames", disable=verbose is not False
-    ) as pbar:
-        while seek < content_frames:
-            time_offset = float(seek * HOP_LENGTH / SAMPLE_RATE)
-            mel_segment = mel[:, seek : seek + N_FRAMES]
-            segment_size = min(N_FRAMES, content_frames - seek)
-            segment_duration = segment_size * HOP_LENGTH / SAMPLE_RATE
-            mel_segment = pad_or_trim(mel_segment, N_FRAMES).to(model.device).to(dtype)
-
-            decode_options["prompt"] = all_tokens[prompt_reset_since:]
-            result: DecodingResult = decode_with_fallback(mel_segment)
-            tokens = torch.tensor(result.tokens)
-            if no_speech_threshold is not None:
-                # no voice activity check
-                should_skip = result.no_speech_prob > no_speech_threshold
-                if (
-                    logprob_threshold is not None
-                    and result.avg_logprob > logprob_threshold
-                ):
-                    # don't skip if the logprob is high enough, despite the no_speech_prob
-                    should_skip = False
-
-                if should_skip:
-                    seek += segment_size  # fast-forward to the next segment boundary
-                    continue
-
-            previous_seek = seek
-            current_segments = []
-
-            timestamp_tokens: torch.Tensor = tokens.ge(tokenizer.timestamp_begin)
-            single_timestamp_ending = timestamp_tokens[-2:].tolist() == [False, True]
-
-            consecutive = torch.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0]
-            consecutive.add_(1)
-            if len(consecutive) > 0:
-                # if the output contains two consecutive timestamp tokens
-                slices = consecutive.tolist()
-                if single_timestamp_ending:
-                    slices.append(len(tokens))
-
-                last_slice = 0
-                for current_slice in slices:
-                    sliced_tokens = tokens[last_slice:current_slice]
-                    start_timestamp_pos = (
-                        sliced_tokens[0].item() - tokenizer.timestamp_begin
-                    )
-                    end_timestamp_pos = (
-                        sliced_tokens[-1].item() - tokenizer.timestamp_begin
-                    )
-                    current_segments.append(
-                        new_segment(
-                            start=time_offset + start_timestamp_pos * time_precision,
-                            end=time_offset + end_timestamp_pos * time_precision,
-                            tokens=sliced_tokens,
-                            result=result,
-                        )
-                    )
-                    last_slice = current_slice
-
-                if single_timestamp_ending:
-                    # single timestamp at the end means no speech after the last timestamp.
-                    seek += segment_size
-                else:
-                    # otherwise, ignore the unfinished segment and seek to the last timestamp
-                    last_timestamp_pos = (
-                        tokens[last_slice - 1].item() - tokenizer.timestamp_begin
-                    )
-                    seek += last_timestamp_pos * input_stride
-            else:
-                duration = segment_duration
-                timestamps = tokens[timestamp_tokens.nonzero().flatten()]
-                if (
-                    len(timestamps) > 0
-                    and timestamps[-1].item() != tokenizer.timestamp_begin
-                ):
-                    # no consecutive timestamps but it has a timestamp; use the last one.
-                    last_timestamp_pos = (
-                        timestamps[-1].item() - tokenizer.timestamp_begin
-                    )
-                    duration = last_timestamp_pos * time_precision
-
-                current_segments.append(
-                    new_segment(
-                        start=time_offset,
-                        end=time_offset + duration,
-                        tokens=tokens,
-                        result=result,
-                    )
-                )
-                seek += segment_size
-
-            if not condition_on_previous_text or result.temperature > 0.5:
-                # do not feed the prompt tokens if a high temperature was used
-                prompt_reset_since = len(all_tokens)
-
-            if word_timestamps:
-                add_word_timestamps(
-                    segments=current_segments,
-                    model=model,
-                    tokenizer=tokenizer,
-                    mel=mel_segment,
-                    num_frames=segment_size,
-                    prepend_punctuations=prepend_punctuations,
-                    append_punctuations=append_punctuations,
-                )
-                word_end_timestamps = [
-                    w["end"] for s in current_segments for w in s["words"]
-                ]
-                if not single_timestamp_ending and len(word_end_timestamps) > 0:
-                    seek_shift = round(
-                        (word_end_timestamps[-1] - time_offset) * FRAMES_PER_SECOND
-                    )
-                    if seek_shift > 0:
-                        seek = previous_seek + seek_shift
-
-            if verbose:
-                for segment in current_segments:
-                    start, end, text = segment["start"], segment["end"], segment["text"]
-                    line = f"[{format_timestamp(start)} --> {format_timestamp(end)}] {text}"
-                    print(make_safe(line))
-
-            # if a segment is instantaneous or does not contain text, clear it
-            for i, segment in enumerate(current_segments):
-                if segment["start"] == segment["end"] or segment["text"].strip() == "":
-                    segment["text"] = ""
-                    segment["tokens"] = []
-                    segment["words"] = []
-
-            all_segments.extend(
-                [
-                    {"id": i, **segment}
-                    for i, segment in enumerate(
-                        current_segments, start=len(all_segments)
-                    )
-                ]
-            )
-            all_tokens.extend(
-                [token for segment in current_segments for token in segment["tokens"]]
-            )
-
-            # update progress bar
-            pbar.update(min(content_frames, seek) - previous_seek)
-
-
-    return dict(
-        text=tokenizer.decode(all_tokens[len(initial_prompt_tokens) :]),
-        segments=all_segments,
-        language=language,
-    )
-
-
-def transcribe_with_vad(
-    model: "Whisper",
-    audio: str,
-    vad_pipeline,
-    mel = None,
-    verbose: Optional[bool] = None,
-    **kwargs
-):
-    """
-    Transcribe per VAD segment
-    """
-
-    vad_segments = vad_pipeline(audio)
-
-    # if not torch.is_tensor(audio):
-        # if isinstance(audio, str):
-    audio = load_audio(audio)
-    audio = torch.from_numpy(audio)
-
-    prev = 0
-    output = {"segments": []}
-
-    # merge segments to approx 30s inputs to make whisper most appropraite
-    vad_segments = merge_chunks(vad_segments, chunk_size=CHUNK_LENGTH)
-    if len(vad_segments) == 0:
-        return output
-
-    print(">>Performing transcription...")
-    for sdx, seg_t in enumerate(vad_segments):
-        if verbose:
-            print(f"~~ Transcribing VAD chunk: ({format_timestamp(seg_t['start'])} --> {format_timestamp(seg_t['end'])}) ~~")
-        seg_f_start, seg_f_end = int(seg_t["start"] * SAMPLE_RATE), int(seg_t["end"] * SAMPLE_RATE)
-        local_f_start, local_f_end = seg_f_start - prev, seg_f_end - prev
-        audio = audio[local_f_start:] # seek forward
-        seg_audio = audio[:local_f_end-local_f_start] # seek forward
-        prev = seg_f_start
-        local_mel = log_mel_spectrogram(seg_audio, padding=N_SAMPLES)
-        # need to pad
-
-        result = transcribe(model, audio, mel=local_mel, verbose=verbose, **kwargs)
-        seg_t["text"] = result["text"]
-        output["segments"].append(
-            {
-                "start": seg_t["start"],
-                "end": seg_t["end"],
-                "language": result["language"],
-                "text": result["text"],
-                "seg-text": [x["text"] for x in result["segments"]],
-                "seg-start": [x["start"] for x in result["segments"]],
-                "seg-end": [x["end"] for x in result["segments"]],
-                }
-            )
-
-    output["language"] = output["segments"][0]["language"]
-
-    return output

whisperx/diarize.py

@@ -1,76 +0,0 @@
-import numpy as np
-import pandas as pd
-from pyannote.audio import Pipeline
-
-class DiarizationPipeline:
-    def __init__(
-        self,
-        model_name="pyannote/[email protected]",
-        use_auth_token=None,
-    ):
-        self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token)
-
-    def __call__(self, audio, min_speakers=None, max_speakers=None):
-        segments = self.model(audio, min_speakers=min_speakers, max_speakers=max_speakers)
-        diarize_df = pd.DataFrame(segments.itertracks(yield_label=True))
-        diarize_df['start'] = diarize_df[0].apply(lambda x: x.start)
-        diarize_df['end'] = diarize_df[0].apply(lambda x: x.end)
-        return diarize_df
-
-def assign_word_speakers(diarize_df, result_segments, fill_nearest=False):
-    for seg in result_segments:
-        wdf = seg['word-segments']
-        if len(wdf['start'].dropna()) == 0:
-            wdf['start'] = seg['start']
-            wdf['end'] = seg['end']
-        speakers = []
-        for wdx, wrow in wdf.iterrows():
-            if not np.isnan(wrow['start']):
-                diarize_df['intersection'] = np.minimum(diarize_df['end'], wrow['end']) - np.maximum(diarize_df['start'], wrow['start'])
-                diarize_df['union'] = np.maximum(diarize_df['end'], wrow['end']) - np.minimum(diarize_df['start'], wrow['start'])
-                # remove no hit
-                if not fill_nearest:
-                    dia_tmp = diarize_df[diarize_df['intersection'] > 0]
-                else:
-                    dia_tmp = diarize_df
-                if len(dia_tmp) == 0:
-                    speaker = None
-                else:
-                    speaker = dia_tmp.sort_values("intersection", ascending=False).iloc[0][2]
-            else:
-                speaker = None
-            speakers.append(speaker)
-        seg['word-segments']['speaker'] = speakers
-
-        speaker_count = pd.Series(speakers).value_counts()
-        if len(speaker_count) == 0:
-            seg["speaker"]= "UNKNOWN"
-        else:
-            seg["speaker"] = speaker_count.index[0]
-
-    # create word level segments for .srt
-    word_seg = []
-    for seg in result_segments:
-        wseg = pd.DataFrame(seg["word-segments"])
-        for wdx, wrow in wseg.iterrows():
-            if wrow["start"] is not None:
-                speaker = wrow['speaker']
-                if speaker is None or speaker == np.nan:
-                    speaker = "UNKNOWN"
-                word_seg.append(
-                    {
-                        "start": wrow["start"],
-                        "end": wrow["end"],
-                        "text": f"[{speaker}]: " + seg["text"][int(wrow["segment-text-start"]):int(wrow["segment-text-end"])]
-                    }
-                )
-
-    # TODO: create segments but split words on new speaker
-
-    return result_segments, word_seg
-
-class Segment:
-    def __init__(self, start, end, speaker=None):
-        self.start = start
-        self.end = end
-        self.speaker = speaker

whisperx/transcribe.py

@@ -1,220 +0,0 @@
-import argparse
-import os
-import gc
-import warnings
-from typing import TYPE_CHECKING, Optional, Tuple, Union
-import numpy as np
-import torch
-import tempfile 
-import ffmpeg
-from whisper.tokenizer import LANGUAGES, TO_LANGUAGE_CODE
-from whisper.audio import SAMPLE_RATE
-from whisper.utils import (
-    optional_float,
-    optional_int,
-    str2bool,
-)
-
-from .alignment import load_align_model, align
-from .asr import transcribe, transcribe_with_vad
-from .diarize import DiarizationPipeline, assign_word_speakers
-from .utils import get_writer
-from .vad import load_vad_model
-
-def cli():
-    from whisper import available_models
-
-    # fmt: off
-    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    parser.add_argument("audio", nargs="+", type=str, help="audio file(s) to transcribe")
-    parser.add_argument("--model", default="small", choices=available_models(), help="name of the Whisper model to use")
-    parser.add_argument("--model_dir", type=str, default=None, help="the path to save model files; uses ~/.cache/whisper by default")
-    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu", help="device to use for PyTorch inference")
-    parser.add_argument("--output_dir", "-o", type=str, default=".", help="directory to save the outputs")
-    parser.add_argument("--output_format", "-f", type=str, default="all", choices=["all", "srt", "srt-word", "vtt", "txt", "tsv", "ass", "ass-char", "pickle", "vad"], help="format of the output file; if not specified, all available formats will be produced")
-    parser.add_argument("--verbose", type=str2bool, default=True, help="whether to print out the progress and debug messages")
-
-    parser.add_argument("--task", type=str, default="transcribe", choices=["transcribe", "translate"], help="whether to perform X->X speech recognition ('transcribe') or X->English translation ('translate')")
-    parser.add_argument("--language", type=str, default=None, choices=sorted(LANGUAGES.keys()) + sorted([k.title() for k in TO_LANGUAGE_CODE.keys()]), help="language spoken in the audio, specify None to perform language detection")
-
-    # alignment params
-    parser.add_argument("--align_model", default=None, help="Name of phoneme-level ASR model to do alignment")
-    parser.add_argument("--align_extend", default=2, type=float, help="Seconds before and after to extend the whisper segments for alignment (if not using VAD).")
-    parser.add_argument("--align_from_prev", default=True, type=bool, help="Whether to clip the alignment start time of current segment to the end time of the last aligned word of the previous segment (if not using VAD)")
-    parser.add_argument("--interpolate_method", default="nearest", choices=["nearest", "linear", "ignore"], help="For word .srt, method to assign timestamps to non-aligned words, or merge them into neighbouring.")
-    parser.add_argument("--no_align", action='store_true', help="Do not perform phoneme alignment")
-
-    # vad params
-    parser.add_argument("--vad_filter", type=str2bool, default=True, help="Whether to pre-segment audio with VAD, highly recommended! Produces more accurate alignment + timestamp see WhisperX paper https://arxiv.org/abs/2303.00747")
-    parser.add_argument("--vad_onset", type=float, default=0.500, help="Onset threshold for VAD (see pyannote.audio), reduce this if speech is not being detected")
-    parser.add_argument("--vad_offset", type=float, default=0.363, help="Offset threshold for VAD (see pyannote.audio), reduce this if speech is not being detected.")
-
-    # diarization params
-    parser.add_argument("--diarize", action="store_true", help="Apply diarization to assign speaker labels to each segment/word")
-    parser.add_argument("--min_speakers", default=None, type=int)
-    parser.add_argument("--max_speakers", default=None, type=int)
-
-    parser.add_argument("--temperature", type=float, default=0, help="temperature to use for sampling")
-    parser.add_argument("--best_of", type=optional_int, default=5, help="number of candidates when sampling with non-zero temperature")
-    parser.add_argument("--beam_size", type=optional_int, default=5, help="number of beams in beam search, only applicable when temperature is zero")
-    parser.add_argument("--patience", type=float, default=None, help="optional patience value to use in beam decoding, as in https://arxiv.org/abs/2204.05424, the default (1.0) is equivalent to conventional beam search")
-    parser.add_argument("--length_penalty", type=float, default=None, help="optional token length penalty coefficient (alpha) as in https://arxiv.org/abs/1609.08144, uses simple length normalization by default")
-
-    parser.add_argument("--suppress_tokens", type=str, default="-1", help="comma-separated list of token ids to suppress during sampling; '-1' will suppress most special characters except common punctuations")
-    parser.add_argument("--initial_prompt", type=str, default=None, help="optional text to provide as a prompt for the first window.")
-    parser.add_argument("--condition_on_previous_text", type=str2bool, default=False, help="if True, provide the previous output of the model as a prompt for the next window; disabling may make the text inconsistent across windows, but the model becomes less prone to getting stuck in a failure loop")
-    parser.add_argument("--fp16", type=str2bool, default=True, help="whether to perform inference in fp16; True by default")
-
-    parser.add_argument("--temperature_increment_on_fallback", type=optional_float, default=0.2, help="temperature to increase when falling back when the decoding fails to meet either of the thresholds below")
-    parser.add_argument("--compression_ratio_threshold", type=optional_float, default=2.4, help="if the gzip compression ratio is higher than this value, treat the decoding as failed")
-    parser.add_argument("--logprob_threshold", type=optional_float, default=-1.0, help="if the average log probability is lower than this value, treat the decoding as failed")
-    parser.add_argument("--no_speech_threshold", type=optional_float, default=0.6, help="if the probability of the <|nospeech|> token is higher than this value AND the decoding has failed due to `logprob_threshold`, consider the segment as silence")
-    parser.add_argument("--word_timestamps", type=str2bool, default=False, help="(experimental) extract word-level timestamps and refine the results based on them")
-    parser.add_argument("--prepend_punctuations", type=str, default="\"\'“¿([{-", help="if word_timestamps is True, merge these punctuation symbols with the next word")
-    parser.add_argument("--append_punctuations", type=str, default="\"\'.。,，!！?？:：”)]}、", help="if word_timestamps is True, merge these punctuation symbols with the previous word")
-    parser.add_argument("--threads", type=optional_int, default=0, help="number of threads used by torch for CPU inference; supercedes MKL_NUM_THREADS/OMP_NUM_THREADS")
-
-    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face Access Token to access PyAnnote gated models")
-    # parser.add_argument("--model_flush", action="store_true", help="Flush memory from each model after use, reduces GPU requirement but slower processing >1 audio file.")
-    parser.add_argument("--tmp_dir", default=None, help="Temporary directory to write audio file if input if not .wav format (only for VAD).")
-    # fmt: on
-
-    args = parser.parse_args().__dict__
-    model_name: str = args.pop("model")
-    model_dir: str = args.pop("model_dir")
-    output_dir: str = args.pop("output_dir")
-    output_format: str = args.pop("output_format")
-    device: str = args.pop("device")
-    # model_flush: bool = args.pop("model_flush")
-    os.makedirs(output_dir, exist_ok=True)
-
-    tmp_dir: str = args.pop("tmp_dir")
-    if tmp_dir is not None:
-        os.makedirs(tmp_dir, exist_ok=True)
-
-    align_model: str = args.pop("align_model")
-    align_extend: float = args.pop("align_extend")
-    align_from_prev: bool = args.pop("align_from_prev")
-    interpolate_method: str = args.pop("interpolate_method")
-    no_align: bool = args.pop("no_align")
-
-    hf_token: str = args.pop("hf_token")
-    vad_filter: bool = args.pop("vad_filter")
-    vad_onset: float = args.pop("vad_onset")
-    vad_offset: float = args.pop("vad_offset")
-
-    diarize: bool = args.pop("diarize")
-    min_speakers: int = args.pop("min_speakers")
-    max_speakers: int = args.pop("max_speakers")
-
-    if vad_filter:
-        from pyannote.audio import Pipeline
-        from pyannote.audio import Model, Pipeline
-        vad_model = load_vad_model(torch.device(device), vad_onset, vad_offset, use_auth_token=hf_token)
-    else:
-        vad_model = None
-
-    # if model_flush:
-    #     print(">>Model flushing activated... Only loading model after ASR stage")
-    #     del align_model
-    #     align_model = ""
-
-
-    if model_name.endswith(".en") and args["language"] not in {"en", "English"}:
-        if args["language"] is not None:
-            warnings.warn(
-                f"{model_name} is an English-only model but receipted '{args['language']}'; using English instead."
-            )
-        args["language"] = "en"
-
-    temperature = args.pop("temperature")
-    if (increment := args.pop("temperature_increment_on_fallback")) is not None:
-        temperature = tuple(np.arange(temperature, 1.0 + 1e-6, increment))
-    else:
-        temperature = [temperature]
-
-    if (threads := args.pop("threads")) > 0:
-        torch.set_num_threads(threads)
-
-    from whisper import load_model
-
-    writer = get_writer(output_format, output_dir)
-
-    # Part 1: VAD & ASR Loop
-    results = []
-    tmp_results = []
-    model = load_model(model_name, device=device, download_root=model_dir)
-    for audio_path in args.pop("audio"):
-        input_audio_path = audio_path
-        tfile = None
-
-        # >> VAD & ASR
-        if vad_model is not None:
-            if not audio_path.endswith(".wav"):
-                print(">>VAD requires .wav format, converting to wav as a tempfile...")
-                audio_basename = os.path.splitext(os.path.basename(audio_path))[0]
-                if tmp_dir is not None:
-                    input_audio_path = os.path.join(tmp_dir, audio_basename + ".wav")
-                else:
-                    input_audio_path = os.path.join(os.path.dirname(audio_path), audio_basename + ".wav")
-                ffmpeg.input(audio_path, threads=0).output(input_audio_path, ac=1, ar=SAMPLE_RATE).run(cmd=["ffmpeg"])
-            print(">>Performing VAD...")
-            result = transcribe_with_vad(model, input_audio_path, vad_model, temperature=temperature, **args)
-        else:
-            print(">>Performing transcription...")
-            result = transcribe(model, input_audio_path, temperature=temperature, **args)
-        
-        results.append((result, input_audio_path))
-
-    # Unload Whisper and VAD
-    del model
-    del vad_model
-    gc.collect()
-    torch.cuda.empty_cache()
-
-    # Part 2: Align Loop
-    if not no_align:
-        tmp_results = results
-        results = []
-        align_language = args["language"] if args["language"] is not None else "en" # default to loading english if not specified
-        align_model, align_metadata = load_align_model(align_language, device, model_name=align_model)
-        for result, input_audio_path in tmp_results:
-            # >> Align
-            if align_model is not None and len(result["segments"]) > 0:
-                if result.get("language", "en") != align_metadata["language"]:
-                    # load new language
-                    print(f"New language found ({result['language']})! Previous was ({align_metadata['language']}), loading new alignment model for new language...")
-                    align_model, align_metadata = load_align_model(result["language"], device)
-                print(">>Performing alignment...")
-                result = align(result["segments"], align_model, align_metadata, input_audio_path, device,
-                    extend_duration=align_extend, start_from_previous=align_from_prev, interpolate_method=interpolate_method)
-            results.append((result, input_audio_path))
-
-        # Unload align model
-        del align_model
-        gc.collect()
-        torch.cuda.empty_cache()
-
-    # >> Diarize
-    if diarize:
-        if hf_token is None:
-            print("Warning, no --hf_token used, needs to be saved in environment variable, otherwise will throw error loading diarization model...")
-        tmp_results = results
-        results = []
-        diarize_model = DiarizationPipeline(use_auth_token=hf_token)
-        for result, input_audio_path in tmp_results:
-            diarize_segments = diarize_model(input_audio_path, min_speakers=min_speakers, max_speakers=max_speakers)
-            results_segments, word_segments = assign_word_speakers(diarize_segments, result["segments"])
-            result = {"segments": results_segments, "word_segments": word_segments}
-            results.append((result, input_audio_path))
-
-    # >> Write
-    for result, audio_path in results:
-        writer(result, audio_path)
-
-        # cleanup
-        if input_audio_path != audio_path:
-            os.remove(input_audio_path)
-
-if __name__ == "__main__":
-    cli()

whisperx/utils.py

@@ -1,317 +0,0 @@
-import os
-import zlib
-from typing import Callable, TextIO, Iterator, Tuple
-import pandas as pd
-import numpy as np
-
-def interpolate_nans(x, method='nearest'):
-    if x.notnull().sum() > 1:
-        return x.interpolate(method=method).ffill().bfill()
-    else:
-        return x.ffill().bfill()
-    
-
-def write_txt(transcript: Iterator[dict], file: TextIO):
-    for segment in transcript:
-        print(segment['text'].strip(), file=file, flush=True)
-
-
-def write_vtt(transcript: Iterator[dict], file: TextIO):
-    print("WEBVTT\n", file=file)
-    for segment in transcript:
-        print(
-            f"{format_timestamp(segment['start'])} --> {format_timestamp(segment['end'])}\n"
-            f"{segment['text'].strip().replace('-->', '->')}\n",
-            file=file,
-            flush=True,
-        )
-
-def write_tsv(transcript: Iterator[dict], file: TextIO):
-    print("start", "end", "text", sep="\t", file=file)
-    for segment in transcript:
-        print(segment['start'], file=file, end="\t")
-        print(segment['end'], file=file, end="\t")
-        print(segment['text'].strip().replace("\t", " "), file=file, flush=True)
-
-
-def write_srt(transcript: Iterator[dict], file: TextIO):
-    """
-    Write a transcript to a file in SRT format.
-
-    Example usage:
-        from pathlib import Path
-        from whisper.utils import write_srt
-
-        result = transcribe(model, audio_path, temperature=temperature, **args)
-
-        # save SRT
-        audio_basename = Path(audio_path).stem
-        with open(Path(output_dir) / (audio_basename + ".srt"), "w", encoding="utf-8") as srt:
-            write_srt(result["segments"], file=srt)
-    """
-    for i, segment in enumerate(transcript, start=1):
-        # write srt lines
-        print(
-            f"{i}\n"
-            f"{format_timestamp(segment['start'], always_include_hours=True, decimal_marker=',')} --> "
-            f"{format_timestamp(segment['end'], always_include_hours=True, decimal_marker=',')}\n"
-            f"{segment['text'].strip().replace('-->', '->')}\n",
-            file=file,
-            flush=True,
-        )
-
-
-def write_ass(transcript: Iterator[dict],
-            file: TextIO,
-            resolution: str = "word",
-            color: str = None, underline=True,
-            prefmt: str = None, suffmt: str = None,
-            font: str = None, font_size: int = 24,
-            strip=True, **kwargs):
-    """
-    Credit: https://github.com/jianfch/stable-ts/blob/ff79549bd01f764427879f07ecd626c46a9a430a/stable_whisper/text_output.py
-        Generate Advanced SubStation Alpha (ass) file from results to
-    display both phrase-level & word-level timestamp simultaneously by:
-     -using segment-level timestamps display phrases as usual
-     -using word-level timestamps change formats (e.g. color/underline) of the word in the displayed segment
-    Note: ass file is used in the same way as srt, vtt, etc.
-    Parameters
-    ----------
-    transcript: dict
-        results from modified model
-    file: TextIO
-        file object to write to
-    resolution: str
-        "word" or "char", timestamp resolution to highlight.
-    color: str
-        color code for a word at its corresponding timestamp
-        <bbggrr> reverse order hexadecimal RGB value (e.g. FF0000 is full intensity blue. Default: 00FF00)
-    underline: bool
-        whether to underline a word at its corresponding timestamp
-    prefmt: str
-        used to specify format for word-level timestamps (must be use with 'suffmt' and overrides 'color'&'underline')
-        appears as such in the .ass file:
-            Hi, {<prefmt>}how{<suffmt>} are you?
-        reference [Appendix A: Style override codes] in http://www.tcax.org/docs/ass-specs.htm
-    suffmt: str
-        used to specify format for word-level timestamps (must be use with 'prefmt' and overrides 'color'&'underline')
-        appears as such in the .ass file:
-            Hi, {<prefmt>}how{<suffmt>} are you?
-        reference [Appendix A: Style override codes] in http://www.tcax.org/docs/ass-specs.htm
-    font: str
-        word font (default: Arial)
-    font_size: int
-        word font size (default: 48)
-    kwargs:
-        used for format styles:
-        'Name', 'Fontname', 'Fontsize', 'PrimaryColour', 'SecondaryColour', 'OutlineColour', 'BackColour', 'Bold',
-        'Italic', 'Underline', 'StrikeOut', 'ScaleX', 'ScaleY', 'Spacing', 'Angle', 'BorderStyle', 'Outline',
-        'Shadow', 'Alignment', 'MarginL', 'MarginR', 'MarginV', 'Encoding'
-
-    """
-
-    fmt_style_dict = {'Name': 'Default', 'Fontname': 'Arial', 'Fontsize': '48', 'PrimaryColour': '&Hffffff',
-                    'SecondaryColour': '&Hffffff', 'OutlineColour': '&H0', 'BackColour': '&H0', 'Bold': '0',
-                    'Italic': '0', 'Underline': '0', 'StrikeOut': '0', 'ScaleX': '100', 'ScaleY': '100',
-                    'Spacing': '0', 'Angle': '0', 'BorderStyle': '1', 'Outline': '1', 'Shadow': '0',
-                    'Alignment': '2', 'MarginL': '10', 'MarginR': '10', 'MarginV': '10', 'Encoding': '0'}
-
-    for k, v in filter(lambda x: 'colour' in x[0].lower() and not str(x[1]).startswith('&H'), kwargs.items()):
-        kwargs[k] = f'&H{kwargs[k]}'
-
-    fmt_style_dict.update((k, v) for k, v in kwargs.items() if k in fmt_style_dict)
-
-    if font:
-        fmt_style_dict.update(Fontname=font)
-    if font_size:
-        fmt_style_dict.update(Fontsize=font_size)
-
-    fmts = f'Format: {", ".join(map(str, fmt_style_dict.keys()))}'
-
-    styles = f'Style: {",".join(map(str, fmt_style_dict.values()))}'
-
-    ass_str = f'[Script Info]\nScriptType: v4.00+\nPlayResX: 384\nPlayResY: 288\nScaledBorderAndShadow: yes\n\n' \
-            f'[V4+ Styles]\n{fmts}\n{styles}\n\n' \
-            f'[Events]\nFormat: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n\n'
-
-    if prefmt or suffmt:
-        if suffmt:
-            assert prefmt, 'prefmt must be used along with suffmt'
-        else:
-            suffmt = r'\r'
-    else:
-        if not color:
-            color = 'HFF00'
-        underline_code = r'\u1' if underline else ''
-
-        prefmt = r'{\1c&' + f'{color.upper()}&{underline_code}' + '}'
-        suffmt = r'{\r}'
-    
-    def secs_to_hhmmss(secs: Tuple[float, int]):
-        mm, ss = divmod(secs, 60)
-        hh, mm = divmod(mm, 60)
-        return f'{hh:0>1.0f}:{mm:0>2.0f}:{ss:0>2.2f}'
-
-
-    def dialogue(chars: str, start: float, end: float, idx_0: int, idx_1: int) -> str:
-        if idx_0 == -1:
-            text = chars
-        else:
-            text = f'{chars[:idx_0]}{prefmt}{chars[idx_0:idx_1]}{suffmt}{chars[idx_1:]}'
-        return f"Dialogue: 0,{secs_to_hhmmss(start)},{secs_to_hhmmss(end)}," \
-               f"Default,,0,0,0,,{text.strip() if strip else text}"
-
-    if resolution == "word":
-        resolution_key = "word-segments"
-    elif resolution == "char":
-        resolution_key = "char-segments"
-    else:
-        raise ValueError(".ass resolution should be 'word' or 'char', not ", resolution)
-    
-    ass_arr = []
-
-    for segment in transcript:
-        # if "12" in segment['text']:
-            # import pdb; pdb.set_trace()
-        if resolution_key in segment:
-            res_segs = pd.DataFrame(segment[resolution_key])
-            prev = segment['start']
-            if "speaker" in segment:
-                speaker_str = f"[{segment['speaker']}]: "
-            else:
-                speaker_str = ""
-            for cdx, crow in res_segs.iterrows():
-                if not np.isnan(crow['start']):
-                    if resolution == "char":
-                        idx_0 = cdx
-                        idx_1 = cdx + 1
-                    elif resolution == "word":
-                        idx_0 = int(crow["segment-text-start"])
-                        idx_1 = int(crow["segment-text-end"])
-                    # fill gap
-                    if crow['start'] > prev:
-                        filler_ts = {
-                            "chars": speaker_str + segment['text'],
-                            "start": prev,
-                            "end": crow['start'],
-                            "idx_0": -1,
-                            "idx_1": -1
-                        }
-
-                        ass_arr.append(filler_ts)
-                    # highlight current word
-                    f_word_ts = {
-                        "chars": speaker_str + segment['text'],
-                        "start": crow['start'],
-                        "end": crow['end'],
-                        "idx_0": idx_0 + len(speaker_str),
-                        "idx_1": idx_1 + len(speaker_str)
-                    }
-                    ass_arr.append(f_word_ts)
-                    prev = crow['end']
-
-    ass_str += '\n'.join(map(lambda x: dialogue(**x), ass_arr))
-
-    file.write(ass_str)
-
-
-from whisper.utils import SubtitlesWriter, ResultWriter, WriteTXT, WriteVTT, WriteSRT, WriteTSV, WriteJSON, format_timestamp
-
-class WriteASS(ResultWriter):
-    extension: str = "ass"
-
-    def write_result(self, result: dict, file: TextIO):
-        write_ass(result["segments"], file, resolution="word")
-
-class WriteASSchar(ResultWriter):
-    extension: str = "ass"
-
-    def write_result(self, result: dict, file: TextIO):
-        write_ass(result["segments"], file, resolution="char")
-
-class WritePickle(ResultWriter):
-    extension: str = "ass"
-
-    def write_result(self, result: dict, file: TextIO):
-        pd.DataFrame(result["segments"]).to_pickle(file)
-
-class WriteSRTWord(ResultWriter):
-    extension: str = "word.srt"
-    always_include_hours: bool = True
-    decimal_marker: str = ","
-
-    def iterate_result(self, result: dict):
-        for segment in result["word_segments"]:
-            segment_start = self.format_timestamp(segment["start"])
-            segment_end = self.format_timestamp(segment["end"])
-            segment_text = segment["text"].strip().replace("-->", "->")
-
-            if word_timings := segment.get("words", None):
-                all_words = [timing["word"] for timing in word_timings]
-                all_words[0] = all_words[0].strip()  # remove the leading space, if any
-                last = segment_start
-                for i, this_word in enumerate(word_timings):
-                    start = self.format_timestamp(this_word["start"])
-                    end = self.format_timestamp(this_word["end"])
-                    if last != start:
-                        yield last, start, segment_text
-
-                    yield start, end, "".join(
-                        [
-                            f"<u>{word}</u>" if j == i else word
-                            for j, word in enumerate(all_words)
-                        ]
-                    )
-                    last = end
-
-                if last != segment_end:
-                    yield last, segment_end, segment_text
-            else:
-                yield segment_start, segment_end, segment_text
-
-    def write_result(self, result: dict, file: TextIO):
-        if "word_segments" not in result:
-            return
-        for i, (start, end, text) in enumerate(self.iterate_result(result), start=1):
-            print(f"{i}\n{start} --> {end}\n{text}\n", file=file, flush=True)
-
-    def format_timestamp(self, seconds: float):
-        return format_timestamp(
-            seconds=seconds,
-            always_include_hours=self.always_include_hours,
-            decimal_marker=self.decimal_marker,
-        )
-
-def get_writer(output_format: str, output_dir: str) -> Callable[[dict, TextIO], None]:
-    writers = {
-        "txt": WriteTXT,
-        "vtt": WriteVTT,
-        "srt": WriteSRT,
-        "tsv": WriteTSV,
-        "ass": WriteASS,
-        "srt-word": WriteSRTWord,
-        # "ass-char": WriteASSchar,
-        # "pickle": WritePickle,
-        # "json": WriteJSON,
-    }
-
-    writers_other = {
-        "pkl": WritePickle,
-        "ass-char": WriteASSchar
-    }
-
-    if output_format == "all":
-        all_writers = [writer(output_dir) for writer in writers.values()]
-
-        def write_all(result: dict, file: TextIO):
-            for writer in all_writers:
-                writer(result, file)
-
-        return write_all
-
-    if output_format in writers:
-        return writers[output_format](output_dir)
-    elif output_format in writers_other:
-        return writers_other[output_format](output_dir)
-    else:
-        raise ValueError(f"Output format '{output_format}' not supported, choose from {writers.keys()} and {writers_other.keys()}")

whisperx/vad.py

@@ -1,305 +0,0 @@
-import os
-import urllib
-import pandas as pd
-import numpy as np
-import torch
-import hashlib
-from tqdm import tqdm
-from typing import Optional, Callable, Union, Text
-from pyannote.audio.core.io import AudioFile
-from pyannote.core import Annotation, Segment, SlidingWindowFeature
-from pyannote.audio.pipelines.utils import PipelineModel
-from pyannote.audio import Model
-from pyannote.audio.pipelines import VoiceActivityDetection
-from .diarize import Segment as SegmentX
-from typing import List, Tuple, Optional
-
-VAD_SEGMENTATION_URL = "https://whisperx.s3.eu-west-2.amazonaws.com/model_weights/segmentation/0b5b3216d60a2d32fc086b47ea8c67589aaeb26b7e07fcbe620d6d0b83e209ea/pytorch_model.bin"
-
-def load_vad_model(device, vad_onset, vad_offset, use_auth_token=None):
-    model_dir = torch.hub._get_torch_home()
-    os.makedirs(model_dir, exist_ok = True)
-    model_fp = os.path.join(model_dir, "whisperx-vad-segmentation.bin")
-    if os.path.exists(model_fp) and not os.path.isfile(model_fp):
-        raise RuntimeError(f"{model_fp} exists and is not a regular file")
-
-    if not os.path.isfile(model_fp):
-        with urllib.request.urlopen(VAD_SEGMENTATION_URL) as source, open(model_fp, "wb") as output:
-            with tqdm(
-                total=int(source.info().get("Content-Length")),
-                ncols=80,
-                unit="iB",
-                unit_scale=True,
-                unit_divisor=1024,
-            ) as loop:
-                while True:
-                    buffer = source.read(8192)
-                    if not buffer:
-                        break
-
-                    output.write(buffer)
-                    loop.update(len(buffer))
-
-    model_bytes = open(model_fp, "rb").read()
-    if hashlib.sha256(model_bytes).hexdigest() != VAD_SEGMENTATION_URL.split('/')[-2]:
-        raise RuntimeError(
-            "Model has been downloaded but the SHA256 checksum does not not match. Please retry loading the model."
-        )
-
-    vad_model = Model.from_pretrained(model_fp, use_auth_token=use_auth_token)
-    hyperparameters = {"onset": vad_onset, 
-                    "offset": vad_offset,
-                    "min_duration_on": 0.1,
-                    "min_duration_off": 0.1}
-    vad_pipeline = VoiceActivitySegmentation(segmentation=vad_model, device=torch.device(device))
-    vad_pipeline.instantiate(hyperparameters)
-
-    return vad_pipeline
-
-class Binarize:
-    """Binarize detection scores using hysteresis thresholding, with min-cut operation
-    to ensure not segments are longer than max_duration.
-
-    Parameters
-    ----------
-    onset : float, optional
-        Onset threshold. Defaults to 0.5.
-    offset : float, optional
-        Offset threshold. Defaults to `onset`.
-    min_duration_on : float, optional
-        Remove active regions shorter than that many seconds. Defaults to 0s.
-    min_duration_off : float, optional
-        Fill inactive regions shorter than that many seconds. Defaults to 0s.
-    pad_onset : float, optional
-        Extend active regions by moving their start time by that many seconds.
-        Defaults to 0s.
-    pad_offset : float, optional
-        Extend active regions by moving their end time by that many seconds.
-        Defaults to 0s.
-    max_duration: float
-        The maximum length of an active segment, divides segment at timestamp with lowest score.
-    Reference
-    ---------
-    Gregory Gelly and Jean-Luc Gauvain. "Minimum Word Error Training of
-    RNN-based Voice Activity Detection", InterSpeech 2015.
-
-    Modified by Max Bain to include WhisperX's min-cut operation 
-    https://arxiv.org/abs/2303.00747
-    
-    Pyannote-audio
-    """
-
-    def __init__(
-        self,
-        onset: float = 0.5,
-        offset: Optional[float] = None,
-        min_duration_on: float = 0.0,
-        min_duration_off: float = 0.0,
-        pad_onset: float = 0.0,
-        pad_offset: float = 0.0,
-        max_duration: float = float('inf')
-    ):
-
-        super().__init__()
-
-        self.onset = onset
-        self.offset = offset or onset
-
-        self.pad_onset = pad_onset
-        self.pad_offset = pad_offset
-
-        self.min_duration_on = min_duration_on
-        self.min_duration_off = min_duration_off
-
-        self.max_duration = max_duration
-
-    def __call__(self, scores: SlidingWindowFeature) -> Annotation:
-        """Binarize detection scores
-        Parameters
-        ----------
-        scores : SlidingWindowFeature
-            Detection scores.
-        Returns
-        -------
-        active : Annotation
-            Binarized scores.
-        """
-
-        num_frames, num_classes = scores.data.shape
-        frames = scores.sliding_window
-        timestamps = [frames[i].middle for i in range(num_frames)]
-
-        # annotation meant to store 'active' regions
-        active = Annotation()
-        for k, k_scores in enumerate(scores.data.T):
-
-            label = k if scores.labels is None else scores.labels[k]
-
-            # initial state
-            start = timestamps[0]
-            is_active = k_scores[0] > self.onset
-            curr_scores = [k_scores[0]]
-            curr_timestamps = [start]
-            for t, y in zip(timestamps[1:], k_scores[1:]):
-                # currently active
-                if is_active: 
-                    curr_duration = t - start
-                    if curr_duration > self.max_duration:
-                        # if curr_duration > 15:
-                            # import pdb; pdb.set_trace()
-                        search_after = len(curr_scores) // 2
-                        # divide segment
-                        min_score_div_idx = search_after + np.argmin(curr_scores[search_after:])
-                        min_score_t = curr_timestamps[min_score_div_idx]
-                        region = Segment(start - self.pad_onset, min_score_t + self.pad_offset)
-                        active[region, k] = label
-                        start = curr_timestamps[min_score_div_idx]
-                        curr_scores = curr_scores[min_score_div_idx+1:]
-                        curr_timestamps = curr_timestamps[min_score_div_idx+1:]
-                    # switching from active to inactive
-                    elif y < self.offset:
-                        region = Segment(start - self.pad_onset, t + self.pad_offset)
-                        active[region, k] = label
-                        start = t
-                        is_active = False
-                        curr_scores = []
-                        curr_timestamps = []
-                # currently inactive
-                else:
-                    # switching from inactive to active
-                    if y > self.onset:
-                        start = t
-                        is_active = True
-                curr_scores.append(y)
-                curr_timestamps.append(t)
-
-            # if active at the end, add final region
-            if is_active:
-                region = Segment(start - self.pad_onset, t + self.pad_offset)
-                active[region, k] = label
-
-        # because of padding, some active regions might be overlapping: merge them.
-        # also: fill same speaker gaps shorter than min_duration_off
-        if self.pad_offset > 0.0 or self.pad_onset > 0.0 or self.min_duration_off > 0.0:
-            if self.max_duration < float("inf"):
-                raise NotImplementedError(f"This would break current max_duration param")
-            active = active.support(collar=self.min_duration_off)
-
-        # remove tracks shorter than min_duration_on
-        if self.min_duration_on > 0:
-            for segment, track in list(active.itertracks()):
-                if segment.duration < self.min_duration_on:
-                    del active[segment, track]
-
-        return active
-
-
-class VoiceActivitySegmentation(VoiceActivityDetection):
-    def __init__(
-        self,
-        segmentation: PipelineModel = "pyannote/segmentation",
-        fscore: bool = False,
-        use_auth_token: Union[Text, None] = None,
-        **inference_kwargs,
-    ):
-
-        super().__init__(segmentation=segmentation, fscore=fscore, use_auth_token=use_auth_token, **inference_kwargs)
-
-    def apply(self, file: AudioFile, hook: Optional[Callable] = None) -> Annotation:
-        """Apply voice activity detection
-
-        Parameters
-        ----------
-        file : AudioFile
-            Processed file.
-        hook : callable, optional
-            Hook called after each major step of the pipeline with the following
-            signature: hook("step_name", step_artefact, file=file)
-
-        Returns
-        -------
-        speech : Annotation
-            Speech regions.
-        """
-
-        # setup hook (e.g. for debugging purposes)
-        hook = self.setup_hook(file, hook=hook)
-
-        # apply segmentation model (only if needed)
-        # output shape is (num_chunks, num_frames, 1)
-        if self.training:
-            if self.CACHED_SEGMENTATION in file:
-                segmentations = file[self.CACHED_SEGMENTATION]
-            else:
-                segmentations = self._segmentation(file)
-                file[self.CACHED_SEGMENTATION] = segmentations
-        else:
-            segmentations: SlidingWindowFeature = self._segmentation(file)
-
-        return segmentations
-
-
-def merge_vad(vad_arr, pad_onset=0.0, pad_offset=0.0, min_duration_off=0.0, min_duration_on=0.0):
-
-    active = Annotation()
-    for k, vad_t in enumerate(vad_arr):
-        region = Segment(vad_t[0] - pad_onset, vad_t[1] + pad_offset)
-        active[region, k] = 1
-
-
-    if pad_offset > 0.0 or pad_onset > 0.0 or min_duration_off > 0.0:
-        active = active.support(collar=min_duration_off)
-    
-    # remove tracks shorter than min_duration_on
-    if min_duration_on > 0:
-        for segment, track in list(active.itertracks()):
-            if segment.duration < min_duration_on:
-                    del active[segment, track]
-    
-    active = active.for_json()
-    active_segs = pd.DataFrame([x['segment'] for x in active['content']])
-    return active_segs
-
-def merge_chunks(segments, chunk_size):
-    """
-    Merge operation described in paper
-    """
-    curr_end = 0
-    merged_segments = []
-    seg_idxs = []
-    speaker_idxs = []
-
-    assert chunk_size > 0
-    binarize = Binarize(max_duration=chunk_size)
-    segments = binarize(segments)
-    segments_list = []
-    for speech_turn in segments.get_timeline():
-        segments_list.append(SegmentX(speech_turn.start, speech_turn.end, "UNKNOWN"))
-
-    if len(segments_list) == 0:
-        print("No active speech found in audio")
-        return []
-    # assert segments_list, "segments_list is empty."
-    # Make sur the starting point is the start of the segment.
-    curr_start = segments_list[0].start
-
-    for seg in segments_list:
-        if seg.end - curr_start > chunk_size and curr_end-curr_start > 0:
-            merged_segments.append({
-                "start": curr_start,
-                "end": curr_end,
-                "segments": seg_idxs,
-            })
-            curr_start = seg.start
-            seg_idxs = []
-            speaker_idxs = []
-        curr_end = seg.end
-        seg_idxs.append((seg.start, seg.end))
-        speaker_idxs.append(seg.speaker)
-    # add final
-    merged_segments.append({ 
-                "start": curr_start,
-                "end": curr_end,
-                "segments": seg_idxs,
-            })    
-    return merged_segments