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

@@ -21,6 +21,7 @@
 from tqdm import tqdm
 
 _VALID_DTYPES = {"float16", "float32"}
+_WEIGHTS_FILE_NAME = "weights.safetensors"
 
 _MODELS = {
     "tiny.en": "https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt",
@@ -109,6 +110,11 @@ def available_models() -> List[str]:
     return list(_MODELS.keys())
 
 
+def save_weights(mlx_path: Path, weights) -> None:
+    """Save converted weights using the filename supported by released clients."""
+    mx.save_safetensors(str(mlx_path / _WEIGHTS_FILE_NAME), weights)
+
+
 def hf_to_pt(weights, config):
     config = {
         "n_mels": config["num_mel_bins"],
@@ -382,7 +388,7 @@ def quantize(weights, config, args):
 
     # Save weights
     print("[INFO] Saving")
-    mx.save_safetensors(str(mlx_path / "model.safetensors"), weights)
+    save_weights(mlx_path, 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
@@ -13,7 +14,7 @@
 import mlx_whisper.load_models as load_models
 import numpy as np
 import torch
-from convert import convert, load_torch_model, quantize
+from convert import convert, load_torch_model, quantize, save_weights
 from mlx.utils import tree_flatten
 
 MODEL_NAME = "tiny"
@@ -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_save_weights_uses_released_client_filename(self):
+        weights = dict(tree_flatten(self.model.parameters()))
+        with tempfile.TemporaryDirectory() as temp_dir:
+            model_dir = Path(temp_dir)
+            save_weights(model_dir, weights)
+            with open(model_dir / "config.json", "w") as f:
+                config = asdict(self.model.dims)
+                config["model_type"] = "whisper"
+                json.dump(config, f)
+
+            weights_path = model_dir / "weights.safetensors"
+            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):