import copy import re import warnings import torch import numpy as np from tqdm import tqdm from typing import TYPE_CHECKING, Union, List, Callable, Optional, Tuple import whisper from whisper.audio import ( SAMPLE_RATE, N_FRAMES, N_SAMPLES, N_FFT, pad_or_trim, log_mel_spectrogram, FRAMES_PER_SECOND, CHUNK_LENGTH ) from .result import WhisperResult, Segment from .timing import add_word_timestamps_stable, split_word_tokens from .audio import prep_audio from .utils import safe_print, format_timestamp from .whisper_compatibility import warn_compatibility_issues, get_tokenizer from .stabilization import get_vad_silence_func, wav2mask, mask2timing if TYPE_CHECKING: from whisper.model import Whisper __all__ = ['align', 'refine', 'locate'] def align( model: "Whisper", audio: Union[str, np.ndarray, torch.Tensor, bytes], text: Union[str, List[int], WhisperResult], language: str = None, *, verbose: Optional[bool] = False, regroup: bool = True, suppress_silence: bool = True, suppress_word_ts: bool = True, use_word_position: bool = True, min_word_dur: bool = 0.1, nonspeech_error: float = 0.3, q_levels: int = 20, k_size: int = 5, vad: bool = False, vad_threshold: float = 0.35, vad_onnx: bool = False, demucs: Union[bool, torch.nn.Module] = False, demucs_output: str = None, demucs_options: dict = None, only_voice_freq: bool = False, prepend_punctuations: str = "\"'“¿([{-", append_punctuations: str = "\"'.。,，!！?？:：”)]}、", progress_callback: Callable = None, ignore_compatibility: bool = False, remove_instant_words: bool = False, token_step: int = 100, original_split: bool = False, word_dur_factor: Optional[float] = 2.0, max_word_dur: Optional[float] = 3.0, nonspeech_skip: Optional[float] = 3.0, fast_mode: bool = False, tokenizer: "Tokenizer" = None ) -> Union[WhisperResult, None]: """ Align plain text or tokens with audio at word-level. Since this is significantly faster than transcribing, it is a more efficient method for testing various settings without re-transcribing. This is also useful for timing a more correct transcript than one that Whisper can produce. Parameters ---------- model : "Whisper" The Whisper ASR model modified instance audio : str or numpy.ndarray or torch.Tensor or bytes Path/URL to the audio file, the audio waveform, or bytes of audio file. If audio is :class:`numpy.ndarray` or :class:`torch.Tensor`, the audio must be already at sampled to 16kHz. text : str or list of int or stable_whisper.result.WhisperResult String of plain-text, list of tokens, or instance of :class:`stable_whisper.result.WhisperResult`. language : str, default None, uses ``language`` in ``text`` if it is a :class:`stable_whisper.result.WhisperResult` Language of ``text``. Required if ``text`` does not contain ``language``. remove_instant_words : bool, default False Whether to truncate any words with zero duration. token_step : int, default 100 Max number of tokens to align each pass. Use higher values to reduce chance of misalignment. original_split : bool, default False Whether to preserve the original segment groupings. Segments are spit by line break if ``text`` is plain-text. max_word_dur : float or None, default 3.0 Global maximum word duration in seconds. Re-align words that exceed the global maximum word duration. word_dur_factor : float or None, default 2.0 Factor to compute the Local maximum word duration, which is ``word_dur_factor`` * local medium word duration. Words that need re-alignment, are re-algined with duration <= local/global maximum word duration. nonspeech_skip : float or None, default 3.0 Skip non-speech sections that are equal or longer than this duration in seconds. Disable skipping if ``None``. fast_mode : bool, default False Whether to speed up alignment by re-alignment with local/global maximum word duration. ``True`` tends produce better timestamps when ``text`` is accurate and there are no large speechless gaps. tokenizer : "Tokenizer", default None, meaning a new tokenizer is created according ``language`` and ``model`` A tokenizer to used tokenizer text and detokenize tokens. verbose : bool or None, default False Whether to display the text being decoded to the console. Displays all the details if ``True``. Displays progressbar if ``False``. Display nothing if ``None``. regroup : bool or str, default True, meaning the default regroup algorithm String for customizing the regrouping algorithm. False disables regrouping. Ignored if ``word_timestamps = False``. suppress_silence : bool, default True Whether to enable timestamps adjustments based on the detected silence. suppress_word_ts : bool, default True Whether to adjust word timestamps based on the detected silence. Only enabled if ``suppress_silence = True``. use_word_position : bool, default True Whether to use position of the word in its segment to determine whether to keep end or start timestamps if adjustments are required. If it is the first word, keep end. Else if it is the last word, keep the start. q_levels : int, default 20 Quantization levels for generating timestamp suppression mask; ignored if ``vad = true``. Acts as a threshold to marking sound as silent. Fewer levels will increase the threshold of volume at which to mark a sound as silent. k_size : int, default 5 Kernel size for avg-pooling waveform to generate timestamp suppression mask; ignored if ``vad = true``. Recommend 5 or 3; higher sizes will reduce detection of silence. demucs : bool or torch.nn.Module, default False Whether to preprocess ``audio`` with Demucs to isolate vocals / remove noise. Set ``demucs`` to an instance of a Demucs model to avoid reloading the model for each run. Demucs must be installed to use. Official repo, https://github.com/facebookresearch/demucs. demucs_output : str, optional Path to save the vocals isolated by Demucs as WAV file. Ignored if ``demucs = False``. Demucs must be installed to use. Official repo, https://github.com/facebookresearch/demucs. demucs_options : dict, optional Options to use for :func:`stable_whisper.audio.demucs_audio`. vad : bool, default False Whether to use Silero VAD to generate timestamp suppression mask. Silero VAD requires PyTorch 1.12.0+. Official repo, https://github.com/snakers4/silero-vad. vad_threshold : float, default 0.35 Threshold for detecting speech with Silero VAD. Low threshold reduces false positives for silence detection. vad_onnx : bool, default False Whether to use ONNX for Silero VAD. min_word_dur : float, default 0.1 Shortest duration each word is allowed to reach for silence suppression. nonspeech_error : float, default 0.3 Relative error of non-speech sections that appear in between a word for silence suppression. only_voice_freq : bool, default False Whether to only use sound between 200 - 5000 Hz, where majority of human speech are. prepend_punctuations : str, default '"'“¿([{-)' Punctuations to prepend to next word. append_punctuations : str, default '.。,，!！?？:：”)]}、)' Punctuations to append to previous word. progress_callback : Callable, optional A function that will be called when transcription progress is updated. The callback need two parameters. The first parameter is a float for seconds of the audio that has been transcribed. The second parameter is a float for total duration of audio in seconds. ignore_compatibility : bool, default False Whether to ignore warnings for compatibility issues with the detected Whisper version. Returns ------- stable_whisper.result.WhisperResult or None All timestamps, words, probabilities, and other data from the alignment of ``audio``. Return None if alignment fails and ``remove_instant_words = True``. Notes ----- If ``token_step`` is less than 1, ``token_step`` will be set to its maximum value, 442. This value is computed with ``whisper.model.Whisper.dims.n_text_ctx`` - 6. IF ``original_split = True`` and a line break is found in middle of a word in ``text``, the split will occur after that word. ``regroup`` is ignored if ``original_split = True``. Examples -------- >>> import stable_whisper >>> model = stable_whisper.load_model('base') >>> result = model.align('helloworld.mp3', 'Hello, World!', 'English') >>> result.to_srt_vtt('helloword.srt') Saved 'helloworld.srt' """ is_faster_model = model.__module__.startswith('faster_whisper.') if demucs_options is None: demucs_options = {} if demucs_output: if 'save_path' not in demucs_options: demucs_options['save_path'] = demucs_output warnings.warn('``demucs_output`` is deprecated. Use ``demucs_options`` with ``save_path`` instead. ' 'E.g. demucs_options=dict(save_path="demucs_output.mp3")', DeprecationWarning, stacklevel=2) max_token_step = (model.max_length if is_faster_model else model.dims.n_text_ctx) - 6 if token_step < 1: token_step = max_token_step elif token_step > max_token_step: raise ValueError(f'The max value for [token_step] is {max_token_step} but got {token_step}.') warn_compatibility_issues(whisper, ignore_compatibility) split_indices_by_char = [] if isinstance(text, WhisperResult): if language is None: language = text.language if original_split and len(text.segments) > 1 and text.has_words: split_indices_by_char = np.cumsum([sum(len(w.word) for w in seg.words) for seg in text.segments]) text = text.all_tokens() if text.has_words and all(w.tokens for w in text.all_words()) else text.text elif isinstance(text, str): if original_split and '\n' in text: text_split = [line if line.startswith(' ') else ' '+line for line in text.splitlines()] split_indices_by_char = np.cumsum([len(seg) for seg in text_split]) text = ''.join(re.sub(r'\s', ' ', seg) for seg in text_split) else: text = re.sub(r'\s', ' ', text) if not text.startswith(' '): text = ' ' + text if language is None: raise TypeError('expected argument for language') if tokenizer is None: tokenizer = get_tokenizer(model, is_faster_model=is_faster_model, language=language, task='transcribe') tokens = tokenizer.encode(text) if isinstance(text, str) else text tokens = [t for t in tokens if t < tokenizer.eot] _, (words, word_tokens), _ = split_word_tokens([dict(tokens=tokens)], tokenizer) audio = prep_audio( audio, demucs=demucs, demucs_options=demucs_options, only_voice_freq=only_voice_freq, verbose=verbose ) sample_padding = int(N_FFT // 2) + 1 seek_sample = 0 total_samples = audio.shape[-1] total_duration = round(total_samples / SAMPLE_RATE, 2) total_words = len(words) if is_faster_model: def timestamp_words(): temp_segment = dict( seek=0, start=0.0, end=round(segment_samples / model.feature_extractor.sampling_rate, 3), tokens=[t for wt in curr_word_tokens for t in wt], ) features = model.feature_extractor(audio_segment.numpy()) encoder_output = model.encode(features[:, : model.feature_extractor.nb_max_frames]) model.add_word_timestamps( segments=[temp_segment], tokenizer=tokenizer, encoder_output=encoder_output, num_frames=round(segment_samples / model.feature_extractor.hop_length), prepend_punctuations=prepend_punctuations, append_punctuations=append_punctuations, last_speech_timestamp=temp_segment['start'], ) cumsum_lens = np.cumsum([len(w) for w in curr_words]).tolist() final_cumsum_lens = np.cumsum([len(w['word']) for w in temp_segment['words']]).tolist() assert not (set(final_cumsum_lens) - set(cumsum_lens)), 'word mismatch' prev_l_idx = 0 for w_idx, cs_len in enumerate(final_cumsum_lens): temp_segment['words'][w_idx]['start'] = round(temp_segment['words'][w_idx]['start'] + time_offset, 3) temp_segment['words'][w_idx]['end'] = round(temp_segment['words'][w_idx]['end'] + time_offset, 3) l_idx = cumsum_lens.index(cs_len)+1 temp_segment['words'][w_idx]['tokens'] = [t for wt in curr_word_tokens[prev_l_idx:l_idx] for t in wt] prev_l_idx = l_idx return temp_segment else: def timestamp_words(): temp_segment = dict( seek=time_offset, tokens=(curr_words, curr_word_tokens) ) mel_segment = log_mel_spectrogram(audio_segment, model.dims.n_mels, padding=sample_padding) mel_segment = pad_or_trim(mel_segment, N_FRAMES).to(device=model.device) add_word_timestamps_stable( segments=[temp_segment], model=model, tokenizer=tokenizer, mel=mel_segment, num_samples=segment_samples, split_callback=(lambda x, _: x), prepend_punctuations=prepend_punctuations, append_punctuations=append_punctuations, gap_padding=None ) return temp_segment def get_curr_words(): nonlocal words, word_tokens curr_tk_count = 0 w, wt = [], [] for _ in range(len(words)): tk_count = len(word_tokens[0]) if curr_tk_count + tk_count > token_step and w: break w.append(words.pop(0)) wt.append(word_tokens.pop(0)) curr_tk_count += tk_count return w, wt result = [] nonspeech_timings = [[], []] nonspeech_vad_timings = None if (suppress_silence or nonspeech_skip is not None) and vad: nonspeech_vad_timings = ( get_vad_silence_func(onnx=vad_onnx, verbose=verbose)(audio, speech_threshold=vad_threshold) ) if nonspeech_vad_timings is not None: nonspeech_timings = nonspeech_vad_timings[0].copy(), nonspeech_vad_timings[1].copy() with tqdm(total=total_duration, unit='sec', disable=verbose is not False, desc='Align') as tqdm_pbar: def update_pbar(finish: bool = False): tqdm_pbar.update((total_duration if finish else min(round(last_ts, 2), total_duration)) - tqdm_pbar.n) if progress_callback is not None: progress_callback(seek=tqdm_pbar.n, total=tqdm_pbar.total) def redo_words(_idx: int = None): nonlocal seg_words, seg_tokens, seg_words, words, word_tokens, curr_words, temp_word if curr_words and temp_word is not None: assert curr_words[0]['word'] == temp_word['word'] if curr_words[0]['probability'] >= temp_word['probability']: temp_word = curr_words[0] if _idx is None: # redo all words = seg_words + words word_tokens = seg_tokens + word_tokens curr_words = [] elif _idx != len(seg_words): # redo from _idx words = seg_words[_idx:] + words word_tokens = seg_tokens[_idx:] + word_tokens curr_words = curr_words[:_idx] if curr_words: if temp_word is not None: curr_words[0] = temp_word temp_word = None words = seg_words[_idx-1:_idx] + words word_tokens = seg_tokens[_idx-1:_idx] + word_tokens temp_word = curr_words.pop(-1) else: if temp_word is not None: curr_words[0] = temp_word temp_word = None n_samples = model.feature_extractor.n_samples if is_faster_model else N_SAMPLES temp_word = None while words and seek_sample < total_samples: time_offset = seek_sample / SAMPLE_RATE seek_sample_end = seek_sample + n_samples audio_segment = audio[seek_sample:seek_sample_end] segment_samples = audio_segment.shape[-1] if nonspeech_skip is not None: segment_nonspeech_timings = None if not vad: ts_token_mask = wav2mask(audio_segment, q_levels=q_levels, k_size=k_size) segment_nonspeech_timings = mask2timing(ts_token_mask, time_offset=time_offset) if segment_nonspeech_timings is not None: nonspeech_timings[0].extend(segment_nonspeech_timings[0]) nonspeech_timings[1].extend(segment_nonspeech_timings[1]) elif nonspeech_vad_timings: timing_indices = np.logical_and( nonspeech_vad_timings[1] > time_offset, nonspeech_vad_timings[0] < time_offset + 30.0 ) if timing_indices.any(): segment_nonspeech_timings = ( nonspeech_vad_timings[0][timing_indices], nonspeech_vad_timings[1][timing_indices] ) else: segment_nonspeech_timings = None if mn := timing_indices.argmax(): nonspeech_vad_timings = (nonspeech_vad_timings[0][mn:], nonspeech_vad_timings[1][mn:]) if segment_nonspeech_timings is not None: # segment has no detectable speech if ( (segment_nonspeech_timings[0][0] <= time_offset + min_word_dur) and (segment_nonspeech_timings[1][0] >= time_offset + segment_samples - min_word_dur) ): seek_sample += segment_samples continue timing_indices = (segment_nonspeech_timings[1] - segment_nonspeech_timings[0]) >= nonspeech_skip if any(timing_indices): nonspeech_starts = segment_nonspeech_timings[0][timing_indices] nonspeech_ends = segment_nonspeech_timings[1][timing_indices] if round(time_offset, 3) >= nonspeech_starts[0]: seek_sample = round(nonspeech_ends[0] * SAMPLE_RATE) if seek_sample + (min_word_dur * SAMPLE_RATE) >= total_samples: seek_sample = total_samples continue time_offset = seek_sample / SAMPLE_RATE if len(nonspeech_starts) > 1: seek_sample_end = ( seek_sample + round((nonspeech_starts[1] - nonspeech_ends[0]) * SAMPLE_RATE) ) audio_segment = audio[seek_sample:seek_sample_end] segment_samples = audio_segment.shape[-1] curr_words, curr_word_tokens = get_curr_words() segment = timestamp_words() curr_words = segment['words'] seg_words = [w['word'] for w in curr_words] seg_tokens = [w['tokens'] for w in curr_words] durations = np.array([w['end'] - w['start'] for w in curr_words]).round(3) nonzero_mask = durations > 0 nonzero_indices = np.flatnonzero(nonzero_mask) if len(nonzero_indices): redo_index = nonzero_indices[-1] + 1 if ( words and redo_index > 1 and curr_words[nonzero_indices[-1]]['end'] >= np.floor(time_offset + segment_samples / SAMPLE_RATE) ): nonzero_mask[nonzero_indices[-1]] = False nonzero_indices = nonzero_indices[:-1] redo_index = nonzero_indices[-1] + 1 med_dur = np.median(durations[:redo_index]) if fast_mode: new_start = None global_max_dur = None else: local_max_dur = round(med_dur * word_dur_factor, 3) if word_dur_factor else None if max_word_dur: local_max_dur = min(local_max_dur, max_word_dur) if local_max_dur else max_word_dur global_max_dur = max_word_dur else: global_max_dur = local_max_dur or None if global_max_dur and med_dur > global_max_dur: med_dur = global_max_dur if ( local_max_dur and durations[nonzero_indices[0]] > global_max_dur ): new_start = round(max( curr_words[nonzero_indices[0]]['end'] - (med_dur * nonzero_indices[0] + local_max_dur), curr_words[nonzero_indices[0]]['start'] ), 3) if new_start <= time_offset: new_start = None else: new_start = None if new_start is None: if global_max_dur: index_offset = nonzero_indices[0] + 1 redo_indices = \ np.flatnonzero(durations[index_offset:redo_index] > global_max_dur) + index_offset if len(redo_indices): redo_index = redo_indices[0] last_ts = curr_words[redo_index - 1]['end'] redo_words(redo_index) else: last_ts = new_start redo_words() seek_sample = round(last_ts * SAMPLE_RATE) else: seek_sample += audio_segment.shape[-1] last_ts = round(seek_sample / SAMPLE_RATE, 2) redo_words() update_pbar() result.extend(curr_words) if verbose: line = '\n'.join( f"[{format_timestamp(word['start'])}] -> " f"[{format_timestamp(word['end'])}] \"{word['word']}\"" for word in curr_words ) safe_print(line) update_pbar(True) if temp_word is not None: result.append(temp_word) if not result: warnings.warn('Failed to align text.', stacklevel=2) elif words: warnings.warn(f'Failed to align the last {len(words)}/{total_words} words after ' f'{format_timestamp(result[-1]["end"])}.', stacklevel=2) if words and not remove_instant_words: result.extend( [ dict(word=w, start=total_duration, end=total_duration, probability=0.0, tokens=wt) for w, wt in zip(words, word_tokens) ] ) if not result: return if len(split_indices_by_char): word_lens = np.cumsum([[len(w['word']) for w in result]]) split_indices = [(word_lens >= i).nonzero()[0][0]+1 for i in split_indices_by_char] result = WhisperResult([result[i:j] for i, j in zip([0]+split_indices[:-1], split_indices)]) else: result = WhisperResult([result]) if suppress_silence: result.suppress_silence( *nonspeech_timings, min_word_dur=min_word_dur, word_level=suppress_word_ts, nonspeech_error=nonspeech_error, use_word_position=use_word_position ) result.update_nonspeech_sections(*nonspeech_timings) if not original_split: result.regroup(regroup) if fail_segs := len([None for s in result.segments if s.end-s.start <= 0]): warnings.warn(f'{fail_segs}/{len(result.segments)} segments failed to align.', stacklevel=2) return result def refine( model: "Whisper", audio: Union[str, np.ndarray, torch.Tensor, bytes], result: WhisperResult, *, steps: str = None, rel_prob_decrease: float = .03, abs_prob_decrease: float = .05, rel_rel_prob_decrease: Optional[float] = None, prob_threshold: float = .5, rel_dur_change: Optional[float] = .5, abs_dur_change: Optional[float] = None, word_level: bool = True, precision: float = None, single_batch: bool = False, inplace: bool = True, demucs: Union[bool, torch.nn.Module] = False, demucs_options: dict = None, only_voice_freq: bool = False, verbose: Optional[bool] = False ) -> WhisperResult: """ Improve existing timestamps. This function iteratively muting portions of the audio and monitoring token probabilities to find the most precise timestamps. This "most precise" in this case means the latest start and earliest end of a word that maintains an acceptable probability determined by the specified arguments. This is useful readjusting timestamps when they start too early or end too late. Parameters ---------- model : "Whisper" The Whisper ASR model modified instance audio : str or numpy.ndarray or torch.Tensor or bytes Path/URL to the audio file, the audio waveform, or bytes of audio file. If audio is :class:`numpy.ndarray` or :class:`torch.Tensor`, the audio must be already at sampled to 16kHz. result : stable_whisper.result.WhisperResult All timestamps, words, probabilities, and other data from the transcription of ``audio``. steps : str, default 'se' Instructions for refinement. A 's' means refine start-timestamps. An 'e' means refine end-timestamps. rel_prob_decrease : float, default 0.3 Maximum percent decrease in probability relative to original probability which is the probability from muting according initial timestamps. abs_prob_decrease : float, default 0.05 Maximum decrease in probability from original probability. rel_rel_prob_decrease : float, optional Maximum percent decrease in probability relative to previous probability which is the probability from previous iteration of muting. prob_threshold : float, default 0.5 Stop refining the timestamp if the probability of its token goes below this value. rel_dur_change : float, default 0.5 Maximum percent change in duration of a word relative to its original duration. abs_dur_change : float, optional Maximum seconds a word is allowed deviate from its original duration. word_level : bool, default True Whether to refine timestamps on word-level. If ``False``, only refine start/end timestamps of each segment. precision : float, default 0.1 Precision of refined timestamps in seconds. The lowest precision is 0.02 second. single_batch : bool, default False Whether to process in only batch size of one to reduce memory usage. inplace : bool, default True, meaning return a deepcopy of ``result`` Whether to alter timestamps in-place. demucs : bool or torch.nn.Module, default False Whether to preprocess ``audio`` with Demucs to isolate vocals / remove noise. Set ``demucs`` to an instance of a Demucs model to avoid reloading the model for each run. Demucs must be installed to use. Official repo, https://github.com/facebookresearch/demucs. demucs_options : dict, optional Options to use for :func:`stable_whisper.audio.demucs_audio`. only_voice_freq : bool, default False Whether to only use sound between 200 - 5000 Hz, where majority of human speech are. verbose : bool or None, default False Whether to display the text being decoded to the console. Displays all the details if ``True``. Displays progressbar if ``False``. Display nothing if ``None``. Returns ------- stable_whisper.result.WhisperResult All timestamps, words, probabilities, and other data from the refinement of ``text`` with ``audio``. Notes ----- The lower the ``precision``, the longer the processing time. Examples -------- >>> import stable_whisper >>> model = stable_whisper.load_model('base') >>> result = model.transcribe('audio.mp3') >>> model.refine('audio.mp3', result) >>> result.to_srt_vtt('audio.srt') Saved 'audio.srt' """ if not steps: steps = 'se' if precision is None: precision = 0.1 if invalid_steps := steps.replace('s', '').replace('e', ''): raise ValueError(f'Invalid step(s): {", ".join(invalid_steps)}') if not result.has_words: raise NotImplementedError(f'Result must have word timestamps.') if not inplace: result = copy.deepcopy(result) audio = prep_audio( audio, demucs=demucs, demucs_options=demucs_options, only_voice_freq=only_voice_freq, verbose=verbose ) max_inference_tokens = model.dims.n_text_ctx - 6 sample_padding = int(N_FFT // 2) + 1 frame_precision = max(round(precision * FRAMES_PER_SECOND), 2) total_duration = round(audio.shape[-1] / SAMPLE_RATE, 3) tokenizer = get_tokenizer(model, language=result.language, task='transcribe') def ts_to_frames(timestamps: Union[np.ndarray, list]) -> np.ndarray: if isinstance(timestamps, list): timestamps = np.array(timestamps) return (timestamps * FRAMES_PER_SECOND).round().astype(int) def curr_segments(): all_words = result.all_words() seg_edge_mask = np.array([ 1 if _i == 0 else (2 if _i == len(seg.words)-1 else 0) for seg in result.segments for _i, w in enumerate(seg.words) ]) start_times = [ max( 0 if abs_dur_change is None else (w.start - abs_dur_change), 0 if rel_dur_change is None else (w.start - w.duration * rel_dur_change), 0 if i == 0 else max(all_words[i - 1].end, w.end - 14.5, 0) ) for i, w in enumerate(all_words) ] end_times = [ min( total_duration if abs_dur_change is None else (w.end + abs_dur_change), total_duration if rel_dur_change is None else (w.end + w.duration * rel_dur_change), total_duration if i == len(all_words) else min(all_words[i].start, w.start + 14.5, total_duration) ) for i, w in enumerate(all_words, 1) ] start = start_times[0] prev_i = 0 curr_words, curr_starts, curr_ends = [], [], [] for i, w in enumerate(all_words, 1): if ( (end_times[0] - start > 30) or (len(curr_words) + 1 > max_inference_tokens) ): if curr_words: yield curr_words, curr_starts, curr_ends, seg_edge_mask[prev_i:prev_i+len(curr_words)] curr_words, curr_starts, curr_ends = [], [], [] start = start_times[0] prev_i = i - 1 curr_words.append(w) curr_starts.append(start_times.pop(0)) curr_ends.append(end_times.pop(0)) if i == len(all_words): yield curr_words, curr_starts, curr_ends, seg_edge_mask[prev_i:prev_i+len(curr_words)] def _refine(_step: str): for words, min_starts, max_ends, edge_mask in curr_segments(): time_offset = min_starts[0] start_sample = round(time_offset * SAMPLE_RATE) end_sample = round(max_ends[-1] * SAMPLE_RATE) audio_segment = audio[start_sample:end_sample + 1].unsqueeze(0) max_starts = ts_to_frames(np.array([w.end for w in words]) - time_offset) min_ends = ts_to_frames(np.array([w.start for w in words]) - time_offset) min_starts = ts_to_frames(np.array(min_starts) - time_offset) max_ends = ts_to_frames(np.array(max_ends) - time_offset) mid_starts = min_starts + ((max_starts - min_starts) / 2).round().astype(int) mid_ends = min_ends + ((max_ends - min_ends) / 2).round().astype(int) text_tokens = [t for w in words for t in w.tokens if t < tokenizer.eot] word_tokens = [[t for t in w.tokens if t < tokenizer.eot] for w in words] orig_mel_segment = log_mel_spectrogram(audio_segment, model.dims.n_mels, padding=sample_padding) orig_mel_segment = pad_or_trim(orig_mel_segment, N_FRAMES).to(device=model.device) def get_prob(): tokens = torch.tensor( [ *tokenizer.sot_sequence, tokenizer.no_timestamps, *text_tokens, tokenizer.eot, ] ).to(model.device) with torch.no_grad(): curr_mel_segment = mel_segment if prob_indices else orig_mel_segment if single_batch: logits = torch.cat( [model(_mel.unsqueeze(0), tokens.unsqueeze(0)) for _mel in curr_mel_segment] ) else: logits = model(curr_mel_segment, tokens.unsqueeze(0)) sampled_logits = logits[:, len(tokenizer.sot_sequence):, : tokenizer.eot] token_probs = sampled_logits.softmax(dim=-1) text_token_probs = token_probs[:, np.arange(len(text_tokens)), text_tokens] token_positions = token_probs[:, np.arange(len(text_tokens))] if logits.shape[0] != 1 and prob_indices is not None: indices1 = np.arange(len(prob_indices)) text_token_probs = text_token_probs[prob_indices, indices1] token_positions = token_positions[prob_indices, indices1] else: text_token_probs.squeeze_(0) text_token_probs = text_token_probs.tolist() token_positions = \ ( token_positions.sort().indices == tokens[len(tokenizer.sot_sequence) + 1:-1][:, None] ).nonzero()[:, -1].tolist() word_boundaries = np.pad(np.cumsum([len(t) for t in word_tokens]), (1, 0)) word_probabilities = np.array([ text_token_probs[j-1] if is_end_ts else text_token_probs[i] for i, j in zip(word_boundaries[:-1], word_boundaries[1:]) ]) token_positions = [ token_positions[j-1] if is_end_ts else token_positions[i] for i, j in zip(word_boundaries[:-1], word_boundaries[1:]) ] return word_probabilities, token_positions def update_ts(): if not is_finish[idx] or changes[idx, -1] == -1: return new_ts = round(time_offset + (changes[idx, -1] / FRAMES_PER_SECOND), 3) if changes[idx, 0] and not changes[idx, 1]: if is_end_ts: if new_ts <= words[idx].end: return elif new_ts >= words[idx].start: return if not verbose: return curr_word = words[idx] word_info = (f'[Word="{curr_word.word}"] ' f'[Segment ID: {curr_word.segment_id}] ' f'[Word ID: {curr_word.id}]') if is_end_ts: print(f'End: {words[idx].end} -> {new_ts} {word_info}') words[idx].end = new_ts else: print(f'Start: {words[idx].start} -> {new_ts} {word_info}') words[idx].start = new_ts mel_segment = orig_mel_segment.clone().repeat_interleave(2, 0) is_end_ts = _step == 'e' prob_indices = [] is_finish = np.less([w.probability for w in words], prob_threshold) is_finish = np.logical_or(is_finish, [w.duration == 0 for w in words]) if not word_level: is_finish[edge_mask != (2 if is_end_ts else 1)] = True for idx, _i in enumerate(max_starts if is_end_ts else min_ends): row = idx % 2 prob_indices.extend([row] * len(words[idx].tokens)) if is_finish[idx]: continue if is_end_ts: _p = mel_segment.shape[-1] if idx == len(words)-1 else mid_ends[idx+1] mel_segment[row, :, _i:_p] = 0 else: _p = 0 if idx == 0 else mid_starts[idx-1] mel_segment[row, :, _p:_i] = 0 orig_probs, orig_tk_poss = get_prob() changes = np.zeros((orig_probs.shape[-1], 3), dtype=int) changes[:, -1] = -1 frame_indices = (mid_ends, max_starts) if is_end_ts else (min_ends, mid_starts) for idx, (_s, _e) in enumerate(zip(*frame_indices)): row = idx % 2 if is_finish[idx]: continue mel_segment[row, :, _s:_e] = 0 new_probs = prev_probs = orig_probs while not np.all(is_finish): probs, tk_poss = get_prob() abs_diffs = orig_probs - probs rel_diffs = abs_diffs / orig_probs rel_change_diffs = (prev_probs - probs) / prev_probs prev_probs = probs for idx, (abs_diff, rel_diff, rel_change_diff, prob) \ in enumerate(zip(abs_diffs, rel_diffs, rel_change_diffs, probs)): if is_finish[idx]: continue if is_end_ts: curr_min, curr_max, curr_mid = min_ends[idx], max_ends[idx], mid_ends[idx] else: curr_min, curr_max, curr_mid = min_starts[idx], max_starts[idx], mid_starts[idx] row = prob_indices[idx] best_tks_changed = orig_tk_poss[idx] > tk_poss[idx] failed_requirements = ( abs_diff > abs_prob_decrease or rel_diff > rel_prob_decrease or (rel_rel_prob_decrease is not None and rel_change_diff > rel_rel_prob_decrease) or prob < prob_threshold or best_tks_changed ) if failed_requirements: changes[idx][0] = 1 if is_end_ts: curr_min = curr_mid else: curr_max = curr_mid else: changes[idx][1] = 1 if is_end_ts: curr_max = curr_mid else: curr_min = curr_mid if (new_mid_change := round((curr_max - curr_min) / 2)) < frame_precision: is_finish[idx] = True update_ts() continue new_mid = curr_min + new_mid_change if failed_requirements: if is_end_ts: mel_segment[row, :, curr_min:new_mid] = orig_mel_segment[0, :, curr_min:new_mid] else: mel_segment[row, :, new_mid:curr_max] = orig_mel_segment[0, :, new_mid:curr_max] else: if is_end_ts: mel_segment[row, :, new_mid:curr_max] = 0 else: mel_segment[row, :, curr_min:new_mid] = 0 if is_end_ts: min_ends[idx], max_ends[idx], mid_ends[idx] = curr_min, curr_max, new_mid else: min_starts[idx], max_starts[idx], mid_starts[idx] = curr_min, curr_max, new_mid if not best_tks_changed: changes[idx][-1] = new_mid new_probs[idx] = prob update_pbar(words[-1].end) with tqdm(total=round(total_duration, 2), unit='sec', disable=verbose is not False, desc='Refine') as tqdm_pbar: def update_pbar(last_ts: float): nonlocal prev_ts tqdm_pbar.update(round(((last_ts - prev_ts) / len(steps)), 2)) prev_ts = last_ts for step_count, step in enumerate(steps, 1): prev_ts = 0 _refine(step) update_pbar(round(tqdm_pbar.total / len(step), 2)) tqdm_pbar.update(tqdm_pbar.total - tqdm_pbar.n) result.update_all_segs_with_words() return result def locate( model: "Whisper", audio: Union[str, np.ndarray, torch.Tensor, bytes], text: Union[str, List[int]], language: str, count: int = 1, duration_window: Union[float, Tuple[float, float]] = 3.0, *, mode: int = 0, start: float = None, end: float = None, probability_threshold: float = 0.5, eots: int = 1, max_token_per_seg: int = 20, exact_token: bool = False, case_sensitive: bool = False, verbose: bool = False, initial_prompt: str = None, suppress_tokens: Union[str, List[int]] = '-1', demucs: Union[bool, torch.nn.Module] = False, demucs_options: dict = None, only_voice_freq: bool = False, ) -> Union[List[Segment], List[dict]]: """ Locate when specific words are spoken in ``audio`` without fully transcribing. This is usefully for quickly finding at what time the specify words or phrases are spoken in an audio. Since it does not need to transcribe the audio to approximate the time, it is significantly faster transcribing then locating the word in the transcript. It can also transcribe few seconds around the approximated time to find out what was said around those words or confirm if the word was even spoken near that time. Parameters ---------- model : whisper.model.Whisper An instance of Whisper ASR model. audio : str or numpy.ndarray or torch.Tensor or bytes Path/URL to the audio file, the audio waveform, or bytes of audio file. If audio is :class:`numpy.ndarray` or :class:`torch.Tensor`, the audio must be already at sampled to 16kHz. text: str or list of int Words/phrase or list of tokens to search for in ``audio``. language : str Language of the ``text``. count : int, default 1, meaning stop search after 1 match Number of matches to find. Use 0 to look for all. duration_window : float or tuple of (float, float), default 3.0, same as (3.0, 3.0) Seconds before and after the end timestamp approximations to transcribe after mode 1. If tuple pair of values, then the 1st value will be seconds before the end and 2nd value will be seconds after. mode : int, default 0 Mode of search. 2, Approximates the end timestamp of ``text`` in the audio. This mode does not confirm whether ``text`` is spoken at the timestamp 1, Completes mode 2 then transcribes audio within ``duration_window`` to confirm whether `text` is a match at the approximated timestamp by checking if ``text`` at that ``duration_window`` is within ``probability_threshold`` or matching the string content if ``text`` with the transcribed text at the ``duration_window``. 0, Completes mode 1 then add word timestamps to the transcriptions of each match. Modes from fastest to slowest: 2, 1, 0 start : float, optional, meaning it starts from 0s Seconds into the audio to start searching for ``text``. end : float, optional Seconds into the audio to stop searching for ``text``. probability_threshold : float, default 0.5 Minimum probability of each token in ``text`` for it to be considered a match. eots : int, default 1 Number of EOTs to reach before stopping transcription at mode 1. When transcription reach a EOT, it usually means the end of the segment or audio. Once ``text`` is found in the ``duration_window``, the transcription will stop immediately upon reaching a EOT. max_token_per_seg : int, default 20 Maximum number of tokens to transcribe in the ``duration_window`` before stopping. exact_token : bool, default False Whether to find a match base on the exact tokens that make up ``text``. case_sensitive : bool, default False Whether to consider the case of ``text`` when matching in string content. verbose : bool or None, default False Whether to display the text being decoded to the console. Displays all the details if ``True``. Displays progressbar if ``False``. Display nothing if ``None``. initial_prompt : 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. suppress_tokens : str or list of int, default '-1', meaning suppress special characters except common punctuations List of tokens to suppress. demucs : bool or torch.nn.Module, default False Whether to preprocess ``audio`` with Demucs to isolate vocals / remove noise. Set ``demucs`` to an instance of a Demucs model to avoid reloading the model for each run. Demucs must be installed to use. Official repo, https://github.com/facebookresearch/demucs. demucs_options : dict, optional Options to use for :func:`stable_whisper.audio.demucs_audio`. only_voice_freq : bool, default False Whether to only use sound between 200 - 5000 Hz, where majority of human speech are. Returns ------- stable_whisper.result.Segment or list of dict or list of float Mode 0, list of instances of :class:`stable_whisper.result.Segment`. Mode 1, list of dictionaries with end timestamp approximation of matches and transcribed neighboring words. Mode 2, list of timestamps in seconds for each end timestamp approximation. Notes ----- For ``text``, the case and spacing matters as 'on', ' on', ' On' are different tokens, therefore chose the one that best suits the context (e.g. ' On' to look for it at the beginning of a sentence). Use a sufficiently large first value of ``duration_window`` i.e. the value > time it is expected to speak ``text``. If ``exact_token = False`` and the string content matches, then ``probability_threshold`` is not used. Examples -------- >>> import stable_whisper >>> model = stable_whisper.load_model('base') >>> matches = model.locate('audio.mp3', 'are', 'English', verbose=True) Some words can sound the same but have different spellings to increase of the chance of finding such words use ``initial_prompt``. >>> matches = model.locate('audio.mp3', ' Nickie', 'English', verbose=True, initial_prompt='Nickie') """ from whisper.timing import median_filter from whisper.decoding import DecodingTask, DecodingOptions, SuppressTokens from .timing import split_word_tokens sample_padding = int(N_FFT // 2) + 1 sec_per_emb = model.dims.n_audio_ctx / CHUNK_LENGTH CHUNK_SAMPLES = round(CHUNK_LENGTH * SAMPLE_RATE) if isinstance(duration_window, (float, int)): duration_window = [duration_window] * 2 window_sum = sum(duration_window) assert CHUNK_SAMPLES > window_sum, \ f'Sum of [duration_window] must be less than {CHUNK_SAMPLES}, got {window_sum}' adjusted_chunk_size = CHUNK_SAMPLES - round(duration_window[0]*SAMPLE_RATE) if initial_prompt: initial_prompt = ' ' + initial_prompt.strip() task = DecodingTask(model, DecodingOptions( language=language, prompt=initial_prompt, suppress_tokens=suppress_tokens, without_timestamps=True, )) tokenizer = task.tokenizer initial_tokens = list(task.initial_tokens) text_tokens, text = (tokenizer.encode(text), text) if isinstance(text, str) else (text, tokenizer.decode(text)) if not exact_token and not case_sensitive: text = text.lower() tk_suppress_masks = [ [i for i in fil.suppress_tokens if i < tokenizer.eot] for fil in task.logit_filters if isinstance(fil, SuppressTokens) ] audio = prep_audio( audio, demucs=demucs, demucs_options=demucs_options, only_voice_freq=only_voice_freq, verbose=verbose ) prev_target_end = None found = 0 if end: audio = audio[:round(end * SAMPLE_RATE)] seek_sample = round(start * SAMPLE_RATE) if start else 0 total_samples = audio.shape[-1] def _locate(): nonlocal seek_sample, found seek = round(seek_sample / SAMPLE_RATE, 3) audio_segment = audio[seek_sample: seek_sample + CHUNK_SAMPLES] mel_segment = log_mel_spectrogram(audio_segment, model.dims.n_mels, padding=sample_padding) mel_segment = pad_or_trim(mel_segment, N_FRAMES).to(device=model.device) QKs = [None] * model.dims.n_text_layer hooks = [ block.cross_attn.register_forward_hook( lambda _, ins, outs, index=i: QKs.__setitem__(index, outs[-1]) ) for i, block in enumerate(model.decoder.blocks) ] tokens = torch.tensor([initial_tokens + text_tokens]).to(model.device) with torch.no_grad(): audio_features = model.encoder(mel_segment.unsqueeze(0)) model.decoder(tokens, audio_features) for hook in hooks: hook.remove() weights = torch.cat([QKs[_l][:, _h] for _l, _h in model.alignment_heads.indices().T], dim=0) weights = weights.softmax(dim=-1) std, mean = torch.std_mean(weights, dim=-2, keepdim=True, unbiased=False) weights = (weights - mean) / std weights = median_filter(weights, 7) matrix = weights.mean(axis=0) target_end = round((matrix[-1].argmax()/sec_per_emb).item(), 3) found_msg = f'"{text}" ending at ~{format_timestamp(target_end+seek)}' if verbose else '' if mode == 2: if found_msg: safe_print('Unconfirmed:' + found_msg) nonlocal prev_target_end found += 1 if ( (seek_sample + CHUNK_SAMPLES >= total_samples) or (count and found >= count) or (prev_target_end == target_end) ): seek_sample = total_samples else: seek_sample += round(target_end * SAMPLE_RATE) prev_target_end = target_end return dict(tokens=[], target_end=target_end+seek) curr_start = round(max(target_end - duration_window[0], 0.), 3) curr_end = round(target_end + duration_window[1], 3) start_frame = round(curr_start * FRAMES_PER_SECOND) end_frame = round(curr_end * FRAMES_PER_SECOND) mel_segment_section = pad_or_trim(mel_segment[..., start_frame:end_frame], N_FRAMES) temp_tokens = torch.tensor([initial_tokens]).to(model.device) predictions = [] target_token_idx = 0 not_end = True found_target = False curr_eots = 0 temp_audio_features = model.encoder(mel_segment_section.unsqueeze(0)) tokens_to_decode = [] replace_found_tokens = [] infer_tokens = [temp_tokens[0]] kv_cache, hooks = model.install_kv_cache_hooks() while not_end: with torch.no_grad(): logits = model.decoder(temp_tokens, temp_audio_features, kv_cache=kv_cache)[0, -1, :tokenizer.eot+1] for tks in tk_suppress_masks: logits[tks] = -np.inf sorted_logits_idxs = logits.sort(dim=-1).indices[-2:] best_token = sorted_logits_idxs[-1] best_non_eot_token = sorted_logits_idxs[-2] if best_token == tokenizer.eot else best_token logits = logits[:tokenizer.eot].softmax(dim=-1) if found_target: target_word_prob = is_match = None else: if exact_token: is_match = False else: tokens_to_decode.append(best_non_eot_token) temp_text = tokenizer.decode(tokens_to_decode) if not case_sensitive: temp_text = temp_text.lower() if is_match := temp_text.endswith(text): tokens_to_decode = [] target_word_prob = logits[text_tokens[target_token_idx]].item() if ( target_word_prob is not None and ( target_word_prob >= probability_threshold or best_non_eot_token == text_tokens[target_token_idx] or is_match ) ): if is_match: best_token = best_non_eot_token token_prob = logits[best_token].item() found_target = True else: best_token[None] = text_tokens[target_token_idx] if len(replace_found_tokens) or best_non_eot_token != text_tokens[target_token_idx]: replace_found_tokens.append(best_non_eot_token) target_token_idx += 1 if target_token_idx == len(text_tokens): found_target = True token_prob = target_word_prob if found_target: found += 1 curr_eots = 0 else: if not found_target: if len(replace_found_tokens): temp_tokens = torch.cat(infer_tokens)[None] temp_tokens = torch.cat( [temp_tokens[..., :-len(replace_found_tokens)], torch.stack(replace_found_tokens)[None]] ) replace_found_tokens = [] kv_cache.clear() target_token_idx = 0 if best_token == tokenizer.eot: if curr_eots >= eots or found_target: not_end = False else: curr_eots += 1 best_token = best_non_eot_token else: curr_eots = 0 token_prob = None if best_token == tokenizer.eot else logits[best_token].item() predictions.append(dict(token=best_token.item(), prob=token_prob)) if len(predictions) > max_token_per_seg: not_end = False if not_end: infer_tokens.append(best_token[None]) temp_tokens = best_token[None, None] kv_cache.clear() for hook in hooks: hook.remove() segment = None if found_target: if found_msg: safe_print('Confirmed: ' + found_msg, tqdm_pbar.write) final_tokens = [p['token'] for p in predictions] if mode == 1: _, (ws, wts), _ = split_word_tokens([dict(tokens=final_tokens)], tokenizer) final_token_probs = [p['prob'] for p in predictions] wps = [float(np.mean([final_token_probs.pop(0) for _ in wt])) for wt in wts] words = [dict(word=w, tokens=wt, probability=wp) for w, wt, wp in zip(ws, wts, wps)] final_end = target_end+seek near_text = "".join(ws) segment = dict(end=final_end, text=text, duration_window_text=near_text, duration_window_word=words) if verbose: safe_print(f'Duration Window: "{near_text}"\n', tqdm_pbar.write) seek_sample += round(curr_end * SAMPLE_RATE) else: segment = dict( seek=0, tokens=final_tokens ) add_word_timestamps_stable( segments=[segment], model=model, tokenizer=tokenizer, mel=mel_segment, num_samples=round(curr_end*SAMPLE_RATE), gap_padding=None ) segment = Segment(0, 0, '', words=segment['words']) segment.update_seg_with_words() seek_sample += round(segment.words[-1].end * SAMPLE_RATE) segment.offset_time(seek) segment.seek = curr_start if verbose: safe_print(segment.to_display_str(), tqdm_pbar.write) else: seek_sample += adjusted_chunk_size if audio_segment.shape[-1] == CHUNK_SAMPLES else audio_segment.shape[-1] return segment total_duration = round(total_samples / SAMPLE_RATE, 2) matches = [] with tqdm(total=total_duration, unit='sec', disable=verbose is None, desc='Locate') as tqdm_pbar: while seek_sample < total_samples and (not count or found < count): if match := _locate(): matches.append(match) tqdm_pbar.update(round(seek_sample/SAMPLE_RATE, 2) - tqdm_pbar.n) tqdm_pbar.update(tqdm_pbar.total - tqdm_pbar.n) if verbose and not matches: safe_print(f'Failed to locate "{text}".') return matches