MLX text decoder - Githubissues

argmaxinc / WhisperKit

On-device Speech Recognition for Apple Silicon

https://takeargmax.com/blog/whisperkit

MIT License

3.17k stars 267 forks source link

MLX text decoder #161

Closed jkrukowski closed 3 months ago

jkrukowski commented 3 months ago

In this PR:

added MLXTextDecoder and relevant MLX layers implementations
added basic MLXTextDecoder tests

Tested with https://huggingface.co/jkrukowski/whisper-base-mlx-safetensors and https://huggingface.co/jkrukowski/whisper-tiny-mlx-safetensors models. Tested different pipelines with combinations of MLX and Core decoder / encoder.

Possible improvements:

refactor decodeText method, it's a bit messy and copy pasted from TextEncoder
add support for weights alignment in MLXTextDecoder
make weight loading more efficient, right now we basically load the full model 2 times (for encoder and decoder)
add more tests