diff --git a/whisper/mlx_whisper/decoding.py b/whisper/mlx_whisper/decoding.py index 814dc95ca..a42b7c649 100644 --- a/whisper/mlx_whisper/decoding.py +++ b/whisper/mlx_whisper/decoding.py @@ -332,54 +332,67 @@ def __init__( self.tokenizer = tokenizer self.sample_begin = sample_begin self.max_initial_timestamp_index = max_initial_timestamp_index + self._vocab = None def apply(self, logits: mx.array, tokens: mx.array) -> mx.array: - mask = np.zeros(logits.shape, np.float32) + if self._vocab is None: + self._vocab = mx.arange(logits.shape[-1])[None, :] + vocab = self._vocab + mask = mx.zeros(logits.shape, dtype=mx.bool_) + # suppress <|notimestamps|> which is handled by without_timestamps if self.tokenizer.no_timestamps is not None: - mask[:, self.tokenizer.no_timestamps] = -np.inf + mask = mask | (vocab == self.tokenizer.no_timestamps) ## timestamps have to appear in pairs, except directly before EOT; mask logits accordingly - tokens = tokens.tolist() - for k in range(len(tokens)): - seq = tokens[k][self.sample_begin :] - last_was_timestamp = ( - len(seq) >= 1 and seq[-1] >= self.tokenizer.timestamp_begin - ) + sampled = tokens.shape[1] - self.sample_begin + if sampled >= 1: + seq = tokens[:, self.sample_begin :] + last_was_timestamp = seq[:, -1] >= self.tokenizer.timestamp_begin penultimate_was_timestamp = ( - len(seq) < 2 or seq[-2] >= self.tokenizer.timestamp_begin + mx.ones_like(last_was_timestamp) + if sampled < 2 + else seq[:, -2] >= self.tokenizer.timestamp_begin + ) + mask = mask | ( + (last_was_timestamp & penultimate_was_timestamp)[:, None] + & (vocab >= self.tokenizer.timestamp_begin) + ) + mask = mask | ( + (last_was_timestamp & ~penultimate_was_timestamp)[:, None] + & (vocab < self.tokenizer.eot) ) - if last_was_timestamp: - if penultimate_was_timestamp: # has to be non-timestamp - mask[k, self.tokenizer.timestamp_begin :] = -np.inf - else: # cannot be normal text tokens - mask[k, : self.tokenizer.eot] = -np.inf + # Preserve the legacy position-based lower bound for output parity. + # Token-value semantics are tracked separately in KT-642. + timestamp_positions = mx.where( + seq > self.tokenizer.timestamp_begin, + mx.arange(sampled)[None, :], + -1, + ) + last_timestamp = timestamp_positions.max(axis=-1) + has_timestamp = last_timestamp >= 0 + last_timestamp = last_timestamp + ( + (last_timestamp == 0) | penultimate_was_timestamp + ) + mask = mask | ( + has_timestamp[:, None] + & (vocab >= self.tokenizer.timestamp_begin) + & (vocab < last_timestamp[:, None]) + ) - timestamps = [ - i for i, v in enumerate(seq) if v > self.tokenizer.timestamp_begin - ] - if len(timestamps) > 0: - # timestamps shouldn't decrease; forbid timestamp tokens smaller than the last - # also force each segment to have a nonzero length, to prevent infinite looping - last_timestamp = timestamps[-1] - if not last_timestamp or penultimate_was_timestamp: - last_timestamp += 1 - mask[k, self.tokenizer.timestamp_begin : last_timestamp] = -np.inf - - if len(tokens[0]) == self.sample_begin: + if tokens.shape[1] == self.sample_begin: # suppress generating non-timestamp tokens at the beginning - mask[:, : self.tokenizer.timestamp_begin] = -np.inf + mask = mask | (vocab < self.tokenizer.timestamp_begin) # apply the `max_initial_timestamp` option if self.max_initial_timestamp_index is not None: last_allowed = ( self.tokenizer.timestamp_begin + self.max_initial_timestamp_index ) - mask[:, last_allowed + 1 :] = -np.inf + mask = mask | (vocab > last_allowed) # if sum of probability over timestamps is above any other token, sample timestamp - mask = mx.array(mask) logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) timestamp_logprob = logprobs[:, self.tokenizer.timestamp_begin :].logsumexp( axis=-1, keepdims=True @@ -387,12 +400,11 @@ def apply(self, logits: mx.array, tokens: mx.array) -> mx.array: max_text_token_logprob = logprobs[:, : self.tokenizer.timestamp_begin].max( axis=-1, keepdims=True ) - mask[:, : self.tokenizer.timestamp_begin] = mx.where( - timestamp_logprob > max_text_token_logprob, - -mx.inf, - mask[:, : self.tokenizer.timestamp_begin], + mask = mask | ( + (timestamp_logprob > max_text_token_logprob) + & (vocab < self.tokenizer.timestamp_begin) ) - return logits + mask + return logits + mx.where(mask, -mx.inf, 0.0) class DecodingTask: diff --git a/whisper/mlx_whisper/timing.py b/whisper/mlx_whisper/timing.py index 07b81186c..8a763e130 100644 --- a/whisper/mlx_whisper/timing.py +++ b/whisper/mlx_whisper/timing.py @@ -2,12 +2,11 @@ import itertools from dataclasses import dataclass -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, Optional import mlx.core as mx import numba import numpy as np -from scipy import signal from .audio import HOP_LENGTH, SAMPLE_RATE, TOKENS_PER_SECOND from .tokenizer import Tokenizer @@ -20,28 +19,33 @@ def median_filter(x: np.ndarray, filter_width: int): """Apply a median filter of width `filter_width` along the last dimension of `x`""" pad_width = filter_width // 2 if x.shape[-1] <= pad_width: - # F.pad requires the padding width to be smaller than the input dimension + # Reflect padding requires the padding width to be smaller than the input. return x - if (ndim := x.ndim) <= 2: - # `F.pad` does not support 1D or 2D inputs for reflect padding but supports 3D and 4D - x = x[None, None, :] - assert ( filter_width > 0 and filter_width % 2 == 1 ), "`filter_width` should be an odd number" - - x = np.pad(x, ((0, 0), (0, 0), (pad_width, pad_width)), mode="reflect") - - # todo: more efficient version in mlx - result = signal.medfilt(x.astype(np.float32), kernel_size=(1, 1, filter_width))[ - ..., pad_width:-pad_width - ] - - if ndim <= 2: - result = result[0, 0] - - return result + if filter_width == 1: + return x.astype(np.float32, copy=False) + + values = mx.array(x.astype(np.float32, copy=False)) + size = values.shape[-1] + left = mx.take( + values, + mx.array(list(range(pad_width, 0, -1))), + axis=-1, + ) + right = mx.take( + values, + mx.array(list(range(size - 2, size - pad_width - 2, -1))), + axis=-1, + ) + padded = mx.concatenate([left, values, right], axis=-1) + windows = mx.stack( + [padded[..., offset : offset + size] for offset in range(filter_width)], + axis=-1, + ) + return np.array(mx.partition(windows, pad_width, axis=-1)[..., pad_width]) @numba.jit(nopython=True) @@ -69,7 +73,7 @@ def backtrace(trace: np.ndarray): return result[::-1, :].T -@numba.jit(nopython=True, parallel=True) +@numba.jit(nopython=True, parallel=True, cache=True) def dtw_cpu(x: np.ndarray): N, M = x.shape cost = np.ones((N + 1, M + 1), dtype=np.float32) * np.inf @@ -116,6 +120,7 @@ def find_alignment( mel: mx.array, num_frames: int, *, + audio_features: Optional[mx.array] = None, medfilt_width: int = 7, qk_scale: float = 1.0, ) -> List[WordTiming]: @@ -131,14 +136,19 @@ def find_alignment( ] ) - logits, cross_qk = model.forward_with_cross_qk(mel[None, :], tokens[None, :]) + if audio_features is not None and audio_features.ndim == 2: + audio_features = audio_features[None, :] + logits, cross_qk = model.forward_with_cross_qk( + mel[None, :], + tokens[None, :], + audio_features=audio_features, + ) # consider only the logits associated with predicting text sampled_logits = logits[0][len(tokenizer.sot_sequence) : -2, : tokenizer.eot] token_probs = mx.softmax(sampled_logits, precise=True, axis=-1) text_token_probs = mx.take_along_axis( token_probs, mx.array(text_tokens)[:, None], axis=1 ).squeeze(1) - text_token_probs = np.array(text_token_probs) # heads * tokens * frames weights = mx.stack( @@ -150,6 +160,8 @@ def find_alignment( mean = mx.mean(weights, axis=-2, keepdims=True) std = mx.var(weights, axis=-2, keepdims=True, ddof=0).sqrt() weights = (weights - mean) / std + mx.eval(text_token_probs, weights) + text_token_probs = np.array(text_token_probs) weights = median_filter(np.array(weights), medfilt_width) matrix = weights.mean(axis=0) diff --git a/whisper/mlx_whisper/transcribe.py b/whisper/mlx_whisper/transcribe.py index bced16a58..831e9f12e 100644 --- a/whisper/mlx_whisper/transcribe.py +++ b/whisper/mlx_whisper/transcribe.py @@ -209,6 +209,7 @@ def decode_with_fallback(segment: mx.array) -> DecodingResult: [temperature] if isinstance(temperature, (int, float)) else temperature ) decode_result = None + audio_features = None for t in temperatures: kwargs = {**decode_options} @@ -221,7 +222,11 @@ def decode_with_fallback(segment: mx.array) -> DecodingResult: kwargs.pop("best_of", None) options = DecodingOptions(**kwargs, temperature=t) - decode_result = model.decode(segment, options) + decode_result = model.decode( + segment if audio_features is None else audio_features, + options, + ) + audio_features = decode_result.audio_features needs_fallback = False if ( @@ -418,6 +423,7 @@ def next_words_segment(segments: List[dict]) -> Optional[dict]: tokenizer=tokenizer, mel=mel_segment, num_frames=segment_size, + audio_features=result.audio_features, prepend_punctuations=prepend_punctuations, append_punctuations=append_punctuations, last_speech_timestamp=last_speech_timestamp, diff --git a/whisper/mlx_whisper/whisper.py b/whisper/mlx_whisper/whisper.py index 5c85195cd..9f3b4d490 100644 --- a/whisper/mlx_whisper/whisper.py +++ b/whisper/mlx_whisper/whisper.py @@ -247,8 +247,10 @@ def embed_audio(self, mel): def logits(self, tokens, audio_features): return self.decoder(tokens, audio_features)[0] - def forward_with_cross_qk(self, mel, tokens): - logits, _, cross_qk = self.decoder(tokens, self.encoder(mel)) + def forward_with_cross_qk(self, mel, tokens, audio_features=None): + if audio_features is None: + audio_features = self.encoder(mel) + logits, _, cross_qk = self.decoder(tokens, audio_features) return logits, cross_qk def __call__(self, mel, tokens):