stable_whisper/timing.py

import string
import torch
import numpy as np
from typing import TYPE_CHECKING, List, Callable, Optional
from itertools import chain
from whisper.audio import TOKENS_PER_SECOND, N_SAMPLES_PER_TOKEN
from whisper.timing import WordTiming, median_filter, dtw, merge_punctuations

if TYPE_CHECKING:
    from whisper.tokenizer import Tokenizer
    from whisper.model import Whisper


# modified version of whisper.timing.find_alignment
def find_alignment_stable(
        model: "Whisper",
        tokenizer: "Tokenizer",
        text_tokens: List[int],
        mel: torch.Tensor,
        num_samples: int,
        *,
        medfilt_width: int = 7,
        qk_scale: float = 1.0,
        ts_num: int = 0,
        ts_noise: float = 0.1,
        token_split=None,
        audio_features: torch.Tensor = None
) -> List[WordTiming]:
    tokens = torch.tensor(
        [
            *tokenizer.sot_sequence,
            tokenizer.no_timestamps,
            *text_tokens,
            tokenizer.eot,
        ]
    ).to(model.device)

    # install hooks on the cross attention layers to retrieve the attention weights
    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)
    ]

    with torch.no_grad():
        if audio_features is None:
            audio_features = model.encoder(mel.unsqueeze(0))
        if ts_num:
            if ts_noise is None:
                ts_noise = 0.1
            extra_audio_features = audio_features.repeat_interleave(ts_num, 0)
            torch.manual_seed(0)
            audio_features = torch.cat([audio_features,
                                        extra_audio_features *
                                        (1 - (torch.rand_like(extra_audio_features) * ts_noise))],
                                       dim=0)
            logits = model.decoder(tokens.unsqueeze(0).repeat_interleave(audio_features.shape[0], 0),
                                   audio_features)
        else:
            logits = model.decoder(tokens.unsqueeze(0), audio_features)

        logits = logits[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]
        text_token_probs = text_token_probs.tolist()

    for hook in hooks:
        hook.remove()

    # heads * tokens * frames
    weights = torch.cat([QKs[_l][:, _h] for _l, _h in model.alignment_heads.indices().T], dim=0)
    weights = weights[:, :, : round(num_samples / N_SAMPLES_PER_TOKEN)]
    weights = (weights * qk_scale).softmax(dim=-1)
    std, mean = torch.std_mean(weights, dim=-2, keepdim=True, unbiased=False)
    weights = (weights - mean) / std
    weights = median_filter(weights, medfilt_width)

    matrix = weights.mean(axis=0)
    matrix = matrix[len(tokenizer.sot_sequence): -1]
    text_indices, time_indices = dtw(-matrix)

    if token_split is None:
        words, word_tokens = tokenizer.split_to_word_tokens(text_tokens + [tokenizer.eot])
    else:
        words, word_tokens = token_split
        words.append(tokenizer.decode([tokenizer.eot]))
        word_tokens.append([tokenizer.eot])
    word_boundaries = np.pad(np.cumsum([len(t) for t in word_tokens[:-1]]), (1, 0))

    jumps = np.pad(np.diff(text_indices), (1, 0), constant_values=1).astype(bool)
    jump_times = time_indices[jumps].clip(min=0) / TOKENS_PER_SECOND
    start_times = jump_times[word_boundaries[:-1]]
    end_times = jump_times[word_boundaries[1:]]
    word_probabilities = [
        np.mean(text_token_probs[i:j])
        for i, j in zip(word_boundaries[:-1], word_boundaries[1:])
    ]

    return [
        WordTiming(word, tokens, start, end, probability)
        for word, tokens, start, end, probability in zip(
            words, word_tokens, start_times, end_times, word_probabilities
        )
    ]


def _split_tokens(tokens: List[int], tokenizer: "Tokenizer"):
    split_by_space = getattr(tokenizer, 'language_code', tokenizer.language) not in {"zh", "ja", "th", "lo", "my"}
    text = tokenizer.decode_with_timestamps(tokens)
    words = []
    word_tokens = []
    curr_tokens = []
    is_append = False
    for token in tokens:
        curr_tokens.append(token)
        curr_text = tokenizer.decode(curr_tokens)
        is_whole = token >= tokenizer.eot
        if not is_whole:
            is_whole = text[:len(curr_text)] == curr_text
            if is_whole and split_by_space:
                is_append = not (curr_text.startswith(" ") or curr_text.strip() in string.punctuation)

        if is_whole:
            if is_append and len(words) != 0:
                words[-1] += curr_text
                word_tokens[-1].extend(curr_tokens)
            else:
                words.append(curr_text)
                word_tokens.append(curr_tokens)
            text = text[len(curr_text):]
            curr_tokens = []

    if len(curr_tokens) != 0:
        words.append(curr_text if len(text) == 0 else text)
        word_tokens.append(curr_tokens)
    elif len(text) != 0:
        words[-1] += text

    return words, word_tokens


def split_word_tokens(segments: List[dict],
                      tokenizer: "Tokenizer",
                      *,
                      padding: (str, int) = None,
                      split_callback: Callable = None):
    if padding is not None:
        if isinstance(padding, str):
            padding = tokenizer.encode(padding)
        else:
            padding = [padding]
    tokens = []
    seg_indices = []
    words = []
    word_tokens = []
    for i, s in enumerate(segments):
        temp_word_tokens = [t for t in s['tokens'] if not isinstance(t, int) or t < tokenizer.eot]
        curr_words, curr_word_tokens = (
            _split_tokens(temp_word_tokens, tokenizer)
            if split_callback is None else
            split_callback(temp_word_tokens, tokenizer)
        )
        assert len(curr_words) == len(curr_word_tokens), \
            f'word count and token group count do not match, {len(curr_words)} and {len(curr_word_tokens)}'
        if (
                padding is not None and
                curr_word_tokens[0][0] != padding and
                (len(tokens) == 0 or tokens[-1] != padding)
        ):
            tokens.extend(padding)
            words.append(None)
            word_tokens.append(padding)
        seg_indices.extend([i] * len(curr_words))
        tokens.extend(list(chain.from_iterable(curr_word_tokens)))
        words.extend(curr_words)
        word_tokens.extend(curr_word_tokens)

    return tokens, (words, word_tokens), seg_indices


def pop_empty_alignment(alignment: List[WordTiming]):
    return list(reversed([alignment.pop(i) for i in reversed(range(len(alignment))) if alignment[i].word is None]))


# modified version of whisper.timing.add_word_timestamps
def add_word_timestamps_stable(
        *,
        segments: List[dict],
        model: "Whisper",
        tokenizer: "Tokenizer",
        mel: torch.Tensor,
        num_samples: int,
        prepend_punctuations: str = "\"'“¿([{-",
        append_punctuations: str = "\"'.。,，!！?？:：”)]}、",
        audio_features: torch.Tensor = None,
        ts_num: int = 0,
        ts_noise: float = 0.1,
        min_word_dur: float = 0.1,
        split_callback: Callable = None,
        gap_padding: Optional[str] = ' ...',
        **kwargs,
):
    if len(segments) == 0:
        return

    if min_word_dur is None:
        min_word_dur = 0

    if prepend_punctuations is None:
        prepend_punctuations = "\"'“¿([{-"

    if append_punctuations is None:
        append_punctuations = "\"'.。,，!！?？:：”)]}、"

    def align():
        for seg in segments:
            seg['words'] = []

        text_tokens, token_split, seg_indices = split_word_tokens(segments, tokenizer,
                                                                  padding=gap_padding, split_callback=split_callback)

        alignment = find_alignment_stable(model, tokenizer, text_tokens, mel, num_samples,
                                          **kwargs,
                                          token_split=token_split,
                                          audio_features=audio_features,
                                          ts_num=ts_num,
                                          ts_noise=ts_noise)
        alt_beginning_alignment = pop_empty_alignment(alignment)

        merge_punctuations(alignment, prepend_punctuations, append_punctuations)

        time_offset = segments[0]["seek"]

        assert len(alignment) == len(seg_indices)
        assert (gap_padding is None or len(segments) == len(alt_beginning_alignment))
        for i, timing in zip(seg_indices, alignment):
            if len(timing.tokens) != 0:
                start = timing.start
                end = timing.end
                if (
                        len(segments[i]['words']) == 0 and
                        ((end - start) < min_word_dur) and
                        len(alt_beginning_alignment)
                ):
                    start = alt_beginning_alignment[i].start
                segments[i]['words'].append(
                    dict(
                        word=timing.word,
                        start=round(time_offset + start, 3),
                        end=round(time_offset + end, 3),
                        probability=timing.probability,
                        tokens=timing.tokens
                    )
                )

    align()
    if (
            gap_padding is not None and
            any(
                (word['end'] - word['start']) < min_word_dur
                for seg in segments
                for word in seg['words']
            )
    ):
        gap_padding = None
        align()

    for segment in segments:
        if len(words := segment["words"]) > 0:
            # adjust the segment-level timestamps based on the word-level timestamps
            segment["start"] = words[0]["start"]
            segment["end"] = words[-1]["end"]