import warnings import re import torch import numpy as np from typing import Union, List, Tuple, Optional, Callable from dataclasses import dataclass from copy import deepcopy from itertools import chain from .stabilization import suppress_silence, get_vad_silence_func, mask2timing, wav2mask from .text_output import * from .utils import str_to_valid_type, format_timestamp, UnsortedException __all__ = ['WhisperResult', 'Segment'] def _combine_attr(obj: object, other_obj: object, attr: str): if (val := getattr(obj, attr)) is not None: other_val = getattr(other_obj, attr) if isinstance(val, list): if other_val is None: setattr(obj, attr, None) else: val.extend(other_val) else: new_val = None if other_val is None else ((val + other_val) / 2) setattr(obj, attr, new_val) def _increment_attr(obj: object, attr: str, val: Union[int, float]): if (curr_val := getattr(obj, attr, None)) is not None: setattr(obj, attr, curr_val + val) @dataclass class WordTiming: word: str start: float end: float probability: float = None tokens: List[int] = None left_locked: bool = False right_locked: bool = False segment_id: Optional[int] = None id: Optional[int] = None def __len__(self): return len(self.word) def __add__(self, other: 'WordTiming'): self_copy = deepcopy(self) self_copy.start = min(self_copy.start, other.start) self_copy.end = max(other.end, self_copy.end) self_copy.word += other.word self_copy.left_locked = self_copy.left_locked or other.left_locked self_copy.right_locked = self_copy.right_locked or other.right_locked _combine_attr(self_copy, other, 'probability') _combine_attr(self_copy, other, 'tokens') return self_copy def __deepcopy__(self, memo=None): return self.copy() def copy(self): return WordTiming( word=self.word, start=self.start, end=self.end, probability=self.probability, tokens=None if self.tokens is None else self.tokens.copy(), left_locked=self.left_locked, right_locked=self.right_locked, segment_id=self.segment_id, id=self.id ) @property def duration(self): return round(self.end - self.start, 3) def round_all_timestamps(self): self.start = round(self.start, 3) self.end = round(self.end, 3) def offset_time(self, offset_seconds: float): self.start = round(self.start + offset_seconds, 3) self.end = round(self.end + offset_seconds, 3) def to_dict(self): dict_ = deepcopy(self).__dict__ dict_.pop('left_locked') dict_.pop('right_locked') return dict_ def lock_left(self): self.left_locked = True def lock_right(self): self.right_locked = True def lock_both(self): self.lock_left() self.lock_right() def unlock_both(self): self.left_locked = False self.right_locked = False def suppress_silence(self, silent_starts: np.ndarray, silent_ends: np.ndarray, min_word_dur: float = 0.1, nonspeech_error: float = 0.3, keep_end: Optional[bool] = True): suppress_silence(self, silent_starts, silent_ends, min_word_dur, nonspeech_error, keep_end) return self def rescale_time(self, scale_factor: float): self.start = round(self.start * scale_factor, 3) self.end = round(self.end * scale_factor, 3) def clamp_max(self, max_dur: float, clip_start: bool = False, verbose: bool = False): if self.duration > max_dur: if clip_start: new_start = round(self.end - max_dur, 3) if verbose: print(f'Start: {self.start} -> {new_start}\nEnd: {self.end}\nText:"{self.word}"\n') self.start = new_start else: new_end = round(self.start + max_dur, 3) if verbose: print(f'Start: {self.start}\nEnd: {self.end} -> {new_end}\nText:"{self.word}"\n') self.end = new_end def set_segment(self, segment: 'Segment'): self._segment = segment def get_segment(self) -> Union['Segment', None]: """ Return instance of :class:`stable_whisper.result.Segment` that this instance is a part of. """ return getattr(self, '_segment', None) def _words_by_lock(words: List[WordTiming], only_text: bool = False, include_single: bool = False): """ Return a nested list of words such that each sublist contains words that are locked together. """ all_words = [] for word in words: if len(all_words) == 0 or not (all_words[-1][-1].right_locked or word.left_locked): all_words.append([word]) else: all_words[-1].append(word) if only_text: all_words = list(map(lambda ws: list(map(lambda w: w.word, ws)), all_words)) if not include_single: all_words = [ws for ws in all_words if len(ws) > 1] return all_words @dataclass class Segment: start: float end: float text: str seek: float = None tokens: List[int] = None temperature: float = None avg_logprob: float = None compression_ratio: float = None no_speech_prob: float = None words: Union[List[WordTiming], List[dict]] = None ori_has_words: bool = None id: int = None def __getitem__(self, index: int) -> WordTiming: if self.words is None: raise ValueError('segment contains no words') return self.words[index] def __delitem__(self, index: int): if self.words is None: raise ValueError('segment contains no words') del self.words[index] self.reassign_ids() self.update_seg_with_words() def __deepcopy__(self, memo=None): return self.copy() def copy(self, new_words: Optional[List[WordTiming]] = None): if new_words is None: words = None if self.words is None else [w.copy() for w in self.words] else: words = [w.copy() for w in new_words] new_seg = Segment( start=self.start, end=self.end, text=self.text, seek=self.seek, tokens=self.tokens, temperature=self.temperature, avg_logprob=self.avg_logprob, compression_ratio=self.compression_ratio, no_speech_prob=self.no_speech_prob, words=words, id=self.id ) new_seg.update_seg_with_words() return new_seg def to_display_str(self, only_segment: bool = False): line = f'[{format_timestamp(self.start)} --> {format_timestamp(self.end)}] "{self.text}"' if self.has_words and not only_segment: line += '\n' + '\n'.join( f"-[{format_timestamp(w.start)}] -> [{format_timestamp(w.end)}] \"{w.word}\"" for w in self.words ) + '\n' return line @property def has_words(self): return bool(self.words) @property def duration(self): return self.end - self.start def word_count(self): if self.has_words: return len(self.words) return -1 def char_count(self): if self.has_words: return sum(len(w) for w in self.words) return len(self.text) def __post_init__(self): if self.has_words: self.words: List[WordTiming] = \ [WordTiming(**word) if isinstance(word, dict) else word for word in self.words] for w in self.words: w.set_segment(self) if self.ori_has_words is None: self.ori_has_words = self.has_words self.round_all_timestamps() def __add__(self, other: 'Segment'): self_copy = deepcopy(self) self_copy.start = min(self_copy.start, other.start) self_copy.end = max(other.end, self_copy.end) self_copy.text += other.text _combine_attr(self_copy, other, 'tokens') _combine_attr(self_copy, other, 'temperature') _combine_attr(self_copy, other, 'avg_logprob') _combine_attr(self_copy, other, 'compression_ratio') _combine_attr(self_copy, other, 'no_speech_prob') if self_copy.has_words: if other.has_words: self_copy.words.extend(other.words) else: self_copy.words = None return self_copy def _word_operations(self, operation: str, *args, **kwargs): if self.has_words: for w in self.words: getattr(w, operation)(*args, **kwargs) def round_all_timestamps(self): self.start = round(self.start, 3) self.end = round(self.end, 3) if self.has_words: for word in self.words: word.round_all_timestamps() def offset_time(self, offset_seconds: float): self.start = round(self.start + offset_seconds, 3) self.end = round(self.end + offset_seconds, 3) _increment_attr(self, 'seek', offset_seconds) self._word_operations('offset_time', offset_seconds) def add_words(self, index0: int, index1: int, inplace: bool = False): if self.has_words: new_word = self.words[index0] + self.words[index1] if inplace: i0, i1 = sorted([index0, index1]) self.words[i0] = new_word del self.words[i1] return new_word def rescale_time(self, scale_factor: float): self.start = round(self.start * scale_factor, 3) self.end = round(self.end * scale_factor, 3) if self.seek is not None: self.seek = round(self.seek * scale_factor, 3) self._word_operations('rescale_time', scale_factor) self.update_seg_with_words() def apply_min_dur(self, min_dur: float, inplace: bool = False): """ Merge any word with adjacent word if its duration is less than ``min_dur``. """ segment = self if inplace else deepcopy(self) if not self.has_words: return segment max_i = len(segment.words) - 1 if max_i == 0: return segment for i in reversed(range(len(segment.words))): if max_i == 0: break if segment.words[i].duration < min_dur: if i == max_i: segment.add_words(i-1, i, inplace=True) elif i == 0: segment.add_words(i, i+1, inplace=True) else: if segment.words[i+1].duration < segment.words[i-1].duration: segment.add_words(i-1, i, inplace=True) else: segment.add_words(i, i+1, inplace=True) max_i -= 1 return segment def _to_reverse_text( self, prepend_punctuations: str = None, append_punctuations: str = None ): """ Return a copy with words reversed order per segment. """ if prepend_punctuations is None: prepend_punctuations = "\"'“¿([{-" if prepend_punctuations and ' ' not in prepend_punctuations: prepend_punctuations += ' ' if append_punctuations is None: append_punctuations = "\"'.。,，!！?？:：”)]}、" self_copy = deepcopy(self) has_prepend = bool(prepend_punctuations) has_append = bool(append_punctuations) if has_prepend or has_append: word_objs = ( self_copy.words if self_copy.has_words else [WordTiming(w, 0, 1, 0) for w in self_copy.text.split(' ')] ) for word in word_objs: new_append = '' if has_prepend: for _ in range(len(word)): char = word.word[0] if char in prepend_punctuations: new_append += char word.word = word.word[1:] else: break new_prepend = '' if has_append: for _ in range(len(word)): char = word.word[-1] if char in append_punctuations: new_prepend += char word.word = word.word[:-1] else: break word.word = f'{new_prepend}{word.word}{new_append[::-1]}' self_copy.text = ''.join(w.word for w in reversed(word_objs)) return self_copy def to_dict(self, reverse_text: Union[bool, tuple] = False): if reverse_text: seg_dict = ( (self._to_reverse_text(*reverse_text) if isinstance(reverse_text, tuple) else self._to_reverse_text()).__dict__ ) else: seg_dict = deepcopy(self).__dict__ seg_dict.pop('ori_has_words') if self.has_words: seg_dict['words'] = [w.to_dict() for w in seg_dict['words']] elif self.ori_has_words: seg_dict['words'] = [] else: seg_dict.pop('words') if self.id is None: seg_dict.pop('id') if reverse_text: seg_dict['reversed_text'] = True return seg_dict def words_by_lock(self, only_text: bool = True, include_single: bool = False): return _words_by_lock(self.words, only_text=only_text, include_single=include_single) @property def left_locked(self): if self.has_words: return self.words[0].left_locked return False @property def right_locked(self): if self.has_words: return self.words[-1].right_locked return False def lock_left(self): if self.has_words: self.words[0].lock_left() def lock_right(self): if self.has_words: self.words[-1].lock_right() def lock_both(self): self.lock_left() self.lock_right() def unlock_all_words(self): self._word_operations('unlock_both') def reassign_ids(self): if self.has_words: for i, w in enumerate(self.words): w.segment_id = self.id w.id = i def update_seg_with_words(self): if self.has_words: self.start = self.words[0].start self.end = self.words[-1].end self.text = ''.join(w.word for w in self.words) self.tokens = ( None if any(w.tokens is None for w in self.words) else [t for w in self.words for t in w.tokens] ) for w in self.words: w.set_segment(self) def suppress_silence(self, silent_starts: np.ndarray, silent_ends: np.ndarray, min_word_dur: float = 0.1, word_level: bool = True, nonspeech_error: float = 0.3, use_word_position: bool = True): if self.has_words: words = self.words if word_level or len(self.words) == 1 else [self.words[0], self.words[-1]] for i, w in enumerate(words, 1): if use_word_position: keep_end = True if i == 1 else (False if i == len(words) else None) else: keep_end = None w.suppress_silence(silent_starts, silent_ends, min_word_dur, nonspeech_error, keep_end) self.update_seg_with_words() else: suppress_silence(self, silent_starts, silent_ends, min_word_dur, nonspeech_error) return self def get_locked_indices(self): locked_indices = [i for i, (left, right) in enumerate(zip(self.words[1:], self.words[:-1])) if left.left_locked or right.right_locked] return locked_indices def get_gaps(self, as_ndarray=False): if self.has_words: s_ts = np.array([w.start for w in self.words]) e_ts = np.array([w.end for w in self.words]) gap = s_ts[1:] - e_ts[:-1] return gap if as_ndarray else gap.tolist() return [] def get_gap_indices(self, max_gap: float = 0.1): # for splitting if not self.has_words or len(self.words) < 2: return [] if max_gap is None: max_gap = 0 indices = (self.get_gaps(True) > max_gap).nonzero()[0].tolist() return sorted(set(indices) - set(self.get_locked_indices())) def get_punctuation_indices(self, punctuation: Union[List[str], List[Tuple[str, str]], str]): # for splitting if not self.has_words or len(self.words) < 2: return [] if isinstance(punctuation, str): punctuation = [punctuation] indices = [] for p in punctuation: if isinstance(p, str): for i, s in enumerate(self.words[:-1]): if s.word.endswith(p): indices.append(i) elif i != 0 and s.word.startswith(p): indices.append(i-1) else: ending, beginning = p indices.extend([i for i, (w0, w1) in enumerate(zip(self.words[:-1], self.words[1:])) if w0.word.endswith(ending) and w1.word.startswith(beginning)]) return sorted(set(indices) - set(self.get_locked_indices())) def get_length_indices(self, max_chars: int = None, max_words: int = None, even_split: bool = True, include_lock: bool = False): # for splitting if not self.has_words or (max_chars is None and max_words is None): return [] assert max_chars != 0 and max_words != 0, \ f'max_chars and max_words must be greater 0, but got {max_chars} and {max_words}' if len(self.words) < 2: return [] indices = [] if even_split: char_count = -1 if max_chars is None else sum(map(len, self.words)) word_count = -1 if max_words is None else len(self.words) exceed_chars = max_chars is not None and char_count > max_chars exceed_words = max_words is not None and word_count > max_words if exceed_chars: splits = np.ceil(char_count / max_chars) chars_per_split = char_count / splits cum_char_count = np.cumsum([len(w.word) for w in self.words[:-1]]) indices = [ (np.abs(cum_char_count-(i*chars_per_split))).argmin() for i in range(1, int(splits)) ] if max_words is not None: exceed_words = any(j-i+1 > max_words for i, j in zip([0]+indices, indices+[len(self.words)])) if exceed_words: splits = np.ceil(word_count / max_words) words_per_split = word_count / splits cum_word_count = np.array(range(1, len(self.words)+1)) indices = [ np.abs(cum_word_count-(i*words_per_split)).argmin() for i in range(1, int(splits)) ] else: curr_words = 0 curr_chars = 0 locked_indices = [] if include_lock: locked_indices = self.get_locked_indices() for i, word in enumerate(self.words): curr_words += 1 curr_chars += len(word) if i != 0: if ( max_chars is not None and curr_chars > max_chars or max_words is not None and curr_words > max_words ) and i-1 not in locked_indices: indices.append(i-1) curr_words = 1 curr_chars = len(word) return indices def get_duration_indices(self, max_dur: float, even_split: bool = True, include_lock: bool = False): if not self.has_words or (total_duration := np.sum([w.duration for w in self.words])) <= max_dur: return [] if even_split: splits = np.ceil(total_duration / max_dur) dur_per_split = total_duration / splits cum_dur = np.cumsum([w.duration for w in self.words[:-1]]) indices = [ (np.abs(cum_dur - (i * dur_per_split))).argmin() for i in range(1, int(splits)) ] else: indices = [] curr_total_dur = 0.0 locked_indices = self.get_locked_indices() if include_lock else [] for i, word in enumerate(self.words): curr_total_dur += word.duration if i != 0: if curr_total_dur > max_dur and i - 1 not in locked_indices: indices.append(i - 1) curr_total_dur = word.duration return indices def split(self, indices: List[int]): if len(indices) == 0: return [] if indices[-1] != len(self.words) - 1: indices.append(len(self.words) - 1) seg_copies = [] prev_i = 0 for i in indices: i += 1 c = deepcopy(self) c.words = c.words[prev_i:i] c.update_seg_with_words() seg_copies.append(c) prev_i = i return seg_copies def set_result(self, result: 'WhisperResult'): self._result = result def get_result(self) -> Union['WhisperResult', None]: """ Return outer instance of :class:`stable_whisper.result.WhisperResult` that ``self`` is a part of. """ return getattr(self, '_result', None) class WhisperResult: def __init__( self, result: Union[str, dict, list], force_order: bool = False, check_sorted: Union[bool, str] = True, show_unsorted: bool = True ): result, self.path = self._standardize_result(result) self.ori_dict = result.get('ori_dict') or result self.language = self.ori_dict.get('language') self._regroup_history = result.get('regroup_history', '') self._nonspeech_sections = result.get('nonspeech_sections', []) segments = deepcopy(result.get('segments', self.ori_dict.get('segments'))) self.segments: List[Segment] = [Segment(**s) for s in segments] if segments else [] self._forced_order = force_order if self._forced_order: self.force_order() self.raise_for_unsorted(check_sorted, show_unsorted) self.remove_no_word_segments(any(seg.has_words for seg in self.segments)) self.update_all_segs_with_words() def __getitem__(self, index: int) -> Segment: return self.segments[index] def __delitem__(self, index: int): del self.segments[index] self.reassign_ids(True) @staticmethod def _standardize_result(result: Union[str, dict, list]): path = None if isinstance(result, str): path = result result = load_result(path) if isinstance(result, list): if isinstance(result[0], list): if not isinstance(result[0][0], dict): raise NotImplementedError(f'Got list of list of {type(result[0])} but expects list of list of dict') result = dict( segments=[ dict( start=words[0]['start'], end=words[-1]['end'], text=''.join(w['word'] for w in words), words=words ) for words in result ] ) elif isinstance(result[0], dict): result = dict(segments=result) else: raise NotImplementedError(f'Got list of {type(result[0])} but expects list of list/dict') return result, path def force_order(self): prev_ts_end = 0 timestamps = self.all_words_or_segments() for i, ts in enumerate(timestamps, 1): if ts.start < prev_ts_end: ts.start = prev_ts_end if ts.start > ts.end: if prev_ts_end > ts.end: warnings.warn('Multiple consecutive timestamps are out of order. Some parts will have no duration.') ts.start = ts.end for j in range(i-2, -1, -1): if timestamps[j].end > ts.end: timestamps[j].end = ts.end if timestamps[j].start > ts.end: timestamps[j].start = ts.end else: if ts.start != prev_ts_end: ts.start = prev_ts_end else: ts.end = ts.start if i == len(timestamps) else timestamps[i].start prev_ts_end = ts.end if self.has_words: self.update_all_segs_with_words() def raise_for_unsorted(self, check_sorted: Union[bool, str] = True, show_unsorted: bool = True): if check_sorted is False: return all_parts = self.all_words_or_segments() has_words = self.has_words timestamps = np.array(list(chain.from_iterable((p.start, p.end) for p in all_parts))) if len(timestamps) > 1 and (unsorted_mask := timestamps[:-1] > timestamps[1:]).any(): if show_unsorted: def get_part_info(idx): curr_part = all_parts[idx] seg_id = curr_part.segment_id if has_words else curr_part.id word_id_str = f'Word ID: {curr_part.id}\n' if has_words else '' return ( f'Segment ID: {seg_id}\n{word_id_str}' f'Start: {curr_part.start}\nEnd: {curr_part.end}\n' f'Text: "{curr_part.word if has_words else curr_part.text}"' ), curr_part.start, curr_part.end for i, unsorted in enumerate(unsorted_mask, 2): if unsorted: word_id = i//2-1 part_info, start, end = get_part_info(word_id) if i % 2 == 1: next_info, next_start, _ = get_part_info(word_id+1) part_info += f'\nConflict: end ({end}) > next start ({next_start})\n{next_info}' else: part_info += f'\nConflict: start ({start}) > end ({end})' print(part_info, end='\n\n') data = self.to_dict() if check_sorted is True: raise UnsortedException(data=data) warnings.warn('Timestamps are not in ascending order. ' 'If data is produced by Stable-ts, please submit an issue with the saved data.') save_as_json(data, check_sorted) def update_all_segs_with_words(self): for seg in self.segments: seg.update_seg_with_words() seg.set_result(self) def update_nonspeech_sections(self, silent_starts, silent_ends): self._nonspeech_sections = [dict(start=s, end=e) for s, e in zip(silent_starts, silent_ends)] def add_segments(self, index0: int, index1: int, inplace: bool = False, lock: bool = False): new_seg = self.segments[index0] + self.segments[index1] new_seg.update_seg_with_words() if lock and self.segments[index0].has_words: lock_idx = len(self.segments[index0].words) new_seg.words[lock_idx - 1].lock_right() if lock_idx < len(new_seg.words): new_seg.words[lock_idx].lock_left() if inplace: i0, i1 = sorted([index0, index1]) self.segments[i0] = new_seg del self.segments[i1] return new_seg def rescale_time(self, scale_factor: float): for s in self.segments: s.rescale_time(scale_factor) def apply_min_dur(self, min_dur: float, inplace: bool = False): """ Merge any word/segment with adjacent word/segment if its duration is less than ``min_dur``. """ result = self if inplace else deepcopy(self) max_i = len(result.segments) - 1 if max_i == 0: return result for i in reversed(range(len(result.segments))): if max_i == 0: break if result.segments[i].duration < min_dur: if i == max_i: result.add_segments(i-1, i, inplace=True) elif i == 0: result.add_segments(i, i+1, inplace=True) else: if result.segments[i+1].duration < result.segments[i-1].duration: result.add_segments(i-1, i, inplace=True) else: result.add_segments(i, i+1, inplace=True) max_i -= 1 result.reassign_ids() for s in result.segments: s.apply_min_dur(min_dur, inplace=True) return result def offset_time(self, offset_seconds: float): for s in self.segments: s.offset_time(offset_seconds) def suppress_silence( self, silent_starts: np.ndarray, silent_ends: np.ndarray, min_word_dur: float = 0.1, word_level: bool = True, nonspeech_error: float = 0.3, use_word_position: bool = True ) -> "WhisperResult": """ Move any start/end timestamps in silence parts of audio to the boundaries of the silence. Parameters ---------- silent_starts : numpy.ndarray An array starting timestamps of silent sections of audio. silent_ends : numpy.ndarray An array ending timestamps of silent sections of audio. min_word_dur : float, default 0.1 Shortest duration each word is allowed to reach for adjustments. word_level : bool, default False Whether to settings to word level timestamps. nonspeech_error : float, default 0.3 Relative error of non-speech sections that appear in between a word for adjustments. 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. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ for s in self.segments: s.suppress_silence( silent_starts, silent_ends, min_word_dur, word_level=word_level, nonspeech_error=nonspeech_error, use_word_position=use_word_position ) return self def adjust_by_silence( self, audio: Union[torch.Tensor, np.ndarray, str, bytes], vad: bool = False, *, verbose: (bool, None) = False, sample_rate: int = None, vad_onnx: bool = False, vad_threshold: float = 0.35, q_levels: int = 20, k_size: int = 5, min_word_dur: float = 0.1, word_level: bool = True, nonspeech_error: float = 0.3, use_word_position: bool = True ) -> "WhisperResult": """ Adjust timestamps base detected speech gaps. This is method combines :meth:`stable_whisper.result.WhisperResult.suppress_silence` with silence detection. Parameters ---------- audio : str or numpy.ndarray or torch.Tensor or bytes Path/URL to the audio file, the audio waveform, or bytes of audio file. 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. verbose : bool or None, default False If ``False``, mute messages about hitting local caches. Note that the message about first download cannot be muted. Only applies if ``vad = True``. sample_rate : int, default None, meaning ``whisper.audio.SAMPLE_RATE``, 16kHZ The sample rate of ``audio``. vad_onnx : bool, default False Whether to use ONNX for Silero VAD. vad_threshold : float, default 0.35 Threshold for detecting speech with Silero VAD. Low threshold reduces false positives for silence detection. 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. min_word_dur : float, default 0.1 Shortest duration each word is allowed to reach from adjustments. word_level : bool, default False Whether to settings to word level timestamps. nonspeech_error : float, default 0.3 Relative error of non-speech sections that appear in between a word for adjustments. 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. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. Notes ----- This operation is already performed by :func:`stable_whisper.whisper_word_level.transcribe_stable` / :func:`stable_whisper.whisper_word_level.transcribe_minimal`/ :func:`stable_whisper.non_whisper.transcribe_any` / :func:`stable_whisper.alignment.align` if ``suppress_silence = True``. """ if vad: silent_timings = get_vad_silence_func( onnx=vad_onnx, verbose=verbose )(audio, speech_threshold=vad_threshold, sr=sample_rate) else: silent_timings = mask2timing( wav2mask(audio, q_levels=q_levels, k_size=k_size, sr=sample_rate) ) if silent_timings is None: return self self.suppress_silence( *silent_timings, min_word_dur=min_word_dur, word_level=word_level, nonspeech_error=nonspeech_error, use_word_position=use_word_position ) self.update_nonspeech_sections(*silent_timings) return self def adjust_by_result( self, other_result: "WhisperResult", min_word_dur: float = 0.1, verbose: bool = False ): """ Minimize the duration of words using timestamps of another result. Parameters ---------- other_result : "WhisperResult" Timing data of the same words in a WhisperResult instance. min_word_dur : float, default 0.1 Prevent changes to timestamps if the resultant word duration is less than ``min_word_dur``. verbose : bool, default False Whether to print out the timestamp changes. """ if not (self.has_words and other_result.has_words): raise NotImplementedError('This operation can only be performed on results with word timestamps') assert [w.word for w in self.all_words()] == [w.word for w in other_result.all_words()], \ 'The words in [other_result] do not match the current words.' for word, other_word in zip(self.all_words(), other_result.all_words()): if word.end > other_word.start: new_start = max(word.start, other_word.start) new_end = min(word.end, other_word.end) if new_end - new_start >= min_word_dur: line = '' if word.start != new_start: if verbose: line += f'[Start:{word.start:.3f}->{new_start:.3f}] ' word.start = new_start if word.end != new_end: if verbose: line += f'[End:{word.end:.3f}->{new_end:.3f}] ' word.end = new_end if line: print(f'{line}"{word.word}"') self.update_all_segs_with_words() def reassign_ids(self, only_segments: bool = False): for i, s in enumerate(self.segments): s.id = i if not only_segments: s.reassign_ids() def remove_no_word_segments(self, ignore_ori=False): for i in reversed(range(len(self.segments))): if (ignore_ori or self.segments[i].ori_has_words) and not self.segments[i].has_words: del self.segments[i] self.reassign_ids() def get_locked_indices(self): locked_indices = [i for i, (left, right) in enumerate(zip(self.segments[1:], self.segments[:-1])) if left.left_locked or right.right_locked] return locked_indices def get_gaps(self, as_ndarray=False): s_ts = np.array([s.start for s in self.segments]) e_ts = np.array([s.end for s in self.segments]) gap = s_ts[1:] - e_ts[:-1] return gap if as_ndarray else gap.tolist() def get_gap_indices(self, min_gap: float = 0.1): # for merging if len(self.segments) < 2: return [] if min_gap is None: min_gap = 0 indices = (self.get_gaps(True) <= min_gap).nonzero()[0].tolist() return sorted(set(indices) - set(self.get_locked_indices())) def get_punctuation_indices(self, punctuation: Union[List[str], List[Tuple[str, str]], str]): # for merging if len(self.segments) < 2: return [] if isinstance(punctuation, str): punctuation = [punctuation] indices = [] for p in punctuation: if isinstance(p, str): for i, s in enumerate(self.segments[:-1]): if s.text.endswith(p): indices.append(i) elif i != 0 and s.text.startswith(p): indices.append(i-1) else: ending, beginning = p indices.extend([i for i, (s0, s1) in enumerate(zip(self.segments[:-1], self.segments[1:])) if s0.text.endswith(ending) and s1.text.startswith(beginning)]) return sorted(set(indices) - set(self.get_locked_indices())) def all_words(self): return list(chain.from_iterable(s.words for s in self.segments)) def all_words_or_segments(self): return self.all_words() if self.has_words else self.segments def all_words_by_lock(self, only_text: bool = True, by_segment: bool = False, include_single: bool = False): if by_segment: return [ segment.words_by_lock(only_text=only_text, include_single=include_single) for segment in self.segments ] return _words_by_lock(self.all_words(), only_text=only_text, include_single=include_single) def all_tokens(self): return list(chain.from_iterable(s.tokens for s in self.all_words())) def to_dict(self): return dict(text=self.text, segments=self.segments_to_dicts(), language=self.language, ori_dict=self.ori_dict, regroup_history=self._regroup_history, nonspeech_sections=self._nonspeech_sections) def segments_to_dicts(self, reverse_text: Union[bool, tuple] = False): return [s.to_dict(reverse_text=reverse_text) for s in self.segments] def _split_segments(self, get_indices, args: list = None, *, lock: bool = False, newline: bool = False): if args is None: args = [] no_words = False for i in reversed(range(0, len(self.segments))): no_words = no_words or not self.segments[i].has_words indices = sorted(set(get_indices(self.segments[i], *args))) if not indices: continue if newline: if indices[-1] == len(self.segments[i].words) - 1: del indices[-1] if not indices: continue for word_idx in indices: if self.segments[i].words[word_idx].word.endswith('\n'): continue self.segments[i].words[word_idx].word += '\n' if lock: self.segments[i].words[word_idx].lock_right() if word_idx + 1 < len(self.segments[i].words): self.segments[i].words[word_idx+1].lock_left() self.segments[i].update_seg_with_words() else: new_segments = self.segments[i].split(indices) if lock: for s in new_segments: if s == new_segments[0]: s.lock_right() elif s == new_segments[-1]: s.lock_left() else: s.lock_both() del self.segments[i] for s in reversed(new_segments): self.segments.insert(i, s) if no_words: warnings.warn('Found segment(s) without word timings. These segment(s) cannot be split.') self.remove_no_word_segments() def _merge_segments(self, indices: List[int], *, max_words: int = None, max_chars: int = None, is_sum_max: bool = False, lock: bool = False): if len(indices) == 0: return for i in reversed(indices): seg = self.segments[i] if ( ( max_words and seg.has_words and ( (seg.word_count() + self.segments[i + 1].word_count() > max_words) if is_sum_max else (seg.word_count() > max_words and self.segments[i + 1].word_count() > max_words) ) ) or ( max_chars and ( (seg.char_count() + self.segments[i + 1].char_count() > max_chars) if is_sum_max else (seg.char_count() > max_chars and self.segments[i + 1].char_count() > max_chars) ) ) ): continue self.add_segments(i, i + 1, inplace=True, lock=lock) self.remove_no_word_segments() def get_content_by_time( self, time: Union[float, Tuple[float, float], dict], within: bool = False, segment_level: bool = False ) -> Union[List[WordTiming], List[Segment]]: """ Return content in the ``time`` range. Parameters ---------- time : float or tuple of (float, float) or dict Range of time to find content. For tuple of two floats, first value is the start time and second value is the end time. For a single float value, it is treated as both the start and end time. within : bool, default False Whether to only find content fully overlaps with ``time`` range. segment_level : bool, default False Whether to look only on the segment level and return instances of :class:`stable_whisper.result.Segment` instead of :class:`stable_whisper.result.WordTiming`. Returns ------- list of stable_whisper.result.WordTiming or list of stable_whisper.result.Segment List of contents in the ``time`` range. The contents are instances of :class:`stable_whisper.result.Segment` if ``segment_level = True`` else :class:`stable_whisper.result.WordTiming`. """ if not segment_level and not self.has_words: raise ValueError('Missing word timestamps in result. Use ``segment_level=True`` instead.') contents = self.segments if segment_level else self.all_words() if isinstance(time, (float, int)): time = [time, time] elif isinstance(time, dict): time = [time['start'], time['end']] start, end = time if within: def is_in_range(c): return start <= c.start and end >= c.end else: def is_in_range(c): return start <= c.end and end >= c.start return [c for c in contents if is_in_range(c)] def split_by_gap( self, max_gap: float = 0.1, lock: bool = False, newline: bool = False ) -> "WhisperResult": """ Split (in-place) any segment where the gap between two of its words is greater than ``max_gap``. Parameters ---------- max_gap : float, default 0.1 Maximum second(s) allowed between two words if the same segment. lock : bool, default False Whether to prevent future splits/merges from altering changes made by this method. newline: bool, default False Whether to insert line break at the split points instead of splitting into separate segments. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ self._split_segments(lambda x: x.get_gap_indices(max_gap), lock=lock, newline=newline) if self._regroup_history: self._regroup_history += '_' self._regroup_history += f'sg={max_gap}+{int(lock)}+{int(newline)}' return self def merge_by_gap( self, min_gap: float = 0.1, max_words: int = None, max_chars: int = None, is_sum_max: bool = False, lock: bool = False ) -> "WhisperResult": """ Merge (in-place) any pair of adjacent segments if the gap between them <= ``min_gap``. Parameters ---------- min_gap : float, default 0.1 Minimum second(s) allow between two segment. max_words : int, optional Maximum number of words allowed in each segment. max_chars : int, optional Maximum number of characters allowed in each segment. is_sum_max : bool, default False Whether ``max_words`` and ``max_chars`` is applied to the merged segment instead of the individual segments to be merged. lock : bool, default False Whether to prevent future splits/merges from altering changes made by this method. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ indices = self.get_gap_indices(min_gap) self._merge_segments(indices, max_words=max_words, max_chars=max_chars, is_sum_max=is_sum_max, lock=lock) if self._regroup_history: self._regroup_history += '_' self._regroup_history += f'mg={min_gap}+{max_words or ""}+{max_chars or ""}+{int(is_sum_max)}+{int(lock)}' return self def split_by_punctuation( self, punctuation: Union[List[str], List[Tuple[str, str]], str], lock: bool = False, newline: bool = False, min_words: Optional[int] = None, min_chars: Optional[int] = None, min_dur: Optional[int] = None ) -> "WhisperResult": """ Split (in-place) segments at words that start/end with ``punctuation``. Parameters ---------- punctuation : list of str of list of tuple of (str, str) or str Punctuation(s) to split segments by. lock : bool, default False Whether to prevent future splits/merges from altering changes made by this method. newline : bool, default False Whether to insert line break at the split points instead of splitting into separate segments. min_words : int, optional Split segments with words >= ``min_words``. min_chars : int, optional Split segments with characters >= ``min_chars``. min_dur : int, optional split segments with duration (in seconds) >= ``min_dur``. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ def _over_max(x: Segment): return ( (min_words and len(x.words) >= min_words) or (min_chars and x.char_count() >= min_chars) or (min_dur and x.duration >= min_dur) ) indices = set(s.id for s in self.segments if _over_max(s)) if any((min_words, min_chars, min_dur)) else None def _get_indices(x: Segment): return x.get_punctuation_indices(punctuation) if indices is None or x.id in indices else [] self._split_segments(_get_indices, lock=lock, newline=newline) if self._regroup_history: self._regroup_history += '_' punct_str = '/'.join(p if isinstance(p, str) else '*'.join(p) for p in punctuation) self._regroup_history += f'sp={punct_str}+{int(lock)}+{int(newline)}' self._regroup_history += f'+{min_words or ""}+{min_chars or ""}+{min_dur or ""}'.rstrip('+') return self def merge_by_punctuation( self, punctuation: Union[List[str], List[Tuple[str, str]], str], max_words: int = None, max_chars: int = None, is_sum_max: bool = False, lock: bool = False ) -> "WhisperResult": """ Merge (in-place) any two segments that has specific punctuations inbetween. Parameters ---------- punctuation : list of str of list of tuple of (str, str) or str Punctuation(s) to merge segments by. max_words : int, optional Maximum number of words allowed in each segment. max_chars : int, optional Maximum number of characters allowed in each segment. is_sum_max : bool, default False Whether ``max_words`` and ``max_chars`` is applied to the merged segment instead of the individual segments to be merged. lock : bool, default False Whether to prevent future splits/merges from altering changes made by this method. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ indices = self.get_punctuation_indices(punctuation) self._merge_segments(indices, max_words=max_words, max_chars=max_chars, is_sum_max=is_sum_max, lock=lock) if self._regroup_history: self._regroup_history += '_' punct_str = '/'.join(p if isinstance(p, str) else '*'.join(p) for p in punctuation) self._regroup_history += f'mp={punct_str}+{max_words or ""}+{max_chars or ""}+{int(is_sum_max)}+{int(lock)}' return self def merge_all_segments(self) -> "WhisperResult": """ Merge all segments into one segment. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ if not self.segments: return self if self.has_words: self.segments[0].words = self.all_words() else: self.segments[0].text += ''.join(s.text for s in self.segments[1:]) if all(s.tokens is not None for s in self.segments): self.segments[0].tokens += list(chain.from_iterable(s.tokens for s in self.segments[1:])) self.segments[0].end = self.segments[-1].end self.segments = [self.segments[0]] self.reassign_ids() self.update_all_segs_with_words() if self._regroup_history: self._regroup_history += '_' self._regroup_history += 'ms' return self def split_by_length( self, max_chars: int = None, max_words: int = None, even_split: bool = True, force_len: bool = False, lock: bool = False, include_lock: bool = False, newline: bool = False ) -> "WhisperResult": """ Split (in-place) any segment that exceeds ``max_chars`` or ``max_words`` into smaller segments. Parameters ---------- max_chars : int, optional Maximum number of characters allowed in each segment. max_words : int, optional Maximum number of words allowed in each segment. even_split : bool, default True Whether to evenly split a segment in length if it exceeds ``max_chars`` or ``max_words``. force_len : bool, default False Whether to force a constant length for each segment except the last segment. This will ignore all previous non-locked segment boundaries. lock : bool, default False Whether to prevent future splits/merges from altering changes made by this method. include_lock: bool, default False Whether to include previous lock before splitting based on max_words, if ``even_split = False``. Splitting will be done after the first non-locked word > ``max_chars`` / ``max_words``. newline: bool, default False Whether to insert line break at the split points instead of splitting into separate segments. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. Notes ----- If ``even_split = True``, segments can still exceed ``max_chars`` and locked words will be ignored to avoid uneven splitting. """ if force_len: self.merge_all_segments() self._split_segments( lambda x: x.get_length_indices( max_chars=max_chars, max_words=max_words, even_split=even_split, include_lock=include_lock ), lock=lock, newline=newline ) if self._regroup_history: self._regroup_history += '_' self._regroup_history += (f'sl={max_chars or ""}+{max_words or ""}+{int(even_split)}+{int(force_len)}' f'+{int(lock)}+{int(include_lock)}+{int(newline)}') return self def split_by_duration( self, max_dur: float, even_split: bool = True, force_len: bool = False, lock: bool = False, include_lock: bool = False, newline: bool = False ) -> "WhisperResult": """ Split (in-place) any segment that exceeds ``max_dur`` into smaller segments. Parameters ---------- max_dur : float Maximum duration (in seconds) per segment. even_split : bool, default True Whether to evenly split a segment in length if it exceeds ``max_dur``. force_len : bool, default False Whether to force a constant length for each segment except the last segment. This will ignore all previous non-locked segment boundaries. lock : bool, default False Whether to prevent future splits/merges from altering changes made by this method. include_lock: bool, default False Whether to include previous lock before splitting based on max_words, if ``even_split = False``. Splitting will be done after the first non-locked word > ``max_dur``. newline: bool, default False Whether to insert line break at the split points instead of splitting into separate segments. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. Notes ----- If ``even_split = True``, segments can still exceed ``max_dur`` and locked words will be ignored to avoid uneven splitting. """ if force_len: self.merge_all_segments() self._split_segments( lambda x: x.get_duration_indices( max_dur=max_dur, even_split=even_split, include_lock=include_lock ), lock=lock, newline=newline ) if self._regroup_history: self._regroup_history += '_' self._regroup_history += (f'sd={max_dur}+{int(even_split)}+{int(force_len)}' f'+{int(lock)}+{int(include_lock)}+{int(newline)}') return self def clamp_max( self, medium_factor: float = 2.5, max_dur: float = None, clip_start: Optional[bool] = None, verbose: bool = False ) -> "WhisperResult": """ Clamp all word durations above certain value. This is most effective when applied before and after other regroup operations. Parameters ---------- medium_factor : float, default 2.5 Clamp durations above (``medium_factor`` * medium duration) per segment. If ``medium_factor = None/0`` or segment has less than 3 words, it will be ignored and use only ``max_dur``. max_dur : float, optional Clamp durations above ``max_dur``. clip_start : bool or None, default None Whether to clamp the start of a word. If ``None``, clamp the start of first word and end of last word per segment. verbose : bool, default False Whether to print out the timestamp changes. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ if not (medium_factor or max_dur): raise ValueError('At least one of following arguments requires non-zero value: medium_factor; max_dur') if not self.has_words: warnings.warn('Cannot clamp due to missing/no word-timestamps') return self for seg in self.segments: curr_max_dur = None if medium_factor and len(seg.words) > 2: durations = np.array([word.duration for word in seg.words]) durations.sort() curr_max_dur = medium_factor * durations[len(durations)//2 + 1] if max_dur and (not curr_max_dur or curr_max_dur > max_dur): curr_max_dur = max_dur if not curr_max_dur: continue if clip_start is None: seg.words[0].clamp_max(curr_max_dur, clip_start=True, verbose=verbose) seg.words[-1].clamp_max(curr_max_dur, clip_start=False, verbose=verbose) else: for i, word in enumerate(seg.words): word.clamp_max(curr_max_dur, clip_start=clip_start, verbose=verbose) seg.update_seg_with_words() if self._regroup_history: self._regroup_history += '_' self._regroup_history += f'cm={medium_factor}+{max_dur or ""}+{clip_start or ""}+{int(verbose)}' return self def lock( self, startswith: Union[str, List[str]] = None, endswith: Union[str, List[str]] = None, right: bool = True, left: bool = False, case_sensitive: bool = False, strip: bool = True ) -> "WhisperResult": """ Lock words/segments with matching prefix/suffix to prevent splitting/merging. Parameters ---------- startswith: str or list of str Prefixes to lock. endswith: str or list of str Suffixes to lock. right : bool, default True Whether prevent splits/merges with the next word/segment. left : bool, default False Whether prevent splits/merges with the previous word/segment. case_sensitive : bool, default False Whether to match the case of the prefixes/suffixes with the words/segments. strip : bool, default True Whether to ignore spaces before and after both words/segments and prefixes/suffixes. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ assert startswith or endswith, 'Must specify [startswith] or/and [endswith].' startswith = [] if startswith is None else ([startswith] if isinstance(startswith, str) else startswith) endswith = [] if endswith is None else ([endswith] if isinstance(endswith, str) else endswith) if not case_sensitive: startswith = [t.lower() for t in startswith] endswith = [t.lower() for t in endswith] if strip: startswith = [t.strip() for t in startswith] endswith = [t.strip() for t in endswith] for part in self.all_words_or_segments(): text = part.word if hasattr(part, 'word') else part.text if not case_sensitive: text = text.lower() if strip: text = text.strip() for prefix in startswith: if text.startswith(prefix): if right: part.lock_right() if left: part.lock_left() for suffix in endswith: if text.endswith(suffix): if right: part.lock_right() if left: part.lock_left() if self._regroup_history: self._regroup_history += '_' startswith_str = (startswith if isinstance(startswith, str) else '/'.join(startswith)) if startswith else "" endswith_str = (endswith if isinstance(endswith, str) else '/'.join(endswith)) if endswith else "" self._regroup_history += (f'l={startswith_str}+{endswith_str}' f'+{int(right)}+{int(left)}+{int(case_sensitive)}+{int(strip)}') return self def remove_word( self, word: Union[WordTiming, Tuple[int, int]], reassign_ids: bool = True, verbose: bool = True ) -> 'WhisperResult': """ Remove a word. Parameters ---------- word : WordTiming or tuple of (int, int) Instance of :class:`stable_whisper.result.WordTiming` or tuple of (segment index, word index). reassign_ids : bool, default True Whether to reassign segment and word ids (indices) after removing ``word``. verbose : bool, default True Whether to print detail of the removed word. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ if isinstance(word, WordTiming): if self[word.segment_id][word.id] is not word: self.reassign_ids() if self[word.segment_id][word.id] is not word: raise ValueError('word not in result') seg_id, word_id = word.segment_id, word.id else: seg_id, word_id = word if verbose: print(f'Removed: {self[seg_id][word_id].to_dict()}') del self.segments[seg_id].words[word_id] if not reassign_ids: return self if self[seg_id].has_words: self[seg_id].reassign_ids() else: self.remove_no_word_segments() return self def remove_segment( self, segment: Union[Segment, int], reassign_ids: bool = True, verbose: bool = True ) -> 'WhisperResult': """ Remove a segment. Parameters ---------- segment : Segment or int Instance :class:`stable_whisper.result.Segment` or segment index. reassign_ids : bool, default True Whether to reassign segment IDs (indices) after removing ``segment``. verbose : bool, default True Whether to print detail of the removed word. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ if isinstance(segment, Segment): if self[segment.id] is not segment: self.reassign_ids() if self[segment.id] is not segment: raise ValueError('segment not in result') segment = segment.id if verbose: print(f'Removed: [id:{self[segment].id}] {self[segment].to_display_str(True)}') del self.segments[segment] if not reassign_ids: return self self.reassign_ids(True) return self def remove_repetition( self, max_words: int = 1, case_sensitive: bool = False, strip: bool = True, ignore_punctuations: str = "\"',.?!", extend_duration: bool = True, verbose: bool = True ) -> 'WhisperResult': """ Remove words that repeat consecutively. Parameters ---------- max_words : int Maximum number of words to look for consecutively. case_sensitive : bool, default False Whether the case of words need to match to be considered as repetition. strip : bool, default True Whether to ignore spaces before and after each word. ignore_punctuations : bool, default '"',.?!' Ending punctuations to ignore. extend_duration: bool, default True Whether to extend the duration of the previous word to cover the duration of the repetition. verbose: bool, default True Whether to print detail of the removed repetitions. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ if not self.has_words: return self for count in range(1, max_words + 1): all_words = self.all_words() if len(all_words) < 2: return self all_words_str = [w.word for w in all_words] if strip: all_words_str = [w.strip() for w in all_words_str] if ignore_punctuations: ptn = f'[{ignore_punctuations}]+$' all_words_str = [re.sub(ptn, '', w) for w in all_words_str] if not case_sensitive: all_words_str = [w.lower() for w in all_words_str] next_i = None changes = [] for i in reversed(range(count*2, len(all_words_str)+1)): if next_i is not None: if next_i != i: continue else: next_i = None s = i - count if all_words_str[s - count:s] != all_words_str[s:i]: continue next_i = s if extend_duration: all_words[s-1].end = all_words[i-1].end temp_changes = [] for j in reversed(range(s, i)): if verbose: temp_changes.append(f'- {all_words[j].to_dict()}') self.remove_word(all_words[j], False, verbose=False) if temp_changes: changes.append( f'Remove: [{format_timestamp(all_words[s].start)} -> {format_timestamp(all_words[i-1].end)}] ' + ''.join(_w.word for _w in all_words[s:i]) + '\n' + '\n'.join(reversed(temp_changes)) + '\n' ) for i0, i1 in zip(range(s - count, s), range(s, i)): if len(all_words[i0].word) < len(all_words[i1].word): all_words[i1].start = all_words[i0].start all_words[i1].end = all_words[i0].end _sid, _wid = all_words[i0].segment_id, all_words[i0].id self.segments[_sid].words[_wid] = all_words[i1] if changes: print('\n'.join(reversed(changes))) self.remove_no_word_segments() self.update_all_segs_with_words() return self def remove_words_by_str( self, words: Union[str, List[str], None], case_sensitive: bool = False, strip: bool = True, ignore_punctuations: str = "\"',.?!", min_prob: float = None, filters: Callable = None, verbose: bool = True ) -> 'WhisperResult': """ Remove words that match ``words``. Parameters ---------- words : str or list of str or None A word or list of words to remove.``None`` for all words to be passed into ``filters``. case_sensitive : bool, default False Whether the case of words need to match to be considered as repetition. strip : bool, default True Whether to ignore spaces before and after each word. ignore_punctuations : bool, default '"',.?!' Ending punctuations to ignore. min_prob : float, optional Acts as the first filter the for the words that match ``words``. Words with probability < ``min_prob`` will be removed if ``filters`` is ``None``, else pass the words into ``filters``. Words without probability will be treated as having probability < ``min_prob``. filters : Callable, optional A function that takes an instance of :class:`stable_whisper.result.WordTiming` as its only argument. This function is custom filter for the words that match ``words`` and were not caught by ``min_prob``. verbose: Whether to print detail of the removed words. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ if not self.has_words: return self if isinstance(words, str): words = [words] all_words = self.all_words() all_words_str = [w.word for w in all_words] if strip: all_words_str = [w.strip() for w in all_words_str] words = [w.strip() for w in words] if ignore_punctuations: ptn = f'[{ignore_punctuations}]+$' all_words_str = [re.sub(ptn, '', w) for w in all_words_str] words = [re.sub(ptn, '', w) for w in words] if not case_sensitive: all_words_str = [w.lower() for w in all_words_str] words = [w.lower() for w in words] changes = [] for i, w in reversed(list(enumerate(all_words_str))): if not (words is None or any(w == _w for _w in words)): continue if ( (min_prob is None or all_words[i].probability is None or min_prob > all_words[i].probability) and (filters is None or filters(all_words[i])) ): if verbose: changes.append(f'Removed: {all_words[i].to_dict()}') self.remove_word(all_words[i], False, verbose=False) if changes: print('\n'.join(reversed(changes))) self.remove_no_word_segments() self.update_all_segs_with_words() return self def fill_in_gaps( self, other_result: Union['WhisperResult', str], min_gap: float = 0.1, case_sensitive: bool = False, strip: bool = True, ignore_punctuations: str = "\"',.?!", verbose: bool = True ) -> 'WhisperResult': """ Fill in segment gaps larger than ``min_gap`` with content from ``other_result`` at the times of gaps. Parameters ---------- other_result : WhisperResult or str Another transcription result as an instance of :class:`stable_whisper.result.WhisperResult` or path to the JSON of the result. min_gap : float, default 0.1 The minimum seconds of a gap between segments that must be exceeded to be filled in. case_sensitive : bool, default False Whether to consider the case of the first and last word of the gap to determine overlapping words to remove before filling in. strip : bool, default True Whether to ignore spaces before and after the first and last word of the gap to determine overlapping words to remove before filling in. ignore_punctuations : bool, default '"',.?!' Ending punctuations to ignore in the first and last word of the gap to determine overlapping words to remove before filling in. verbose: Whether to print detail of the filled content. Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. """ if len(self.segments) < 2: return self if isinstance(other_result, str): other_result = WhisperResult(other_result) if strip: def strip_space(w): return w.strip() else: def strip_space(w): return w if ignore_punctuations: ptn = f'[{ignore_punctuations}]+$' def strip_punctuations(w): return re.sub(ptn, '', strip_space(w)) else: strip_punctuations = strip_space if case_sensitive: strip = strip_punctuations else: def strip(w): return strip_punctuations(w).lower() seg_pairs = list(enumerate(zip(self.segments[:-1], self.segments[1:]))) seg_pairs.insert(0, (-1, (None, self.segments[0]))) seg_pairs.append((seg_pairs[-1][0]+1, (self.segments[-1], None))) changes = [] for i, (seg0, seg1) in reversed(seg_pairs): first_word = None if seg0 is None else seg0.words[-1] last_word = None if seg1 is None else seg1.words[0] start = (other_result[0].start if first_word is None else first_word.end) end = other_result[-1].end if last_word is None else last_word.start if end - start <= min_gap: continue gap_words = other_result.get_content_by_time((start, end)) if first_word is not None and gap_words and strip(first_word.word) == strip(gap_words[0].word): first_word.end = gap_words[0].end gap_words = gap_words[1:] if last_word is not None and gap_words and strip(last_word.word) == strip(gap_words[-1].word): last_word.start = gap_words[-1].start gap_words = gap_words[:-1] if not gap_words: continue if last_word is not None and last_word.start < gap_words[-1].end: last_word.start = gap_words[-1].end new_segments = [other_result[gap_words[0].segment_id].copy([])] for j, new_word in enumerate(gap_words): new_word = deepcopy(new_word) if j == 0 and first_word is not None and first_word.end > gap_words[0].start: new_word.start = first_word.end if new_segments[-1].id != new_word.segment_id: new_segments.append(other_result[new_word.segment_id].copy([])) new_segments[-1].words.append(new_word) if verbose: changes.append('\n'.join('Added: ' + s.to_display_str(True) for s in new_segments)) self.segments = self.segments[:i+1] + new_segments + self.segments[i+1:] if changes: print('\n'.join(reversed(changes))) self.reassign_ids() self.update_all_segs_with_words() return self def regroup( self, regroup_algo: Union[str, bool] = None, verbose: bool = False, only_show: bool = False ) -> "WhisperResult": """ Regroup (in-place) words into segments. Parameters ---------- regroup_algo: str or bool, default 'da' String representation of a custom regrouping algorithm or ``True`` use to the default algorithm 'da'. verbose : bool, default False Whether to show all the methods and arguments parsed from ``regroup_algo``. only_show : bool, default False Whether to show the all methods and arguments parsed from ``regroup_algo`` without running the methods Returns ------- stable_whisper.result.WhisperResult The current instance after the changes. Notes ----- Syntax for string representation of custom regrouping algorithm. Method keys: sg: split_by_gap sp: split_by_punctuation sl: split_by_length sd: split_by_duration mg: merge_by_gap mp: merge_by_punctuation ms: merge_all_segment cm: clamp_max l: lock us: unlock_all_segments da: default algorithm (cm_sp=.* /。/?/？/,* /，_sg=.5_mg=.3+3_sp=.* /。/?/？) rw: remove_word rs: remove_segment rp: remove_repetition rws: remove_words_by_str fg: fill_in_gaps Metacharacters: = separates a method key and its arguments (not used if no argument) _ separates method keys (after arguments if there are any) + separates arguments for a method key / separates an argument into list of strings * separates an item in list of strings into a nested list of strings Notes: -arguments are parsed positionally -if no argument is provided, the default ones will be used -use 1 or 0 to represent True or False Example 1: merge_by_gap(.2, 10, lock=True) mg=.2+10+++1 Note: [lock] is the 5th argument hence the 2 missing arguments inbetween the three + before 1 Example 2: split_by_punctuation([('.', ' '), '。', '?', '？'], True) sp=.* /。/?/？+1 Example 3: merge_all_segments().split_by_gap(.5).merge_by_gap(.15, 3) ms_sg=.5_mg=.15+3 """ if regroup_algo is False: return self if regroup_algo is None or regroup_algo is True: regroup_algo = 'da' for method, kwargs, msg in self.parse_regroup_algo(regroup_algo, include_str=verbose or only_show): if msg: print(msg) if not only_show: method(**kwargs) return self def parse_regroup_algo(self, regroup_algo: str, include_str: bool = True) -> List[Tuple[Callable, dict, str]]: methods = dict( sg=self.split_by_gap, sp=self.split_by_punctuation, sl=self.split_by_length, sd=self.split_by_duration, mg=self.merge_by_gap, mp=self.merge_by_punctuation, ms=self.merge_all_segments, cm=self.clamp_max, us=self.unlock_all_segments, l=self.lock, rw=self.remove_word, rs=self.remove_segment, rp=self.remove_repetition, rws=self.remove_words_by_str, fg=self.fill_in_gaps, ) if not regroup_algo: return [] calls = regroup_algo.split('_') if 'da' in calls: default_calls = 'cm_sp=.* /。/?/？/,* /，_sg=.5_mg=.3+3_sp=.* /。/?/？'.split('_') calls = chain.from_iterable(default_calls if method == 'da' else [method] for method in calls) operations = [] for method in calls: method, args = method.split('=', maxsplit=1) if '=' in method else (method, '') if method not in methods: raise NotImplementedError(f'{method} is not one of the available methods: {tuple(methods.keys())}') args = [] if len(args) == 0 else list(map(str_to_valid_type, args.split('+'))) kwargs = {k: v for k, v in zip(methods[method].__code__.co_varnames[1:], args) if v is not None} if include_str: kwargs_str = ', '.join(f'{k}="{v}"' if isinstance(v, str) else f'{k}={v}' for k, v in kwargs.items()) op_str = f'{methods[method].__name__}({kwargs_str})' else: op_str = None operations.append((methods[method], kwargs, op_str)) return operations def find(self, pattern: str, word_level=True, flags=None) -> "WhisperResultMatches": """ Find segments/words and timestamps with regular expression. Parameters ---------- pattern : str RegEx pattern to search for. word_level : bool, default True Whether to search at word-level. flags : optional RegEx flags. Returns ------- stable_whisper.result.WhisperResultMatches An instance of :class:`stable_whisper.result.WhisperResultMatches` with word/segment that match ``pattern``. """ return WhisperResultMatches(self).find(pattern, word_level=word_level, flags=flags) @property def text(self): return ''.join(s.text for s in self.segments) @property def regroup_history(self): # same syntax as ``regroup_algo`` for :meth:``result.WhisperResult.regroup` return self._regroup_history @property def nonspeech_sections(self): return self._nonspeech_sections def show_regroup_history(self): """ Print details of all regrouping operations that been performed on data. """ if not self._regroup_history: print('Result has no history.') for *_, msg in self.parse_regroup_algo(self._regroup_history): print(f'.{msg}') def __len__(self): return len(self.segments) def unlock_all_segments(self): for s in self.segments: s.unlock_all_words() return self def reset(self): """ Restore all values to that at initialization. """ self.language = self.ori_dict.get('language') self._regroup_history = '' segments = self.ori_dict.get('segments') self.segments: List[Segment] = [Segment(**s) for s in segments] if segments else [] if self._forced_order: self.force_order() self.remove_no_word_segments(any(seg.has_words for seg in self.segments)) self.update_all_segs_with_words() @property def has_words(self): return all(seg.has_words for seg in self.segments) to_srt_vtt = result_to_srt_vtt to_ass = result_to_ass to_tsv = result_to_tsv to_txt = result_to_txt save_as_json = save_as_json class SegmentMatch: def __init__( self, segments: Union[List[Segment], Segment], _word_indices: List[List[int]] = None, _text_match: str = None ): self.segments = [segments] if isinstance(segments, Segment) else segments self.word_indices = [] if _word_indices is None else _word_indices self.words = [self.segments[i].words[j] for i, indices in enumerate(self.word_indices) for j in indices] if len(self.words) != 0: self.text = ''.join( self.segments[i].words[j].word for i, indices in enumerate(self.word_indices) for j in indices ) else: self.text = ''.join(seg.text for seg in self.segments) self.text_match = _text_match @property def start(self): return ( self.words[0].start if len(self.words) != 0 else (self.segments[0].start if len(self.segments) != 0 else None) ) @property def end(self): return ( self.words[-1].end if len(self.words) != 0 else (self.segments[-1].end if len(self.segments) != 0 else None) ) def __len__(self): return len(self.segments) def __repr__(self): return self.__dict__.__repr__() def __str__(self): return self.__dict__.__str__() class WhisperResultMatches: """ RegEx matches for WhisperResults. """ # Use WhisperResult.find() instead of instantiating this class directly. def __init__( self, matches: Union[List[SegmentMatch], WhisperResult], _segment_indices: List[List[int]] = None ): if isinstance(matches, WhisperResult): self.matches = list(map(SegmentMatch, matches.segments)) self._segment_indices = [[i] for i in range(len(matches.segments))] else: self.matches = matches assert _segment_indices is not None assert len(self.matches) == len(_segment_indices) assert all(len(match.segments) == len(_segment_indices[i]) for i, match in enumerate(self.matches)) self._segment_indices = _segment_indices @property def segment_indices(self): return self._segment_indices def _curr_seg_groups(self) -> List[List[Tuple[int, Segment]]]: seg_groups, curr_segs = [], [] curr_max = -1 for seg_indices, match in zip(self._segment_indices, self.matches): for i, seg in zip(sorted(seg_indices), match.segments): if i > curr_max: curr_segs.append((i, seg)) if i - 1 != curr_max: seg_groups.append(curr_segs) curr_segs = [] curr_max = i if curr_segs: seg_groups.append(curr_segs) return seg_groups def find(self, pattern: str, word_level=True, flags=None) -> "WhisperResultMatches": """ Find segments/words and timestamps with regular expression. Parameters ---------- pattern : str RegEx pattern to search for. word_level : bool, default True Whether to search at word-level. flags : optional RegEx flags. Returns ------- stable_whisper.result.WhisperResultMatches An instance of :class:`stable_whisper.result.WhisperResultMatches` with word/segment that match ``pattern``. """ seg_groups = self._curr_seg_groups() matches: List[SegmentMatch] = [] match_seg_indices: List[List[int]] = [] if word_level: if not all(all(seg.has_words for seg in match.segments) for match in self.matches): warnings.warn('Cannot perform word-level search with segment(s) missing word timestamps.') word_level = False for segs in seg_groups: if word_level: idxs = list(chain.from_iterable( [(i, j)]*len(word.word) for (i, seg) in segs for j, word in enumerate(seg.words) )) text = ''.join(word.word for (_, seg) in segs for word in seg.words) else: idxs = list(chain.from_iterable([(i, None)]*len(seg.text) for (i, seg) in segs)) text = ''.join(seg.text for (_, seg) in segs) assert len(idxs) == len(text) for curr_match in re.finditer(pattern, text, flags=flags or 0): start, end = curr_match.span() curr_idxs = idxs[start: end] curr_seg_idxs = sorted(set(i[0] for i in curr_idxs)) if word_level: curr_word_idxs = [ sorted(set(j for i, j in curr_idxs if i == seg_idx)) for seg_idx in curr_seg_idxs ] else: curr_word_idxs = None matches.append(SegmentMatch( segments=[s for i, s in segs if i in curr_seg_idxs], _word_indices=curr_word_idxs, _text_match=curr_match.group() )) match_seg_indices.append(curr_seg_idxs) return WhisperResultMatches(matches, match_seg_indices) def __len__(self): return len(self.matches) def __bool__(self): return self.__len__() != 0 def __getitem__(self, idx): return self.matches[idx]