zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-02 05:04:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Whisper improvements (#1080)

Browse Source 
* use safetensors in whisper

* speed up decoder

* version
This commit is contained in:
Awni Hannun
2024-11-01 10:52:28 -07:00
committed by GitHub
parent 85ffd2c96a
commit 8160e0c4e5
6 changed files with 85 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
whisper/convert.py
View File

@@ -181,7 +181,7 @@ def load_torch_weights_and_config(
)
if name_or_path.endswith(".pt"):
checkpoint = torch.load(name_or_path, map_location="cpu")
checkpoint = torch.load(name_or_path, map_location="cpu", weights_only=False)
weights, config = checkpoint["model_state_dict"], checkpoint["dims"]
else:
name_or_path = Path(name_or_path)
@@ -387,7 +387,7 @@ if __name__ == "__main__":
# Save weights
print("[INFO] Saving")
np.savez(str(mlx_path / "weights.npz"), **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: