Fix Whisper safetensors conversion compatibility

Author: Sahil Faizal

whisper/README.md

@@ -112,7 +112,7 @@ python convert.py --help
 ```
 
 By default, the conversion script will make the directory `mlx_models`
-and save the converted `weights.npz` and `config.json` there.
+and save the converted `weights.safetensors` and `config.json` there.
 
 Each time it is run, `convert.py` will overwrite any model in the provided
 path. To save different models, make sure to set `--mlx-path` to a unique

whisper/convert.py

@@ -382,7 +382,7 @@ def quantize(weights, config, args):
 
     # Save weights
     print("[INFO] Saving")
-    mx.save_safetensors(str(mlx_path / "model.safetensors"), weights)
+    mx.save_safetensors(str(mlx_path / "weights.safetensors"), weights)
 
     # Save config.json with model_type
     with open(str(mlx_path / "config.json"), "w") as f:

whisper/test.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+import tempfile
 import unittest
 from dataclasses import asdict
 from pathlib import Path
@@ -75,6 +76,28 @@ def forward_mlx(model, mels, tokens):
     return np.array(logits)
 
 
+class TestWeightSerialization(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = convert(MODEL_NAME, dtype=mx.float32)
+
+    def test_weights_safetensors_are_loadable(self):
+        weights = dict(tree_flatten(self.model.parameters()))
+        with tempfile.TemporaryDirectory() as temp_dir:
+            model_dir = Path(temp_dir)
+            weights_path = model_dir / "weights.safetensors"
+            mx.save_safetensors(str(weights_path), weights)
+            with open(model_dir / "config.json", "w") as f:
+                config = asdict(self.model.dims)
+                config["model_type"] = "whisper"
+                json.dump(config, f)
+
+            self.assertTrue(weights_path.is_file())
+            self.assertEqual(set(mx.load(str(weights_path))), set(weights))
+            loaded_model = load_models.load_model(str(model_dir))
+            self.assertEqual(loaded_model.dims, self.model.dims)
+
+
 class TestWhisper(unittest.TestCase):
     @classmethod
     def setUpClass(cls):