import os from typing import List, Optional, Union from dataclasses import replace import ctranslate2 import faster_whisper import numpy as np import torch from faster_whisper.tokenizer import Tokenizer from faster_whisper.transcribe import TranscriptionOptions, get_ctranslate2_storage from transformers import Pipeline from transformers.pipelines.pt_utils import PipelineIterator from .audio import N_SAMPLES, SAMPLE_RATE, load_audio, log_mel_spectrogram from .types import SingleSegment, TranscriptionResult from .vads import Vad, Silero, Pyannote def find_numeral_symbol_tokens(tokenizer): """ Finds tokens that represent numeral and symbols. """ numeral_symbol_tokens = [] for i in range(tokenizer.eot): token = tokenizer.decode([i]).removeprefix(" ") has_numeral_symbol = any(c in "0123456789%$£" for c in token) if has_numeral_symbol: numeral_symbol_tokens.append(i) return numeral_symbol_tokens class WhisperModel(faster_whisper.WhisperModel): """ Wrapper around faster-whisper's WhisperModel to enable batched inference. Currently, it only supports non-timestamp mode and a fixed prompt for all samples in a batch. """ def generate_segment_batched( self, features: np.ndarray, tokenizer: Tokenizer, options: TranscriptionOptions, encoder_output=None, ): """ Generates transcription for a batch of audio segments. Args: features: The input audio features. tokenizer: The tokenizer used to decode the generated tokens. options: Transcription options. encoder_output: Output from the encoder model. Returns: The decoded transcription text. """ batch_size = features.shape[0] # Initialize tokens and prompt for the generation process. all_tokens = [] prompt_reset_since = 0 # Check if an initial prompt is provided and handle it. 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) # Prepare the prompt for the current batch. previous_tokens = all_tokens[prompt_reset_since:] prompt = self.get_prompt( tokenizer, previous_tokens, without_timestamps=options.without_timestamps, prefix=options.prefix, ) # Encode the features to obtain the encoder output. encoder_output = self.encode(features) # Determine the maximum initial timestamp index based on the options. max_initial_timestamp_index = int( round(options.max_initial_timestamp / self.time_precision) ) # Generate the transcription result for the batch. result = self.model.generate( encoder_output, [prompt] * batch_size, beam_size=options.beam_size, patience=options.patience, length_penalty=options.length_penalty, max_length=self.max_length, suppress_blank=options.suppress_blank, suppress_tokens=options.suppress_tokens, ) # Extract the token sequences from the result. tokens_batch = [x.sequences_ids[0] for x in result] # Define an inner function to decode the tokens for each batch. def decode_batch(tokens: List[List[int]]) -> str: res = [] for tk in tokens: res.append([token for token in tk if token < tokenizer.eot]) return tokenizer.tokenizer.decode_batch(res) # Decode the tokens to get the transcription text. text = decode_batch(tokens_batch) return text def encode(self, features: np.ndarray) -> ctranslate2.StorageView: """ Encodes the audio features using the CTranslate2 storage. 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. """ # When the model is running on multiple GPUs, the encoder output should be moved to the CPU. to_cpu = self.model.device == "cuda" and len(self.model.device_index) > 1 # If the batch size is 1, unsqueeze the features to ensure it is a 3D array. if len(features.shape) == 2: features = np.expand_dims(features, 0) features = get_ctranslate2_storage(features) # call the model return self.model.encode(features, to_cpu=to_cpu) def get_iterator( self, inputs, num_workers: int, batch_size: int, preprocess_params: dict, forward_params: dict, postprocess_params: dict, ): 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: Optional[int] = None, num_workers=0, language: Optional[str] = None, task: Optional[str] = None, chunk_size=30, print_progress=False, combined_progress=False, verbose=False, ) -> TranscriptionResult: 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]} # Pre-process audio and merge chunks as defined by the respective VAD child class # In case vad_model is manually assigned (see 'load_model') follow the functionality of pyannote toolkit if issubclass(type(self.vad_model), Vad): waveform = self.vad_model.preprocess_audio(audio) merge_chunks = self.vad_model.merge_chunks else: waveform = Pyannote.preprocess_audio(audio) merge_chunks = Pyannote.merge_chunks vad_segments = self.vad_model({"waveform": waveform, "sample_rate": SAMPLE_RATE}) vad_segments = merge_chunks( vad_segments, chunk_size, onset=self._vad_params["vad_onset"], offset=self._vad_params["vad_offset"], ) if self.tokenizer is None: language = language or self.detect_language(audio) task = task or "transcribe" self.tokenizer = Tokenizer( self.model.hf_tokenizer, self.model.model.is_multilingual, task=task, language=language, ) else: language = language or self.tokenizer.language_code task = task or self.tokenizer.task if task != self.tokenizer.task or language != self.tokenizer.language_code: self.tokenizer = Tokenizer( self.model.hf_tokenizer, self.model.model.is_multilingual, task=task, language=language, ) if self.suppress_numerals: previous_suppress_tokens = self.options.suppress_tokens numeral_symbol_tokens = find_numeral_symbol_tokens(self.tokenizer) print(f"Suppressing numeral and symbol tokens") new_suppressed_tokens = numeral_symbol_tokens + self.options.suppress_tokens new_suppressed_tokens = list(set(new_suppressed_tokens)) self.options = replace(self.options, suppress_tokens=new_suppressed_tokens) segments: List[SingleSegment] = [] batch_size = batch_size or self._batch_size total_segments = len(vad_segments) for idx, out in enumerate(self.__call__(data(audio, vad_segments), batch_size=batch_size, num_workers=num_workers)): if print_progress: base_progress = ((idx + 1) / total_segments) * 100 percent_complete = base_progress / 2 if combined_progress else base_progress print(f"Progress: {percent_complete:.2f}%...") text = out['text'] if batch_size in [0, 1, None]: text = text[0] if verbose: print(f"Transcript: [{round(vad_segments[idx]['start'], 3)} --> {round(vad_segments[idx]['end'], 3)}] {text}") segments.append( { "text": text, "start": round(vad_segments[idx]['start'], 3), "end": round(vad_segments[idx]['end'], 3) } ) # revert the tokenizer if multilingual inference is enabled if self.preset_language is None: self.tokenizer = None # revert suppressed tokens if suppress_numerals is enabled if self.suppress_numerals: self.options = replace(self.options, suppress_tokens=previous_suppress_tokens) return {"segments": segments, "language": language} def detect_language(self, audio: np.ndarray) -> str: if audio.shape[0] < N_SAMPLES: print("Warning: audio is shorter than 30s, language detection may be inaccurate.") model_n_mels = self.model.feat_kwargs.get("feature_size") segment = log_mel_spectrogram(audio[: N_SAMPLES], n_mels=model_n_mels if model_n_mels is not None else 80, padding=0 if audio.shape[0] >= N_SAMPLES else N_SAMPLES - audio.shape[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 def load_model( whisper_arch: str, device: str, device_index=0, compute_type="float16", asr_options: Optional[dict] = None, language: Optional[str] = None, vad_model: Optional[Vad]= None, vad_method: Optional[str] = "pyannote", vad_options: Optional[dict] = None, model: Optional[WhisperModel] = None, task="transcribe", download_root: Optional[str] = None, local_files_only=False, threads=4, ) -> FasterWhisperPipeline: """Load a Whisper model for inference. Args: whisper_arch - The name of the Whisper model to load. device - The device to load the model on. compute_type - The compute type to use for the model. vad_method - The vad method to use. vad_model has higher priority if is not None. options - A dictionary of options to use for the model. language - The language of the model. (use English for now) model - The WhisperModel instance to use. download_root - The root directory to download the model to. local_files_only - If `True`, avoid downloading the file and return the path to the local cached file if it exists. threads - The number of cpu threads to use per worker, e.g. will be multiplied by num workers. Returns: A Whisper pipeline. """ if whisper_arch.endswith(".en"): language = "en" model = model or WhisperModel(whisper_arch, device=device, device_index=device_index, compute_type=compute_type, download_root=download_root, local_files_only=local_files_only, cpu_threads=threads) if language is not None: tokenizer = Tokenizer(model.hf_tokenizer, model.model.is_multilingual, task=task, 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, "repetition_penalty": 1, "no_repeat_ngram_size": 0, "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, "prompt_reset_on_temperature": 0.5, "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": "\"'.。,，!！?？:：”)]}、", "multilingual": model.model.is_multilingual, "suppress_numerals": False, "max_new_tokens": None, "clip_timestamps": None, "hallucination_silence_threshold": None, "hotwords": None, } if asr_options is not None: default_asr_options.update(asr_options) suppress_numerals = default_asr_options["suppress_numerals"] del default_asr_options["suppress_numerals"] default_asr_options = TranscriptionOptions(**default_asr_options) default_vad_options = { "chunk_size": 30, # needed by silero since binarization happens before merge_chunks "vad_onset": 0.500, "vad_offset": 0.363 } if vad_options is not None: default_vad_options.update(vad_options) # Note: manually assigned vad_model has higher priority than vad_method! if vad_model is not None: print("Use manually assigned vad_model. vad_method is ignored.") vad_model = vad_model else: if vad_method == "silero": vad_model = Silero(**default_vad_options) elif vad_method == "pyannote": vad_model = Pyannote(torch.device(device), use_auth_token=None, **default_vad_options) else: raise ValueError(f"Invalid vad_method: {vad_method}") return FasterWhisperPipeline( model=model, vad=vad_model, options=default_asr_options, tokenizer=tokenizer, language=language, suppress_numerals=suppress_numerals, vad_params=default_vad_options, )