run formatter

Author: mm65x

mlx_audio/stt/models/mms/mms.py

@@ -7,10 +7,7 @@
 import mlx.nn as nn
 
 from ..base import STTOutput
-from ..wav2vec.wav2vec import (
-    ModelConfig,
-    Wav2Vec2Model,
-)
+from ..wav2vec.wav2vec import ModelConfig, Wav2Vec2Model
 
 
 class Model(nn.Module):
@@ -74,6 +71,7 @@ def generate(
 
         if isinstance(audio, (str, Path)):
             from mlx_audio.stt.utils import load_audio
+
             audio = load_audio(str(audio), sr=self.sample_rate, dtype=dtype)
         elif not isinstance(audio, mx.array):
             audio = mx.array(audio)
@@ -84,7 +82,9 @@ def generate(
         if audio.dtype != dtype:
             audio = audio.astype(dtype)
 
-        audio = (audio - mx.mean(audio, axis=-1, keepdims=True)) / (mx.std(audio, axis=-1, keepdims=True) + 1e-7)
+        audio = (audio - mx.mean(audio, axis=-1, keepdims=True)) / (
+            mx.std(audio, axis=-1, keepdims=True) + 1e-7
+        )
 
         logits = self(audio)
         mx.eval(logits)
@@ -147,13 +147,16 @@ def post_load_hook(cls, model: "Model", model_path: Path) -> "Model":
             with open(vocab_path) as f:
                 vocab = json.load(f)
             if isinstance(next(iter(vocab.values())), dict):
-                lang_vocab = vocab.get("eng", vocab.get("en", next(iter(vocab.values()))))
+                lang_vocab = vocab.get(
+                    "eng", vocab.get("en", next(iter(vocab.values())))
+                )
                 model._vocab = {v: k for k, v in lang_vocab.items()}
             else:
                 model._vocab = {v: k for k, v in vocab.items()}
 
         try:
             from transformers import AutoProcessor
+
             model._processor = AutoProcessor.from_pretrained(str(model_path))
         except Exception:
             model._processor = None

mlx_audio/stt/models/mms/tests/test_mms.py

@@ -85,12 +85,16 @@ def test_keeps_lm_head(self):
         self.assertIn("lm_head.weight", sanitized)
 
     def test_keeps_wav2vec2_prefix(self):
-        weights = {"wav2vec2.encoder.layers.0.attention.q_proj.weight": mx.zeros((32, 32))}
+        weights = {
+            "wav2vec2.encoder.layers.0.attention.q_proj.weight": mx.zeros((32, 32))
+        }
         sanitized = self.model.sanitize(weights)
         self.assertIn("wav2vec2.encoder.layers.0.attention.q_proj.weight", sanitized)
 
     def test_conv_transpose(self):
-        weights = {"wav2vec2.feature_extractor.conv_layers.0.conv.weight": mx.zeros((16, 1, 4))}
+        weights = {
+            "wav2vec2.feature_extractor.conv_layers.0.conv.weight": mx.zeros((16, 1, 4))
+        }
         sanitized = self.model.sanitize(weights)
         key = "wav2vec2.feature_extractor.conv_layers.0.conv.weight"
         self.assertEqual(sanitized[key].shape, (16, 4, 1))