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

mixtral runs a bit faster

Browse Source
This commit is contained in:
Awni Hannun
2023-12-12 08:36:40 -08:00
parent e42682dced
commit 2ffd0da009
5 changed files with 24 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mistral/mistral.py
View File

@@ -196,7 +196,7 @@ def load_model(folder: str, dtype=mx.float16):
config = json.loads(f.read())
config.pop("sliding_window")
model_args = ModelArgs(**config)
weights = mx.load(str(model_path / "mlx_mistral_7b.npz"))
weights = mx.load(str(model_path / "weights.npz"))
weights = tree_unflatten(list(weights.items()))
weights = tree_map(lambda p: p.astype(dtype), weights)
model = Mistral(model_args)