word-level timestamps in transcribe()

requirements.txt

@@ -4,3 +4,4 @@ tqdm
 more-itertools
 transformers>=4.19.0
 ffmpeg-python==0.2.0
+dtw-python==1.3.0

whisper/audio.py

@@ -18,6 +18,10 @@
 N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000: number of samples in a chunk
 N_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH)  # 3000: number of frames in a mel spectrogram input
 
+N_SAMPLES_PER_TOKEN = HOP_LENGTH * 2  # the initial convolutions has stride 2
+FRAMES_PER_SECOND = exact_div(SAMPLE_RATE, HOP_LENGTH)  # 100 mel frames in 1s (10ms each)
+TOKENS_PER_SECOND = exact_div(SAMPLE_RATE, N_SAMPLES_PER_TOKEN)  # 50 audio tokens in 1s (20ms each)
+
 
 def load_audio(file: str, sr: int = SAMPLE_RATE):
     """

whisper/tokenizer.py

@@ -1,4 +1,5 @@
 import os
+import string
 from dataclasses import dataclass
 from functools import lru_cache
 from typing import List, Optional, Tuple, Union
@@ -269,6 +270,39 @@ def _get_single_token_id(self, text) -> int:
         assert len(tokens) == 1, f"{text} is not encoded as a single token"
         return tokens[0]
 
+    def split_tokens_on_unicode(self, tokens: List[int]):
+        words = []
+        word_tokens = []
+        current_tokens = []
+
+        for token in tokens:
+            current_tokens.append(token)
+            decoded = self.decode_with_timestamps(current_tokens)
+            if "\ufffd" not in decoded:
+                words.append(decoded)
+                word_tokens.append(current_tokens)
+                current_tokens = []
+
+        return words, word_tokens
+
+    def split_tokens_on_spaces(self, tokens: List[int]):
+        subwords, subword_tokens_list = self.split_tokens_on_unicode(tokens)
+        words = []
+        word_tokens = []
+
+        for subword, subword_tokens in zip(subwords, subword_tokens_list):
+            special = subword_tokens[0] >= self.eot
+            with_space = subword.startswith(" ")
+            punctuation = subword.strip() in string.punctuation
+            if special or with_space or punctuation:
+                words.append(subword)
+                word_tokens.append(subword_tokens)
+            else:
+                words[-1] = words[-1] + subword
+                word_tokens[-1].extend(subword_tokens)
+
+        return words, word_tokens
+
 
 @lru_cache(maxsize=None)
 def build_tokenizer(name: str = "gpt2"):

whisper/transcribe.py

@@ -8,15 +8,104 @@
 import torch
 import tqdm
 
-from .audio import SAMPLE_RATE, N_FRAMES, HOP_LENGTH, pad_or_trim, log_mel_spectrogram
+from .audio import HOP_LENGTH, N_FRAMES, SAMPLE_RATE, FRAMES_PER_SECOND, TOKENS_PER_SECOND, log_mel_spectrogram, pad_or_trim
 from .decoding import DecodingOptions, DecodingResult
-from .tokenizer import LANGUAGES, TO_LANGUAGE_CODE, get_tokenizer
+from .tokenizer import LANGUAGES, TO_LANGUAGE_CODE, Tokenizer, get_tokenizer
 from .utils import exact_div, format_timestamp, optional_int, optional_float, str2bool, write_txt, write_vtt, write_srt
 
 if TYPE_CHECKING:
     from .model import Whisper
 
 
+def add_word_timestamps(
+    model: "Whisper",
+    tokenizer: Tokenizer,
+    mel: torch.Tensor,
+    num_frames: int,
+    segments: List[dict],
+    *,
+    medfilt_width: int = 7,
+    qk_scale: float = 1.0,
+):
+    if len(segments) == 0:
+        return
+
+    from dtw import dtw
+    from scipy.ndimage import median_filter
+
+    # 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)
+    ]
+
+    tokens = torch.tensor(
+        [
+            *tokenizer.sot_sequence,
+            tokenizer.timestamp_begin,
+            *[t for segment in segments for t in segment["tokens"]],
+            tokenizer.timestamp_begin + mel.shape[-1] // 2,
+            tokenizer.eot,
+        ]
+    ).to(model.device)
+
+    with torch.no_grad():
+        model(mel.unsqueeze(0), tokens.unsqueeze(0))
+
+    for hook in hooks:
+        hook.remove()
+
+    weights = torch.cat(QKs)  # layers * heads * tokens * frames
+    weights = weights[:, :, :, : num_frames // 2].cpu()
+    weights = median_filter(weights, (1, 1, 1, medfilt_width))
+    weights = torch.tensor(weights * qk_scale).softmax(dim=-1)
+
+    w = weights / weights.norm(dim=-2, keepdim=True)
+    matrix = w.mean(axis=(0, 1))
+
+    alignment = dtw(-matrix.double().numpy())
+
+    jumps = np.pad(np.diff(alignment.index1s), (1, 0), constant_values=1).astype(bool)
+    jump_times = alignment.index2s[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
+        split_tokens = tokenizer.split_tokens_on_unicode
+    else:
+        split_tokens = tokenizer.split_tokens_on_spaces
+
+    words, word_tokens = split_tokens(tokens[1:].tolist())
+
+    token_sources = np.repeat(np.arange(len(segments)), [len(s["tokens"]) for s in segments])
+    token_sources = [None] * len(tokenizer.sot_sequence) + list(token_sources)
+
+    time_offset = segments[0]["seek"] * HOP_LENGTH / SAMPLE_RATE
+    word_boundaries = np.pad(np.cumsum([len(t) for t in word_tokens]), (1, 0))
+    start_times = time_offset + jump_times[word_boundaries[:-1]]
+    end_times = time_offset + jump_times[word_boundaries[1:]]
+
+    for segment in segments:
+        segment["words"] = []
+
+    for i, (word, start, end) in enumerate(zip(words, start_times, end_times)):
+        if word.startswith("<|") or word.strip() in ".,!?、。":
+            continue
+
+        segment = segments[token_sources[word_boundaries[i]]]
+        segment["words"].append(dict(word=word, start=round(start, 2), end=round(end, 2)))
+
+    # 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"]
+
+
 def transcribe(
     model: "Whisper",
     audio: Union[str, np.ndarray, torch.Tensor],
@@ -27,6 +116,7 @@ def transcribe(
     logprob_threshold: Optional[float] = -1.0,
     no_speech_threshold: Optional[float] = 0.6,
     condition_on_previous_text: bool = True,
+    word_level_timestamps: bool = False,
     **decode_options,
 ):
     """
@@ -90,14 +180,14 @@ def transcribe(
         else:
             if verbose:
                 print("Detecting language using up to the first 30 seconds. Use `--language` to specify the language")
-            segment = pad_or_trim(mel, N_FRAMES).to(model.device).to(dtype)
-            _, probs = model.detect_language(segment)
+            mel_segment = pad_or_trim(mel, N_FRAMES).to(model.device).to(dtype)
+            _, probs = model.detect_language(mel_segment)
             decode_options["language"] = max(probs, key=probs.get)
             if verbose is not None:
                 print(f"Detected language: {LANGUAGES[decode_options['language']].title()}")
 
-    language = decode_options["language"]
-    task = decode_options.get("task", "transcribe")
+    language: str = decode_options["language"]
+    task: str = decode_options.get("task", "transcribe")
     tokenizer = get_tokenizer(model.is_multilingual, language=language, task=task)
 
     def decode_with_fallback(segment: torch.Tensor) -> DecodingResult:
@@ -147,7 +237,8 @@ def decode_with_fallback(segment: torch.Tensor) -> DecodingResult:
     def add_segment(
         *, start: float, end: float, text_tokens: torch.Tensor, result: DecodingResult
     ):
-        text = tokenizer.decode([token for token in text_tokens if token < tokenizer.eot])
+        text_tokens = [token for token in text_tokens.tolist() if token < tokenizer.eot]
+        text = tokenizer.decode(text_tokens)
         if len(text.strip()) == 0:  # skip empty text output
             return
 
@@ -158,32 +249,28 @@ def add_segment(
                 "start": start,
                 "end": end,
                 "text": text,
-                "tokens": text_tokens.tolist(),
+                "tokens": text_tokens,
                 "temperature": result.temperature,
                 "avg_logprob": result.avg_logprob,
                 "compression_ratio": result.compression_ratio,
                 "no_speech_prob": result.no_speech_prob,
             }
         )
-        if verbose:
-            line = f"[{format_timestamp(start)} --> {format_timestamp(end)}] {text}\n"
-            # compared to just `print(line)`, this replaces any character not representable using
-            # the system default encoding with an '?', avoiding UnicodeEncodeError.
-            sys.stdout.buffer.write(line.encode(sys.getdefaultencoding(), errors="replace"))
-            sys.stdout.flush()
 
     # show the progress bar when verbose is False (otherwise the transcribed text will be printed)
     num_frames = mel.shape[-1]
-    previous_seek_value = seek
+    previous_seek = seek
 
     with tqdm.tqdm(total=num_frames, unit='frames', disable=verbose is not False) as pbar:
         while seek < num_frames:
-            timestamp_offset = float(seek * HOP_LENGTH / SAMPLE_RATE)
-            segment = pad_or_trim(mel[:, seek:], N_FRAMES).to(model.device).to(dtype)
-            segment_duration = segment.shape[-1] * HOP_LENGTH / SAMPLE_RATE
+            time_offset = float(seek * HOP_LENGTH / SAMPLE_RATE)
+            mel_segment = mel[:, seek:]
+            segment_size = min(mel_segment.shape[-1], N_FRAMES)
+            segment_duration = segment_size * HOP_LENGTH / SAMPLE_RATE
+            mel_segment = pad_or_trim(mel_segment, N_FRAMES).to(model.device).to(dtype)
 
             decode_options["prompt"] = all_tokens[prompt_reset_since:]
-            result: DecodingResult = decode_with_fallback(segment)
+            result: DecodingResult = decode_with_fallback(mel_segment)
             tokens = torch.tensor(result.tokens)
 
             if no_speech_threshold is not None:
@@ -194,9 +281,10 @@ def add_segment(
                     should_skip = False
 
                 if should_skip:
-                    seek += segment.shape[-1]  # fast-forward to the next segment boundary
+                    seek += segment_size  # fast-forward to the next segment boundary
                     continue
 
+            last_segment_index = len(all_segments)
             timestamp_tokens: torch.Tensor = tokens.ge(tokenizer.timestamp_begin)
             consecutive = torch.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0].add_(1)
             if len(consecutive) > 0:  # if the output contains two consecutive timestamp tokens
@@ -210,8 +298,8 @@ def add_segment(
                         sliced_tokens[-1].item() - tokenizer.timestamp_begin
                     )
                     add_segment(
-                        start=timestamp_offset + start_timestamp_position * time_precision,
-                        end=timestamp_offset + end_timestamp_position * time_precision,
+                        start=time_offset + start_timestamp_position * time_precision,
+                        end=time_offset + end_timestamp_position * time_precision,
                         text_tokens=sliced_tokens[1:-1],
                         result=result,
                     )
@@ -231,22 +319,45 @@ def add_segment(
                     duration = last_timestamp_position * time_precision
 
                 add_segment(
-                    start=timestamp_offset,
-                    end=timestamp_offset + duration,
+                    start=time_offset,
+                    end=time_offset + duration,
                     text_tokens=tokens,
                     result=result,
                 )
 
-                seek += segment.shape[-1]
+                seek += segment_size
                 all_tokens.extend(tokens.tolist())
 
             if not condition_on_previous_text or result.temperature > 0.5:
                 # do not feed the prompt tokens if a high temperature was used
                 prompt_reset_since = len(all_tokens)
 
+            if word_level_timestamps:
+                current_segments = all_segments[last_segment_index:]
+                add_word_timestamps(
+                    model,
+                    tokenizer,
+                    mel=mel_segment,
+                    num_frames=segment_size,
+                    segments=current_segments,
+                )
+                word_end_timestamps = [w["end"] for s in current_segments for w in s["words"]]
+                if len(word_end_timestamps) > 0:
+                    seek_shift = (word_end_timestamps[-1] - time_offset) * FRAMES_PER_SECOND
+                    seek = previous_seek + round(seek_shift)
+
+            if verbose:
+                for segment in all_segments[last_segment_index:]:
+                    start, end, text = segment["start"], segment["end"], segment["text"]
+                    line = f"[{format_timestamp(start)} --> {format_timestamp(end)}] {text}\n"
+                    # compared to just `print(line)`, this replaces any character not representable using
+                    # the system default encoding with an '?', avoiding UnicodeEncodeError.
+                    sys.stdout.buffer.write(line.encode(sys.getdefaultencoding(), errors="replace"))
+                    sys.stdout.flush()
+
             # update progress bar
-            pbar.update(min(num_frames, seek) - previous_seek_value)
-            previous_seek_value = seek
+            pbar.update(min(num_frames, seek) - previous_seek)
+            previous_seek = seek
 
     return dict(text=tokenizer.decode(all_tokens[len(initial_prompt):]), segments=all_segments, language=language)
 
@@ -280,6 +391,7 @@ def cli():
     parser.add_argument("--compression_ratio_threshold", type=optional_float, default=2.4, help="if the gzip compression ratio is higher than this value, treat the decoding as failed")
     parser.add_argument("--logprob_threshold", type=optional_float, default=-1.0, help="if the average log probability is lower than this value, treat the decoding as failed")
     parser.add_argument("--no_speech_threshold", type=optional_float, default=0.6, help="if the probability of the <|nospeech|> token is higher than this value AND the decoding has failed due to `logprob_threshold`, consider the segment as silence")
+    parser.add_argument("--word_level_timestamps", type=str2bool, default=False, help="Extract word-level timestamps and refine the results based on them")
     parser.add_argument("--threads", type=optional_int, default=0, help="number of threads used by torch for CPU inference; supercedes MKL_NUM_THREADS/OMP_NUM_THREADS")
 
     args = parser.parse_args().__dict__