diff --git a/whisper/README.md b/whisper/README.md index cd3bc684a..799ba5af8 100644 --- a/whisper/README.md +++ b/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 diff --git a/whisper/convert.py b/whisper/convert.py index 9cc8b861b..7369fafa3 100644 --- a/whisper/convert.py +++ b/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: diff --git a/whisper/test.py b/whisper/test.py index f0acb3cd9..9d6c6ab3d 100644 --- a/whisper/test.py +++ b/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):