Add Moonshine to KerasHub

[harshaljanjani]

Moonshine ASR Model Implementation in Keras

This PR introduces the Moonshine Automatic Speech Recognition (ASR) model into the Keras ecosystem. The Moonshine model, originally developed by UsefulSensors and available via Hugging Face, is a transformer-based architecture designed to transcribe audio inputs into text. This implementation ports the model into Keras, complete with support for pre-trained weights from Hugging Face.

Overview

The Moonshine ASR model employs an encoder-decoder architecture. The encoder processes audio features, while the decoder generates text transcriptions. This implementation includes custom layers and components to mirror the original model's behavior, validated against the Hugging Face version for accuracy.

Files Added

The following files have been added to implement the Moonshine ASR model:

• moonshine_backbone.py defines the MoonshineBackbone class, the core of the model. It integrates the encoder and decoder blocks, embeddings, and layer normalization, forming the complete encoder-decoder pipeline.

• moonshine_decoder.py contains the MoonshineDecoderBlock class, a custom decoder block with self-attention (causal), cross-attention, and feedforward layers. It supports caching for efficient generation and uses SwiGLU activation by default.

• moonshine_encoder.py implements the MoonshineEncoderBlock class, the encoder component with self-attention and feedforward layers. It optionally uses SwiGLU activation, matching the original model's configuration.

• moonshine_multi_head_attention.py provides a custom multi-head attention layer, the MoonshineMultiHeadAttention class, that is used in three ways:

  • Decoder Self-Attention: Causal, rotary embeddings, and dynamic caching (essentially, self-attention caches grow with sequence length.
  • Decoder Cross-Attention: Non-causal, no rotary embeddings, and static caching.
  • Encoder Self-Attention: Non-causal, rotary embeddings, and no caching.

• moonshine_layers.py includes utility layers, which are:

  • MoonshineRotaryEmbedding: Rotary positional embeddings with dynamic scaling support.
  • MoonshineMLP: Can be configured to use SwiGLU activation for feedforward networks or as a linear layer with GeLU activation.

• moonshine_audio_converter.py implements the MoonshineAudioConverter class, a specialized audio preprocessing layer that converts raw audio waveforms into feature representations suitable for the Moonshine ASR model. It includes downsampling and feature extraction, normalization, and handling of attention masks.

• moonshine_tokenizer.py provides the MoonshineTokenizer class, which extends the LlamaTokenizer to handle text tokenization for the Moonshine model. It incorporates Moonshine-specific special tokens, including position embedding tokens, hex tokens, and empty tokens, and manages the conversion between raw text and token IDs.

• moonshine_audio_to_text.py implements the MoonshineAudioToText class, a task model that extends the Seq2SeqLM base class. This class integrates the audio converter, backbone, and tokenizer components to create a complete end-to-end ASR pipeline. It includes methods for text generation from audio inputs, with support for customizable generation parameters and built-in trimming of output sequences.

• Weights Conversion Script

  • Converts pre-trained weights from Hugging Face into a Keras-compatible format.
  • Loads them into the MoonshineBackbone model.
  • Validates the Keras implementation by comparing outputs with the Hugging Face model using random inputs.

Dependencies

• Keras 3: Required for backend-agnostic operations.
• Hugging Face Transformers: Needed by the weights conversion script for loading the original model.
• Librosa: Required for audio processing.

Notes for Reviewers

• The implementation is fully functional with pre-trained weights and ready for immediate use.
• The modular design allows for easy extension or modification of individual components (e.g., attention layers or embeddings).
• All custom layers are serializable with get_config() and registered with @keras.saving.register_keras_serializable, ensuring compatibility with Keras model saving/loading.
• End-To-End Demo and Component Functionality Validation Notebook: Colab Notebook.

Closes issue #2083.

[divyashreepathihalli]

Thank you for the PR! I left some initial comments.
I would suggest following the format, structure and naming conventions similar to teh Whisper model here - https://github.com/keras-team/keras-hub/tree/master/keras_hub/src/models/whisper

• add docstrings
• convert backbone to a functional model
• add a moonshine_audio_converter.py
• Add a numerics verification colab to verify the implementation

[harshaljanjani]

Will make the changes at the earliest, thanks for the review!

[divyashreepathihalli]

you will need to run shell/api_gen.sh and also shell/format.sh at root to resolve the code formatting error

[harshaljanjani]

Thanks for the review, made the changes! The issue regarding the build still persists.

[harshaljanjani]

Summary of Changes:

1. Added MoonshineDecoderBlock (passes numeric checks, facing a few issues in the reversible embeddings, which keeps me from integrating the whole decoder, but I'll try to fix that and get back).
2. Made a testable component for the encoder subclassed from keras.Model separate from the MoonshineBackbone class, as it's easier to test loading weights this way since each of the preprocessor, decoder and encoder has separate weight files.

[harshaljanjani]

TODO:

1. Verify the build methods, as the sanity checks for serialization don’t pass, even though the numerics are aligned.
2. Write weight conversion scripts.

[harshaljanjani]

Status of the PR:
Weight assignment works, but the numerics differ.

Outputs of the convert_moonshine_checkpoints.py script:

MD5 Checksum Comparison
Decoder Weights Assignment
Preprocessor Weights Assignment
Encoder Weights Assignment

[mattdangerw]

What does librosa buy us? We definitely can't add it as a hard dependency. Most people using KerasHub will not be using it for audio modeling today, so adding a hard dep for them will just create headaches.

However, we could add a conditional import of librosa in our API -- attempt to use moonshine in KerasHub, get an error message asking for librosa. And of course if we just need this for data preparation outside of our API, that is easiest, just use librosa in any guides and example (but ultimately leave it up to the user).

[harshaljanjani]

@mattdangerw, that sounds about right, a conditional import seems like the way to go.
Additionally, I've written a generate() method as part of a MoonshineForConditionalGeneration task model. It looks something like this, I'll be updating the PR and the Colab Notebook soon with results from the end-to-end example using the task model.

[harshaljanjani]

I've incorporated all review feedback and completed the following updates:

1. Added component-wise tests.
2. Finalized all docstrings.
3. Ensured our generate() API matches Hugging Face's generate() API functionality and incorporated caching for speedup.
4. Cleaned up the codebase.
5. Added comprehensive test cases covering all features.
6. Included both the task model and an end-to-end example in the PR/Notebook.

Classes Merged/Removed:

1. MoonshinePrecomputed + MoonshineCausal + MoonshineMultiHead → MoonshineMultiHeadAttention.
2. MoonshineSwiGLU + MoonshineLinearGeLU → MoonshineMLP.
3. Removed MoonshineInvFreqInitializer and MoonshineArange while ensuring functional equivalence.

Additionally, there's a bug in Hugging Face's implementation of the tiny preset, which has been corrected in the Keras Hub version. I suspect the issue lies in the application of rotary positional embeddings in the HF implementation, in any case, the following code cell demonstrates the bug, and the notebook examples for each backend show how the Keras Hub implementation's tiny preset overcomes this bug for the same audio sample: Code cell showing HF bug for the tiny preset. I've chosen five samples, including the buggy one mentioned above, that have been tested for all backends and both presets in the notebook.

All bells and whistles are covered with the new test suites, caching, forward passes, standardized tests, training flows, and beyond.

Notebook: Colab Notebook

I've also shared the model weights for each preset so you can test them yourself, just as a Keras Hub user would (this would be the end-to-end example you asked about, @divyashreepathihalli and @mattdangerw):
Notebook Cell
Model Weights

I’d love to make a few updates based on your reviews and wrap up the project!

[mattdangerw]

@harshaljanjani Thanks! I'll take a look more closely soon, but I think at a high-level this is too close to huggingface's abstractions and not quite congruous with KerasHub's yet. We probably want to keep this most closely modeled off the Seq2Seq task we have in KerasHub today. We might want a new name for this AudioToText or Transcribe, but we would probably expect the following use cases roughly...

audio_tensor = load_audio_tensor_with_any_lib()
audio_batch = ... # With a batch dim
audio_dataset = ... # Paired audio tensors and strings as a tf.data.Dataset.

# Load model arch and preprocessing.
audio_to_text = keras_hub.models.AudioToText.from_preset(
    "moonshine_preset_name_blah"
)
# Equivalent, no auto class functionality.
audio_to_text = keras_hub.models.MoonshineAudioToText.from_preset(
    "moonshine_preset_name_blah"
)
# Direct string output!
audio_to_text.generate(audio_tensor)
# List of strings output!
audio_to_text.generate(audio_batch)

# Change the generation sampler and regenerate
audio_to_text.compile(sampler="top_k")
audio_to_text.generate(audio_tensor)
audio_to_text.compile(sampler=keras_hub.samplers.Greedy())
audio_to_text.generate(audio_tensor)

# Fine-tune with a dataset dataset!
audio_to_text.compile(optimizer=...)
audio_to_text.enable_lora(4)  # Optional.
audio_to_text.fit(audio_dataset)

# Set max sequence length for encoder and decoder inputs.
audio_to_text.preprocessor.encoder_sequence_length = 1024
audio_to_text.preprocessor.decoder_sequence_length = 512
audio_to_text.generate(audio_tensor)

# Strip preprocessing form the generated function entirely.
preprocessor = audio_to_text.preprocessor
audio_to_text.preprocessor = None
# Run preprocessing separately.
preprocessed_batch = preprocessor.generate_preprocess(audio_batch)
# Returns a token id tensor!
generated_batch = audio_to_text.generate(preprocessed_batch)
# Converts to strings!
preprocessor.generate_preprocess(generated_batch)

I'd maybe start by trying to move the generation from this huggingface port to our infra. Can you compile() your task with a keras_hub.samplers object and change the generation strategy? When we get there we are getting on the right track.

[harshaljanjani]

Please do look into it regarding the issue, thanks!

[divyashreepathihalli]

@harshaljanjani caching logic is complicated. It is out of scope for the maintainers to debug contribution code. Please take your time to debug the model to get matching outputs with that of the HF model.