whisperX merged in

.gitignore

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

EXAMPLES.md

@@ -0,0 +1,37 @@
+# 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

@@ -0,0 +1,27 @@
+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

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

requirements.txt

@@ -0,0 +1,10 @@
+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

@@ -0,0 +1,28 @@
+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

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

whisperx/__main__.py

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

whisperx/alignment.py

@@ -0,0 +1,548 @@
+""""
+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

@@ -0,0 +1,429 @@
+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

@@ -0,0 +1,76 @@
+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

@@ -0,0 +1,220 @@
+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

@@ -0,0 +1,317 @@
+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

@@ -0,0 +1,305 @@
+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