mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-25 01:41:19 +08:00
format
This commit is contained in:
parent
e434e7e5c2
commit
b863e7cca0
@ -315,7 +315,7 @@ def load_model(model_path):
|
|||||||
config["hidden_dim"] = weights["layers.0.feed_forward.w1.weight"].shape[0]
|
config["hidden_dim"] = weights["layers.0.feed_forward.w1.weight"].shape[0]
|
||||||
if config.get("vocab_size", -1) < 0:
|
if config.get("vocab_size", -1) < 0:
|
||||||
config["vocab_size"] = weights["output.weight"].shape[-1]
|
config["vocab_size"] = weights["output.weight"].shape[-1]
|
||||||
unused = ["multiple_of", "ffn_dim_multiplier", 'rope_theta']
|
unused = ["multiple_of", "ffn_dim_multiplier", "rope_theta"]
|
||||||
for k in unused:
|
for k in unused:
|
||||||
if k in config:
|
if k in config:
|
||||||
config.pop(k)
|
config.pop(k)
|
||||||
|
@ -62,6 +62,12 @@ For more options including how to prompt the model, run:
|
|||||||
python mixtral.py --help
|
python mixtral.py --help
|
||||||
```
|
```
|
||||||
|
|
||||||
[^mixtral]: Refer to Mistral's [blog post](https://mistral.ai/news/mixtral-of-experts/) for more details.
|
For the Instruction model, make sure to follow the prompt format:
|
||||||
|
|
||||||
|
```
|
||||||
|
[INST] Instruction prompt [/INST]
|
||||||
|
```
|
||||||
|
|
||||||
|
[^mixtral]: Refer to Mistral's [blog post](https://mistral.ai/news/mixtral-of-experts/) and the [Hugging Face blog post](https://huggingface.co/blog/mixtral) for more details.
|
||||||
[^instruc]: Refer to the [Hugging Face repo](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) for more
|
[^instruc]: Refer to the [Hugging Face repo](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) for more
|
||||||
details
|
details
|
||||||
|
@ -1,6 +1,7 @@
|
|||||||
from transformers import AutoModelForCausalLM
|
from transformers import AutoModelForCausalLM
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
def replace_key(key: str) -> str:
|
def replace_key(key: str) -> str:
|
||||||
if "wte.weight" in key:
|
if "wte.weight" in key:
|
||||||
key = "wte.weight"
|
key = "wte.weight"
|
||||||
|
@ -65,7 +65,6 @@ class TestWhisper(unittest.TestCase):
|
|||||||
logits = mlx_model(mels, tokens)
|
logits = mlx_model(mels, tokens)
|
||||||
self.assertEqual(logits.dtype, mx.float16)
|
self.assertEqual(logits.dtype, mx.float16)
|
||||||
|
|
||||||
|
|
||||||
def test_decode_lang(self):
|
def test_decode_lang(self):
|
||||||
options = decoding.DecodingOptions(task="lang_id", fp16=False)
|
options = decoding.DecodingOptions(task="lang_id", fp16=False)
|
||||||
result = decoding.decode(self.model, self.mels, options)
|
result = decoding.decode(self.model, self.mels, options)
|
||||||
|
@ -44,7 +44,7 @@ _ALIGNMENT_HEADS = {
|
|||||||
"large-v1": b"ABzY8r9j$a0{>%R7#4sLmoOs{s)o3~84-RPdcFk!JR<kSfC2yj",
|
"large-v1": b"ABzY8r9j$a0{>%R7#4sLmoOs{s)o3~84-RPdcFk!JR<kSfC2yj",
|
||||||
"large-v2": b"ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj",
|
"large-v2": b"ABzY8zd+h!0{>%R7=D0pU<_bnWW*tkYAhobTNnu$jnkEkXqp)j;w1Tzk)UH3X%SZd&fFZ2fC2yj",
|
||||||
"large-v3": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00",
|
"large-v3": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00",
|
||||||
"large": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00"
|
"large": b"ABzY8gWO1E0{>%R7(9S+Kn!D~%ngiGaR?*L!iJG9p-nab0JQ=-{D1-g00",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@ -166,7 +166,8 @@ def convert(model, rules=None):
|
|||||||
|
|
||||||
|
|
||||||
def torch_to_mlx(
|
def torch_to_mlx(
|
||||||
torch_model: torch_whisper.Whisper, dtype: mx.Dtype = mx.float16,
|
torch_model: torch_whisper.Whisper,
|
||||||
|
dtype: mx.Dtype = mx.float16,
|
||||||
) -> whisper.Whisper:
|
) -> whisper.Whisper:
|
||||||
def convert_rblock(model, rules):
|
def convert_rblock(model, rules):
|
||||||
children = dict(model.named_children())
|
children = dict(model.named_children())
|
||||||
|
@ -37,6 +37,7 @@ def sinusoids(length, channels, max_timescale=10000):
|
|||||||
scaled_time = mx.arange(length)[:, None] * inv_timescales[None, :]
|
scaled_time = mx.arange(length)[:, None] * inv_timescales[None, :]
|
||||||
return mx.concatenate([mx.sin(scaled_time), mx.cos(scaled_time)], axis=1)
|
return mx.concatenate([mx.sin(scaled_time), mx.cos(scaled_time)], axis=1)
|
||||||
|
|
||||||
|
|
||||||
class LayerNorm(nn.LayerNorm):
|
class LayerNorm(nn.LayerNorm):
|
||||||
def __call__(self, x: mx.array) -> mx.array:
|
def __call__(self, x: mx.array) -> mx.array:
|
||||||
return super().__call__(x.astype(mx.float32)).astype(x.dtype)
|
return super().__call__(x.astype(mx.float32)).astype(x.dtype)
|
||||||
@ -123,7 +124,13 @@ class ResidualAttentionBlock(nn.Module):
|
|||||||
|
|
||||||
class AudioEncoder(nn.Module):
|
class AudioEncoder(nn.Module):
|
||||||
def __init__(
|
def __init__(
|
||||||
self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int, dtype: mx.Dtype = mx.float16,
|
self,
|
||||||
|
n_mels: int,
|
||||||
|
n_ctx: int,
|
||||||
|
n_state: int,
|
||||||
|
n_head: int,
|
||||||
|
n_layer: int,
|
||||||
|
dtype: mx.Dtype = mx.float16,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.conv1 = nn.Conv1d(n_mels, n_state, kernel_size=3, padding=1)
|
self.conv1 = nn.Conv1d(n_mels, n_state, kernel_size=3, padding=1)
|
||||||
@ -148,7 +155,13 @@ class AudioEncoder(nn.Module):
|
|||||||
|
|
||||||
class TextDecoder(nn.Module):
|
class TextDecoder(nn.Module):
|
||||||
def __init__(
|
def __init__(
|
||||||
self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int, dtype: mx.Dtype = mx.float16,
|
self,
|
||||||
|
n_vocab: int,
|
||||||
|
n_ctx: int,
|
||||||
|
n_state: int,
|
||||||
|
n_head: int,
|
||||||
|
n_layer: int,
|
||||||
|
dtype: mx.Dtype = mx.float16,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
|
||||||
@ -160,7 +173,9 @@ class TextDecoder(nn.Module):
|
|||||||
for _ in range(n_layer)
|
for _ in range(n_layer)
|
||||||
]
|
]
|
||||||
self.ln = LayerNorm(n_state)
|
self.ln = LayerNorm(n_state)
|
||||||
self._mask = nn.MultiHeadAttention.create_additive_causal_mask(n_ctx).astype(dtype)
|
self._mask = nn.MultiHeadAttention.create_additive_causal_mask(n_ctx).astype(
|
||||||
|
dtype
|
||||||
|
)
|
||||||
|
|
||||||
def __call__(self, x, xa, kv_cache=None):
|
def __call__(self, x, xa, kv_cache=None):
|
||||||
"""
|
"""
|
||||||
|
Loading…
Reference in New Issue
Block a user