miscellaneous improvements

Author: jongwook

tests/conftest.py

@@ -4,6 +4,10 @@
 import pytest
 
 
+def pytest_configure(config):
+    config.addinivalue_line("markers", "requires_cuda")
+
+
 @pytest.fixture
 def random():
     rand.seed(42)

tests/test_timing.py

@@ -10,7 +10,7 @@
     (10, 20), (32, 16), (123, 1500), (234, 189),
 ]
 shapes = [
-    (4, 5, 20, 345), (6, 12, 240, 512),
+    (10,), (1, 15),  (4, 5, 345), (6, 12, 240, 512),
 ]
 
 
@@ -65,7 +65,12 @@ def test_median_filter(shape):
 
     for filter_width in [3, 5, 7, 13]:
         filtered = median_filter(x, filter_width)
-        scipy_filtered = scipy.ndimage.median_filter(x, (1, 1, 1, filter_width), mode="nearest")
+
+        # using np.pad to reflect-pad, because Scipy's behavior is different near the edges.
+        pad_width = filter_width // 2
+        padded_x = np.pad(x, [(0, 0)] * (x.ndim - 1) + [(pad_width, pad_width)], mode="reflect")
+        scipy_filtered = scipy.ndimage.median_filter(padded_x, [1] * (x.ndim - 1) + [filter_width])
+        scipy_filtered = scipy_filtered[..., pad_width:-pad_width]
 
         assert np.allclose(filtered, scipy_filtered)
 

tests/test_transcribe.py

@@ -25,9 +25,10 @@ def test_transcribe(model_name: str):
     for segment in result["segments"]:
         for timing in segment["words"]:
             assert timing["start"] < timing["end"]
-            if timing["word"].strip() == "Americans":
-                assert timing["start"] <= 1.75
-                assert timing["end"] >= 2.05
+            if timing["word"].strip(" ,") == "Americans":
+                assert timing["start"] <= 1.8
+                assert timing["end"] >= 1.8
+                print(timing)
                 timing_checked = True
 
     assert timing_checked

whisper/__init__.py

@@ -29,6 +29,23 @@
     "large": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
 }
 
+# base85-encoded (n_layers, n_heads) boolean arrays indicating the cross-attention heads that are
+# highly correlated to the word-level timing, i.e. the alignment between audio and text tokens.
+_ALIGNMENT_HEADS = {
+    "tiny.en": b"ABzY8J1N>@0{>%R00Bk>$p{7v037`oCl~+#00",
+    "tiny": b"ABzY8bu8Lr0{>%RKn9Fp%m@SkK7Kt=7ytkO",
+    "base.en": b"ABzY8;40c<0{>%RzzG;p*o+Vo09|#PsxSZm00",
+    "base": b"ABzY8KQ!870{>%RzyTQH3`Q^yNP!>##QT-<FaQ7m",
+    "small.en": b"ABzY8>?_)10{>%RpeA61k&I|OI3I$65C{;;pbCHh0B{qLQ;+}v00",
+    "small": b"ABzY8DmU6=0{>%Rpa?J`kvJ6qF(V^F86#Xh7JUGMK}P<N0000",
+    "medium.en": b"ABzY8usPae0{>%R7<zz_OvQ{)4kMa0BMw6u5rT}kRKX;$NfYBv00*Hl@qhsU00",
+    "medium": b"ABzY8B0Jh+0{>%R7}kK1fFL7w6%<-Pf*t^=N)Qr&0RR9",
+    "large-v1": b"ABzY8r9j$a0{>%R7#4sLmoOs{s)o3~84-RPdcFk!JR<kSfC2yj",
+    "large-v2": b'ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj',
+    "large": b'ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj',
+}
+
+
 
 def _download(url: str, root: str, in_memory: bool) -> Union[bytes, str]:
     os.makedirs(root, exist_ok=True)
@@ -106,8 +123,10 @@ def load_model(name: str, device: Optional[Union[str, torch.device]] = None, dow
 
     if name in _MODELS:
         checkpoint_file = _download(_MODELS[name], download_root, in_memory)
+        alignment_heads = _ALIGNMENT_HEADS[name]
     elif os.path.isfile(name):
         checkpoint_file = open(name, "rb").read() if in_memory else name
+        alignment_heads = None
     else:
         raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
 
@@ -119,4 +138,7 @@ def load_model(name: str, device: Optional[Union[str, torch.device]] = None, dow
     model = Whisper(dims)
     model.load_state_dict(checkpoint["model_state_dict"])
 
+    if alignment_heads is not None:
+        model.set_alignment_heads(alignment_heads)
+
     return model.to(device)

whisper/model.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
 from typing import Dict
 from typing import Iterable, Optional
-
+import gzip
+import base64
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -213,6 +214,15 @@ def __init__(self, dims: ModelDimensions):
             self.dims.n_text_head,
             self.dims.n_text_layer,
         )
+        # use the last half layers for alignment by default; see `set_alignment_heads()` below
+        all_heads = torch.zeros(self.dims.n_text_layer, self.dims.n_text_head, dtype=torch.bool)
+        all_heads[self.dims.n_text_layer // 2:] = True
+        self.register_buffer("alignment_heads", all_heads.to_sparse(), persistent=False)
+
+    def set_alignment_heads(self, dump: bytes):
+        array = np.frombuffer(gzip.decompress(base64.b85decode(dump)), dtype=bool).copy()
+        mask = torch.from_numpy(array).reshape(self.dims.n_text_layer, self.dims.n_text_head)
+        self.register_buffer("alignment_heads", mask.to_sparse(), persistent=False)
 
     def embed_audio(self, mel: torch.Tensor):
         return self.encoder(mel)

whisper/timing.py

@@ -1,5 +1,9 @@
-from typing import List, TYPE_CHECKING
+import time
+import subprocess
+import warnings
 from dataclasses import dataclass
+from typing import List, TYPE_CHECKING
+
 import numba
 import numpy as np
 import torch
@@ -14,17 +18,31 @@
 
 def median_filter(x: torch.Tensor, filter_width: int):
     """Apply a median filter of width `filter_width` along the last dimension of `x`"""
-    assert 3 <= x.ndim <= 4, "`median_filter()` is implemented for only 3D or 4D tensors"
+    if (ndim := x.ndim) <= 2:  # `F.pad` does not support 1D or 2D inputs for reflect padding
+        x = x[None, None, :]
+
     assert filter_width > 0 and filter_width % 2 == 1, "`filter_width` should be an odd number"
 
-    x = F.pad(x, (filter_width // 2, filter_width // 2, 0, 0), mode='replicate')
+    result = None
+    x = F.pad(x, (filter_width // 2, filter_width // 2, 0, 0), mode="reflect")
     if x.is_cuda:
-        from .triton_ops import median_filter_cuda
-        return median_filter_cuda(x, filter_width)
+        try:
+            from .triton_ops import median_filter_cuda
+            result = median_filter_cuda(x, filter_width)
+        except (RuntimeError, subprocess.CalledProcessError):
+            warnings.warn(
+                "Failed to launch Triton kernels, likely due to missing CUDA toolkit; "
+                "falling back to a slower median kernel implementation..."
+            )
+
+    if result is None:
+        # sort() is faster than torch.median (https://github.com/pytorch/pytorch/issues/51450)
+        result = x.unfold(-1, filter_width, 1).sort()[0][..., filter_width // 2]
 
-    # sort() is faster than torch.median (https://github.com/pytorch/pytorch/issues/51450)
-    return x.unfold(-1, filter_width, 1).sort()[0][..., filter_width // 2]
+    if ndim <= 2:
+        result = result[0, 0]
 
+    return result
 
 @numba.jit
 def backtrace(trace: np.ndarray):
@@ -108,17 +126,24 @@ def dtw_cuda(x, BLOCK_SIZE=1024):
 
 def dtw(x: torch.Tensor) -> np.ndarray:
     if x.is_cuda:
-        return dtw_cuda(x)
+        try:
+            return dtw_cuda(x)
+        except (RuntimeError, subprocess.CalledProcessError):
+            warnings.warn(
+                "Failed to launch Triton kernels, likely due to missing CUDA toolkit; "
+                "falling back to a slower DTW implementation..."
+            )
 
     return dtw_cpu(x.double().cpu().numpy())
 
 
 @dataclass
-class Alignment:
-    words: List[str]
-    word_tokens: List[List[int]]
-    start_times: np.ndarray
-    end_times: np.ndarray
+class WordTiming:
+    word: str
+    tokens: List[int]
+    start: float
+    end: float
+    probability: float
 
 
 def find_alignment(
@@ -128,16 +153,14 @@ def find_alignment(
     mel: torch.Tensor,
     num_frames: int,
     *,
-    max_qk_layers: int = 6,
     medfilt_width: int = 7,
     qk_scale: float = 1.0,
-) -> Alignment:
+) -> List[WordTiming]:
     tokens = torch.tensor(
         [
             *tokenizer.sot_sequence,
-            tokenizer.timestamp_begin,
+            tokenizer.no_timestamps,
             *text_tokens,
-            tokenizer.timestamp_begin + num_frames // 2,
             tokenizer.eot,
         ]
     ).to(model.device)
@@ -146,78 +169,132 @@ def find_alignment(
     QKs = [None] * model.dims.n_text_layer
     hooks = [
         block.cross_attn.register_forward_hook(
-            lambda _, ins, outs, index=i: QKs.__setitem__(index, outs[-1])
+            lambda _, ins, outs, index=i: QKs.__setitem__(index, outs[-1][0])
         )
         for i, block in enumerate(model.decoder.blocks)
     ]
 
     with torch.no_grad():
-        model(mel.unsqueeze(0), tokens.unsqueeze(0))
+        logits = model(mel.unsqueeze(0), tokens.unsqueeze(0))[0]
+        token_probs = logits[len(tokenizer.sot_sequence):, :tokenizer.eot].softmax(dim=-1)
+        text_token_probs = token_probs[np.arange(len(text_tokens)), text_tokens].tolist()
 
     for hook in hooks:
         hook.remove()
 
-    weights = torch.cat(QKs[-max_qk_layers:])  # layers * heads * tokens * frames
-    weights = weights[:, :, :, : num_frames // 2]
-    weights = median_filter(weights, medfilt_width)
+    # heads * tokens * frames
+    weights = torch.stack([QKs[l][h] for l, h in model.alignment_heads.indices().T])
+    weights = weights[:, :, : num_frames // 2]
     weights = (weights * qk_scale).softmax(dim=-1)
-    weights = weights / weights.norm(dim=-2, keepdim=True)
-    matrix = weights.mean(axis=(0, 1)).neg()
+    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)
 
-    text_indices, time_indices = dtw(matrix)
+    words, word_tokens = tokenizer.split_to_word_tokens(text_tokens + [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] / TOKENS_PER_SECOND
-
-    if tokenizer.language in {"zh", "ja", "th", "lo", "my"}:
-        # These languages don't typically use spaces, so it is difficult to split words
-        # without morpheme analysis. Here, we instead split words at any
-        # position where the tokens are decoded as valid unicode points
-        words, word_tokens = tokenizer.split_tokens_on_unicode(tokens[1:].tolist())
-    else:
-        words, word_tokens = tokenizer.split_tokens_on_spaces(tokens[1:].tolist())
-
-    word_boundaries = np.pad(np.cumsum([len(t) for t in word_tokens]), (1, 0))
     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:])
+    ]
+
+    # hack: ensure the first and second word is not longer than twice the median word duration.
+    # a better segmentation algorithm based on VAD should be able to replace this.
+    word_durations = end_times - start_times
+    word_durations = word_durations[word_durations.nonzero()]
+    if len(word_durations) > 0:
+        median_duration = np.median(word_durations)
+        max_duration = median_duration * 2
+        if len(word_durations) >= 2 and word_durations[1] > max_duration:
+            end_times[0] = start_times[1] = max(end_times[2] / 2, end_times[2] - max_duration)
+        if len(word_durations) >= 1 and end_times[0] - start_times[0] > max_duration:
+            start_times[0] = max(0, end_times[0] - max_duration)
+
+    return [
+        WordTiming(word, tokens, start, end, probability)
+        for word, tokens, start, end, probability in zip(
+            words, word_tokens, start_times, end_times, word_probabilities
+        )
+    ]
+
 
-    return Alignment(words, word_tokens, start_times, end_times)
+def merge_punctuations(alignment: List[WordTiming], prepended: str, appended: str):
+    # merge prepended punctuations
+    i = len(alignment) - 2
+    j = len(alignment) - 1
+    while i >= 0:
+        previous = alignment[i]
+        following = alignment[j]
+        if previous.word.startswith(" ") and previous.word.strip() in prepended:
+            # prepend it to the following word
+            following.word = previous.word + following.word
+            following.tokens = previous.tokens + following.tokens
+            previous.word = ""
+            previous.tokens = []
+        else:
+            j = i
+        i -= 1
+
+    # merge appended punctuations
+    i = 0
+    j = 1
+    while j < len(alignment):
+        previous = alignment[i]
+        following = alignment[j]
+        if not previous.word.endswith(" ") and following.word in appended:
+            # append it to the previous word
+            previous.word = previous.word + following.word
+            previous.tokens = previous.tokens + following.tokens
+            following.word = ""
+            following.tokens = []
+        else:
+            i = j
+        j += 1
 
 
 def add_word_timestamps(
+    *,
     segments: List[dict],
     model: "Whisper",
     tokenizer: Tokenizer,
     mel: torch.Tensor,
     num_frames: int,
+    prepend_punctuations: str = "\"\'“¿([{-",
+    append_punctuations: str = "\"\'.。,，!！?？:：”)]}、",
     **hyperparams,
 ):
     if len(segments) == 0:
         return
 
     text_tokens = [t for segment in segments for t in segment["tokens"]]
     alignment = find_alignment(model, tokenizer, text_tokens, mel, num_frames, **hyperparams)
+    merge_punctuations(alignment, prepend_punctuations, append_punctuations)
 
     time_offset = segments[0]["seek"] * HOP_LENGTH / SAMPLE_RATE
-    alignment.start_times = time_offset + alignment.start_times
-    alignment.end_times = time_offset + alignment.end_times
-
     token_sources = np.repeat(np.arange(len(segments)), [len(s["tokens"]) for s in segments])
-    token_sources: List[int] = [None] * len(tokenizer.sot_sequence) + list(token_sources)
 
     for segment in segments:
         segment["words"] = []
 
-    word_boundaries = np.pad(np.cumsum([len(t) for t in alignment.word_tokens]), (1, 0))
-    for i, (word, start, end) in enumerate(zip(alignment.words, alignment.start_times, alignment.end_times)):
-        if word.startswith("<|") or word.strip() in ".,!?、。":  # TODO: expand
-            continue
-
-        segment = segments[token_sources[word_boundaries[i]]]
-        segment["words"].append(dict(word=word, start=round(start, 2), end=round(end, 2)))
+    word_boundaries = np.pad(np.cumsum([len(w.tokens) for w in alignment]), (1, 0))
+    for i, timing in enumerate(alignment):
+        if timing.word:
+            segment = segments[token_sources[word_boundaries[i]]]
+            start = round(time_offset + timing.start, 2)
+            end = round(time_offset + timing.end, 2)
+            segment["words"].append(
+                dict(word=timing.word, start=start, end=end, probability=timing.probability)
+            )
 
-    # adjust the segment-level timestamps based on the word-level timestamps
     for segment in segments:
-        if len(segment["words"]) > 0:
-            segment["start"] = segment["words"][0]["start"]
-            segment["end"] = segment["words"][-1]["end"]
+        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"]