whisperX/whisperx/asr.py

import os
import warnings
from typing import List, Union

import ctranslate2
import faster_whisper
import numpy as np
import torch
from transformers import Pipeline
from transformers.pipelines.pt_utils import PipelineIterator

from .audio import N_SAMPLES, SAMPLE_RATE, load_audio, log_mel_spectrogram
from .vad import load_vad_model, merge_chunks


def load_model(whisper_arch, device, compute_type="float16", asr_options=None, language=None,
               vad_options=None, model=None):
    '''Load a Whisper model for inference.
    Args:
        whisper_arch: str - The name of the Whisper model to load.
        device: str - The device to load the model on.
        compute_type: str - The compute type to use for the model.
        options: dict - A dictionary of options to use for the model.
        language: str - The language of the model. (use English for now)
    Returns:
        A Whisper pipeline.
    '''    

    if whisper_arch.endswith(".en"):
        language = "en"

    model = WhisperModel(whisper_arch, device=device, compute_type=compute_type)
    if language is not None:
        tokenizer = faster_whisper.tokenizer.Tokenizer(model.hf_tokenizer, model.model.is_multilingual, task="transcribe", language=language)
    else:
        print("No language specified, language will be first be detected for each audio file (increases inference time).")
        tokenizer = None

    default_asr_options =  {
        "beam_size": 5,
        "best_of": 5,
        "patience": 1,
        "length_penalty": 1,
        "temperatures": [0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
        "compression_ratio_threshold": 2.4,
        "log_prob_threshold": -1.0,
        "no_speech_threshold": 0.6,
        "condition_on_previous_text": False,
        "initial_prompt": None,
        "prefix": None,
        "suppress_blank": True,
        "suppress_tokens": [-1],
        "without_timestamps": True,
        "max_initial_timestamp": 0.0,
        "word_timestamps": False,
        "prepend_punctuations": "\"'“¿([{-",
        "append_punctuations": "\"'.。,，!！?？:：”)]}、"
    }

    if asr_options is not None:
        default_asr_options.update(asr_options)
    default_asr_options = faster_whisper.transcribe.TranscriptionOptions(**default_asr_options)

    default_vad_options = {
        "vad_onset": 0.500,
        "vad_offset": 0.363
    }

    if vad_options is not None:
        default_vad_options.update(vad_options)

    vad_model = load_vad_model(torch.device(device), use_auth_token=None, **default_vad_options)

    return FasterWhisperPipeline(model, vad_model, default_asr_options, tokenizer)



class WhisperModel(faster_whisper.WhisperModel):
    '''
    FasterWhisperModel provides batched inference for faster-whisper.
    Currently only works in non-timestamp mode.
    '''

    def generate_segment_batched(self, features: np.ndarray, tokenizer: faster_whisper.tokenizer.Tokenizer, options: faster_whisper.transcribe.TranscriptionOptions, encoder_output = None):
        batch_size = features.shape[0]
        all_tokens = []
        prompt_reset_since = 0
        if options.initial_prompt is not None:
            initial_prompt = " " + options.initial_prompt.strip()
            initial_prompt_tokens = tokenizer.encode(initial_prompt)
            all_tokens.extend(initial_prompt_tokens)
        previous_tokens = all_tokens[prompt_reset_since:]
        prompt = self.get_prompt(
            tokenizer,
            previous_tokens,
            without_timestamps=options.without_timestamps,
            prefix=options.prefix,
        )

        encoder_output = self.encode(features)

        max_initial_timestamp_index = int(
            round(options.max_initial_timestamp / self.time_precision)
        )

        result = self.model.generate(
                encoder_output,
                [prompt] * batch_size,
                # length_penalty=options.length_penalty,
                # max_length=self.max_length,
                # return_scores=True,
                # return_no_speech_prob=True,
                # suppress_blank=options.suppress_blank,
                # suppress_tokens=options.suppress_tokens,
                # max_initial_timestamp_index=max_initial_timestamp_index,
            )
    
        tokens_batch = [x.sequences_ids[0] for x in result]

        def decode_batch(tokens: List[List[int]]) -> str:
            res = []
            for tk in tokens:
                res.append([token for token in tk if token < tokenizer.eot])
            # text_tokens = [token for token in tokens if token < self.eot]
            return tokenizer.tokenizer.decode_batch(res)

        text = decode_batch(tokens_batch)

        return text
    
    def encode(self, features: np.ndarray) -> ctranslate2.StorageView:
        # When the model is running on multiple GPUs, the encoder output should be moved
        # to the CPU since we don't know which GPU will handle the next job.
        to_cpu = self.model.device == "cuda" and len(self.model.device_index) > 1
        # unsqueeze if batch size = 1
        if len(features.shape) == 2:
            features = np.expand_dims(features, 0)
        features = faster_whisper.transcribe.get_ctranslate2_storage(features)
        
        return self.model.encode(features, to_cpu=to_cpu)
    
class FasterWhisperPipeline(Pipeline):
    def __init__(
            self,
            model,
            vad,
            options,
            tokenizer=None,
            device: Union[int, str, "torch.device"] = -1,
            framework = "pt",
            **kwargs
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.options = options
        self._batch_size = kwargs.pop("batch_size", None)
        self._num_workers = 1
        self._preprocess_params, self._forward_params, self._postprocess_params = self._sanitize_parameters(**kwargs)
        self.call_count = 0
        self.framework = framework
        if self.framework == "pt":
            if isinstance(device, torch.device):
                self.device = device
            elif isinstance(device, str):
                self.device = torch.device(device)
            elif device < 0:
                self.device = torch.device("cpu")
            else:
                self.device = torch.device(f"cuda:{device}")
        else:
            self.device = device
        
        super(Pipeline, self).__init__()
        self.vad_model = vad

    def _sanitize_parameters(self, **kwargs):
        preprocess_kwargs = {}
        if "tokenizer" in kwargs:
            preprocess_kwargs["maybe_arg"] = kwargs["maybe_arg"]
        return preprocess_kwargs, {}, {}

    def preprocess(self, audio):
        audio = audio['inputs']
        features = log_mel_spectrogram(audio, padding=N_SAMPLES - audio.shape[0])
        return {'inputs': features}

    def _forward(self, model_inputs):
        outputs = self.model.generate_segment_batched(model_inputs['inputs'], self.tokenizer, self.options)
        return {'text': outputs}
    
    def postprocess(self, model_outputs):
        return model_outputs

    def get_iterator(
        self, inputs, num_workers: int, batch_size: int, preprocess_params, forward_params, postprocess_params
    ):
        dataset = PipelineIterator(inputs, self.preprocess, preprocess_params)
        if "TOKENIZERS_PARALLELISM" not in os.environ:
            os.environ["TOKENIZERS_PARALLELISM"] = "false"
        # TODO hack by collating feature_extractor and image_processor

        def stack(items):
            return {'inputs': torch.stack([x['inputs'] for x in items])}
        dataloader = torch.utils.data.DataLoader(dataset, num_workers=num_workers, batch_size=batch_size, collate_fn=stack)
        model_iterator = PipelineIterator(dataloader, self.forward, forward_params, loader_batch_size=batch_size)
        final_iterator = PipelineIterator(model_iterator, self.postprocess, postprocess_params)
        return final_iterator

    def transcribe(
        self, audio: Union[str, np.ndarray], batch_size=None
    ):
        if isinstance(audio, str):
            audio = load_audio(audio)
        
        def data(audio, segments):
            for seg in segments:
                f1 = int(seg['start'] * SAMPLE_RATE)
                f2 = int(seg['end'] * SAMPLE_RATE)
                # print(f2-f1)
                yield {'inputs': audio[f1:f2]}

        vad_segments = self.vad_model({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": SAMPLE_RATE})
        vad_segments = merge_chunks(vad_segments, 30)

        del_tokenizer = False
        if self.tokenizer is None:
            language = self.detect_language(audio)
            self.tokenizer = faster_whisper.tokenizer.Tokenizer(self.model.hf_tokenizer, self.model.model.is_multilingual, task="transcribe", language=language)
            del_tokenizer = True
        else:
            language = self.tokenizer.language_code

        segments = []
        batch_size = batch_size or self._batch_size
        for idx, out in enumerate(self.__call__(data(audio, vad_segments), batch_size=batch_size)):
            text = out['text']
            if batch_size in [0, 1, None]:
                text = text[0]
            segments.append(
                {
                    "text": out['text'],
                    "start": round(vad_segments[idx]['start'], 3),
                    "end": round(vad_segments[idx]['end'], 3)
                }
            )
        
        if del_tokenizer:
            self.tokenizer = None

        return {"segments": segments, "language": language}


    def detect_language(self, audio: np.ndarray):
        segment = log_mel_spectrogram(audio[: N_SAMPLES], padding=0)
        encoder_output = self.model.encode(segment)
        results = self.model.model.detect_language(encoder_output)
        language_token, language_probability = results[0][0]
        language = language_token[2:-2]
        print(f"Detected language: {language} ({language_probability:.2f}) in first 30s of audio...")
        return language

if __name__ == "__main__":
    main_type = "simple"
    import time

    import jiwer
    from tqdm import tqdm
    from whisper.normalizers import EnglishTextNormalizer

    from benchmark.tedlium import parse_tedlium_annos

    if main_type == "complex":
        from faster_whisper.tokenizer import Tokenizer
        from faster_whisper.transcribe import TranscriptionOptions
        from faster_whisper.vad import (SpeechTimestampsMap,
                                        get_speech_timestamps)

        from whisperx.vad import load_vad_model, merge_chunks

        from .audio import SAMPLE_RATE, load_audio, log_mel_spectrogram
        faster_t_options = TranscriptionOptions(
        beam_size=5,
        best_of=5,
        patience=1,
        length_penalty=1,
        temperatures=[0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
        compression_ratio_threshold=2.4,
        log_prob_threshold=-1.0,
        no_speech_threshold=0.6,
        condition_on_previous_text=False,
        initial_prompt=None,
        prefix=None,
        suppress_blank=True,
        suppress_tokens=[-1],
        without_timestamps=True,
        max_initial_timestamp=0.0,
        word_timestamps=False,
        prepend_punctuations="\"'“¿([{-",
        append_punctuations="\"'.。,，!！?？:：”)]}、"
    )
        whisper_arch = "large-v2"
        device = "cuda"
        batch_size = 16
        model = WhisperModel(whisper_arch, device="cuda", compute_type="float16",)
        tokenizer = Tokenizer(model.hf_tokenizer, model.model.is_multilingual, task="transcribe", language="en")
        model = FasterWhisperPipeline(model, tokenizer, faster_t_options, device=-1)
        fn = "DanielKahneman_2010.wav"
        wav_dir = f"/tmp/test/wav/"
        vad_model = load_vad_model("cuda", 0.6, 0.3)
        audio = load_audio(os.path.join(wav_dir, fn))
        vad_segments = vad_model({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": SAMPLE_RATE})
        vad_segments = merge_chunks(vad_segments, 30)

        def data(audio, segments):
            for seg in segments:
                f1 = int(seg['start'] * SAMPLE_RATE)
                f2 = int(seg['end'] * SAMPLE_RATE)
                # print(f2-f1)
                yield {'inputs': audio[f1:f2]}
        vad_method="pyannote"

        wav_dir = f"/tmp/test/wav/"
        wer_li = []
        time_li = []
        for fn in os.listdir(wav_dir):
            if fn == "RobertGupta_2010U.wav":
                continue
            base_fn = fn.split('.')[0]
            audio_fp = os.path.join(wav_dir, fn)

            audio = load_audio(audio_fp)
            t1 = time.time()
            if vad_method == "pyannote":
                vad_segments = vad_model({"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": SAMPLE_RATE})
                vad_segments = merge_chunks(vad_segments, 30)
            elif vad_method == "silero":
                vad_segments = get_speech_timestamps(audio, threshold=0.5, max_speech_duration_s=30)
                vad_segments = [{"start": x["start"] / SAMPLE_RATE, "end": x["end"] / SAMPLE_RATE} for x in vad_segments]
                new_segs = []
                curr_start = vad_segments[0]['start']
                curr_end = vad_segments[0]['end']
                for seg in vad_segments[1:]:
                    if seg['end'] - curr_start > 30:
                        new_segs.append({"start": curr_start, "end": curr_end})
                        curr_start = seg['start']
                        curr_end = seg['end']
                    else:
                        curr_end = seg['end']
                new_segs.append({"start": curr_start, "end": curr_end})
                vad_segments = new_segs
            text = []
            # for idx, out in tqdm(enumerate(model(data(audio_fp, vad_segments), batch_size=batch_size)), total=len(vad_segments)):
            for idx, out in enumerate(model(data(audio, vad_segments), batch_size=batch_size)):
                text.append(out['text'])
            t2 = time.time()
            if batch_size == 1:
                text = [x[0] for x in text]
            text = " ".join(text)

            normalizer = EnglishTextNormalizer()
            text = normalizer(text)
            gt_corpus = normalizer(parse_tedlium_annos(base_fn, "/tmp/test/"))

            wer_result = jiwer.wer(gt_corpus, text)
            print("WER: %.2f \t time: %.2f \t [%s]" % (wer_result * 100, t2-t1, fn))

            wer_li.append(wer_result)
            time_li.append(t2-t1)
        print("# Avg Mean...")
        print("WER: %.2f" % (sum(wer_li) * 100/len(wer_li)))
        print("Time: %.2f" % (sum(time_li)/len(time_li)))
    elif main_type == "simple":
        model = load_model(
            "large-v2",
            device="cuda",
            language="en",
        )

        wav_dir = f"/tmp/test/wav/"
        wer_li = []
        time_li = []
        for fn in os.listdir(wav_dir):
            if fn == "RobertGupta_2010U.wav":
                continue
            # fn = "DanielKahneman_2010.wav"
            base_fn = fn.split('.')[0]
            audio_fp = os.path.join(wav_dir, fn)

            audio = load_audio(audio_fp)
            t1 = time.time()
            out = model.transcribe(audio_fp, batch_size=8)["segments"]
            t2 = time.time()

            text = " ".join([x['text'] for x in out])
            normalizer = EnglishTextNormalizer()
            text = normalizer(text)
            gt_corpus = normalizer(parse_tedlium_annos(base_fn, "/tmp/test/"))

            wer_result = jiwer.wer(gt_corpus, text)
            print("WER: %.2f \t time: %.2f \t [%s]" % (wer_result * 100, t2-t1, fn))

            wer_li.append(wer_result)
            time_li.append(t2-t1)
        print("# Avg Mean...")
        print("WER: %.2f" % (sum(wer_li) * 100/len(wer_li)))
        print("Time: %.2f" % (sum(time_li)/len(time_li)))