format

whisper/whisper/decoding.py

@@ -112,7 +112,7 @@ class DecodingOptions:
     max_initial_timestamp: Optional[float] = 1.0
 
     # implementation details
-    fp16: bool = True # use fp16 for most of the calculation
+    fp16: bool = True  # use fp16 for most of the calculation
 
 
 @dataclass(frozen=True)

whisper/whisper/load_models.py

@@ -44,7 +44,7 @@ _ALIGNMENT_HEADS = {
     "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-v3": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00",
-    "large": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00"
+    "large": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00",
 }
 
 
@@ -166,7 +166,8 @@ def convert(model, rules=None):
 
 
 def torch_to_mlx(
-    torch_model: torch_whisper.Whisper, dtype: mx.Dtype = mx.float16,
+    torch_model: torch_whisper.Whisper,
+    dtype: mx.Dtype = mx.float16,
 ) -> whisper.Whisper:
     def convert_rblock(model, rules):
         children = dict(model.named_children())
@@ -194,6 +195,6 @@ def torch_to_mlx(
 def load_model(
     name: str,
     download_root: str = None,
-    dtype : mx.Dtype = mx.float32,
+    dtype: mx.Dtype = mx.float32,
 ) -> whisper.Whisper:
     return torch_to_mlx(load_torch_model(name, download_root), dtype)

whisper/whisper/transcribe.py

@@ -43,7 +43,7 @@ class ModelHolder:
     model_name = None
 
     @classmethod
-    def get_model(cls, model: str, dtype : mx.Dtype):
+    def get_model(cls, model: str, dtype: mx.Dtype):
         if cls.model is None or model != cls.model_name:
             cls.model = load_model(model, dtype=dtype)
             cls.model_name = model

whisper/whisper/whisper.py

@@ -37,6 +37,7 @@ def sinusoids(length, channels, max_timescale=10000):
     scaled_time = mx.arange(length)[:, None] * inv_timescales[None, :]
     return mx.concatenate([mx.sin(scaled_time), mx.cos(scaled_time)], axis=1)
 
+
 class LayerNorm(nn.LayerNorm):
     def __call__(self, x: mx.array) -> mx.array:
         return super().__call__(x.astype(mx.float32)).astype(x.dtype)
@@ -123,7 +124,13 @@ class ResidualAttentionBlock(nn.Module):
 
 class AudioEncoder(nn.Module):
     def __init__(
-        self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int, dtype: mx.Dtype = mx.float16,
+        self,
+        n_mels: int,
+        n_ctx: int,
+        n_state: int,
+        n_head: int,
+        n_layer: int,
+        dtype: mx.Dtype = mx.float16,
     ):
         super().__init__()
         self.conv1 = nn.Conv1d(n_mels, n_state, kernel_size=3, padding=1)
@@ -148,7 +155,13 @@ class AudioEncoder(nn.Module):
 
 class TextDecoder(nn.Module):
     def __init__(
-        self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int, dtype: mx.Dtype = mx.float16,
+        self,
+        n_vocab: int,
+        n_ctx: int,
+        n_state: int,
+        n_head: int,
+        n_layer: int,
+        dtype: mx.Dtype = mx.float16,
     ):
         super().__init__()
 
@@ -160,7 +173,9 @@ class TextDecoder(nn.Module):
             for _ in range(n_layer)
         ]
         self.ln = LayerNorm(n_state)
-        self._mask = nn.MultiHeadAttention.create_additive_causal_mask(n_ctx).astype(dtype)
+        self._mask = nn.MultiHeadAttention.create_additive_causal_mask(n_ctx).astype(
+            dtype
+        )
 
     def __call__(self, x, xa, kv_cache=None):
         """