Add Moonshine to KerasHub

[harshaljanjani]

Adds Moonshine ASR model to KerasHub:

• Encoder-decoder ASR model from Moonshine (UsefulSensors)
• Pre-trained weights from Hugging Face
• Demo notebook
• Functionality Tests notebook (no HF)

Closes #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

[harshaljanjani]

Updated the Colab notebook with results from the latest commit. The PR is now open for review.
What's New?
The task model has been integrated across all three backends into the KerasHub infra, including the custom caching strategy used by Moonshine.

[divyashreepathihalli]

I don't see the demo notebook with the KerasHub model implemented here, I am seeing a demo from the Huggingface model in the colab
please add the demo with KH model - and verify that the outputs match with model.generate

[harshaljanjani]

I don't see the demo notebook with the KerasHub model implemented here, I am seeing a demo from the Huggingface model in the colab
please add the demo with KH model - and verify that the outputs match with model.generate

@divyashreepathihalli The outputs you see across the first three cells are the KH model outputs for four test samples for each preset, using the generate() function. I've run tools/checkpoint_conversion/convert_moonshine_checkpoints.py in each of the cells across the three backends, which both, verifies the numerics, and contains the end-to-end example.

The cell links are:

1. PyTorch Backend Output of keras_model.generate()
2. TensorFlow Backend Output of keras_model.generate()
3. JAX Backend Output of keras_model.generate()

You may also review the checkpoint conversion file to verify the same.

The HF model is only used in the last cell, where I point out a bug in the HF implementation and show how for the same sample, the KH model presets give good transcripts across all three backends. (The sample used in this test is the "Female Clear Voice (Maximum Length - 64 Sec)" one.)

[harshaljanjani]

@mattdangerw / @abheesht17 / @divyashreepathihalli Whenever you have a chance, could you please take a look at this PR and the notebook, thanks!

[mattdangerw]

Dropping a few comments. I think we need still need to get the generation here working similar to other models, make the preprocessing be actual preprocessing (no weights!). I still think a clearer high level colab with intended usage might help clarify things.

• Do some weight conversion, upload to huggingface or kaggle (doesn't matter which) on your own user.
• Make a colab that does not touch huggingface at all that shows the intended usage here.
• Try to show some of the usages here Add Moonshine to KerasHub #2093 (comment)

How much of this is working today? Have we tried running fine-tuning? That will run preprocessing via a tf.data.Dataset map, does that work?

!pip install git+https://github.com/harshaljanjani/keras-hub@moonshine

import os
os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch" with zero other changes.

import keras
import keras_hub

audio_to_text = keras_hub.models.AudioToText.from_preset(
    "hf://harshaljanjani/keras-moonshine",
)

audio_to_text.generate(audio_tensor)
audio_to_text.generate(audio_batch)

audio_to_text.compile(sampler="top_k")
audio_to_text.generate(audio_tensor)

audio_to_text.compile(...)
audio_to_text.enable_lora(4)  # Optional.
audio_to_text.fit(audio_dataset)
autio_to_text.generate(audio_batch)

[harshaljanjani]

Will check the comments out, thanks for the review @mattdangerw. I left a few replies, I'd love to hear your opinion on a few non-trivial things as mentioned in the replies; I'll proceed to make changes on the others.

How much of this is working today? Have we tried running fine-tuning? That will run preprocessing via a tf.data.Dataset map, does that work?

I haven't tested fine-tuning yet, but I'll see what I can do. Since you mentioned that the change in the generate() strategy was key, I focused on it for this round.

[harshaljanjani]

Addressed reviews - (JIT compile + dynamic shapes issue). Looking forward to guidance regarding the same, I'll try to see if I can solve it in the mean time.

[harshaljanjani]

The PR should be ready for the next round of reviews @mattdangerw. Here's the new Colab you mentioned. I've tested the functionality with dummy inputs for now; hope you don't mind! I'll check the weights upload thing and the presets once the design is approved.

1. Functionality Tests Notebook Independent From HF.
2. Same Outputs Notebook, Updated To The Current PR's Version.

[divyashreepathihalli]

The PR should be ready for the next round of reviews @mattdangerw. Here's the new Colab you mentioned. I've tested the functionality with dummy inputs for now; hope you don't mind! I'll check the weights upload thing and the presets once the design is approved.

1. Functionality Tests Notebook Independent From HF.
2. Same Outputs Notebook, Updated To The Current PR's Version.

please add demo colabs, verifications etc to PR descriptions so that it is easier to find

[harshaljanjani]

please add demo colabs, verifications etc to PR descriptions so that it is easier to find

Apologies, the end-to-end demo notebook has been linked in the PR description from the beginning. I've just linked the functionality tests I added today in the PR description!