Skip to content

Add Qwen3 ASR to Keras-Hub - #2942

Open
buildwithsuhana wants to merge 21 commits into
keras-team:masterfrom
buildwithsuhana:feat/qwen3-asr
Open

buildwithsuhana wants to merge 21 commits into
keras-team:masterfrom
buildwithsuhana:feat/qwen3-asr

Conversation

@buildwithsuhana

@buildwithsuhana buildwithsuhana commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Fixes #2600

Description of the change

This PR implements the Qwen3-ASR model family (0.6B and 1.7B) in KerasHub. Qwen3-ASR is a Large Audio-Language Model (LALM) that leverages the Qwen3-Omni foundation for state-of-the-art multilingual speech recognition and language identification.

Colab Notebook

https://colab.research.google.com/drive/1P7hyLcYX0_132Y7JB4REpyInCiYEEq2h?resourcekey=0-8YE-PpzIJjNgNP0Xx8Y3dg&usp=sharing

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have followed the Keras Hub Model contribution guidelines in making these changes.
  • I have followed the Keras Hub API design guidelines in making these changes.
  • I have signed the Contributor License Agreement.

@buildwithsuhana
buildwithsuhana marked this pull request as draft August 8, 2026 13:59

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces the Qwen3-ASR model architecture to KerasHub, adding components for audio preprocessing, encoding, tokenization, and causal language modeling, along with a checkpoint conversion script. The code review identifies several critical issues for backend compatibility and robustness: in Qwen3ASRAudioEncoder, dynamic shape retrieval and dynamic attention mask construction must be avoided to prevent JAX JIT compilation failures, and positional embeddings should be computed dynamically to avoid PyTorch device placement issues. In Qwen3ASRCausalLMPreprocessor, a batch-safe contiguous block approach should replace the fragile ops.where logic to prevent runtime errors with varying batch lengths. Additionally, the audio converter should register its window and filters as weights in build(), and the checkpoint conversion script must remove the invalid gated MLP weight mapping to prevent crashes.

Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_causal_lm_preprocessor.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py Outdated
Comment thread tools/checkpoint_conversion/convert_qwen3_asr_checkpoints.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_converter.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_converter.py Outdated
@buildwithsuhana
buildwithsuhana force-pushed the feat/qwen3-asr branch 4 times, most recently from cb79ad5 to 045d436 Compare August 17, 2026 13:45
…ocessing

- Update Qwen3ASRTokenizer to include <|im_end|> and <|endoftext|> as end tokens.
- Add unsplittable_tokens to Qwen3ASRTokenizer for correct special token handling.
- Fix Qwen3ASRPreprocessor to avoid adding EOS token to prompt during generation.
- Override Qwen3ASRCausalLM.generate to handle multiple stop tokens automatically.
- Refactor conversion script for cleaner output comparison and faster testing.
@buildwithsuhana
buildwithsuhana force-pushed the feat/qwen3-asr branch 6 times, most recently from ee5be5f to cd2ddd9 Compare September 4, 2026 07:09
@buildwithsuhana
buildwithsuhana marked this pull request as ready for review September 4, 2026 07:35
@buildwithsuhana

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request implements the Qwen3-ASR model architecture, including the audio converter, audio encoder, backbone, causal LM, and tokenizer components. The review feedback highlights several technical improvements: correcting a method name typo in the tokenizer, optimizing audio sequence padding to avoid slow map functions, adding input validation to the causal LM for better error handling, simplifying shape manipulation in the audio converter, improving JIT-compatibility in the audio encoder, and refactoring the decoder layer to fix a bug in the base class rather than using subclassing.

Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_tokenizer.py
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_converter.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_causal_lm.py
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_converter.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_backbone.py Outdated
@buildwithsuhana
buildwithsuhana force-pushed the feat/qwen3-asr branch 2 times, most recently from 434d102 to 73e593c Compare September 4, 2026 10:12
@gemini-code-assist

Copy link
Copy Markdown
Contributor

It looks like you didn't specify a command after /gemini. I'm here to help! You can use one of the following commands to get assistance with this pull request:

  • /gemini review: Performs a full code review of the current pull request.
  • /gemini summary: Provides a summary of the current pull request.
  • /gemini help: Displays the list of available commands.

Please let me know if you would like me to perform one of these actions!

@buildwithsuhana

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces the Qwen3-ASR model to KerasHub, adding the audio converter, audio encoder, multi-modal projector, backbone, causal LM, preprocessor, tokenizer, and HuggingFace conversion utilities. The review feedback focuses on improving robustness and compatibility with JAX JIT compilation. Specifically, the reviewer recommends using static shapes instead of symbolic operations in the audio encoder, simplifying the audio scattering logic in the backbone using ops.take_along_axis to avoid dynamic arange operations, adopting backend-agnostic ops.squeeze in the preprocessor, and adding explicit input validation in the causal LM.

Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_backbone.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_preprocessor.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_causal_lm.py
@buildwithsuhana
buildwithsuhana force-pushed the feat/qwen3-asr branch 2 times, most recently from b2c29fa to 905b0d2 Compare September 4, 2026 10:28
@buildwithsuhana

Copy link
Copy Markdown
Contributor Author

/gemini review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces the Qwen3-ASR model architecture to KerasHub, adding the necessary components for multimodal speech recognition and causal language modeling. This includes the audio converter, audio encoder, multimodal projector, backbone, preprocessor, tokenizer, and causal LM task model, along with presets, conversion utilities, and tests. The review feedback focuses on improving robustness and backend compatibility, specifically by clipping indices in the audio scattering logic to prevent out-of-bounds crashes, using -1 in ops.reshape to avoid JAX/JIT tracing issues, optimizing static constants, and cleaning up overly verbose mathematical operations (AI slop) in favor of standard Python operators.

Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_backbone.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_tokenizer.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_encoder.py
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_converter.py Outdated
Comment thread keras_hub/src/models/qwen3_asr/qwen3_asr_audio_converter.py Outdated
@buildwithsuhana
buildwithsuhana force-pushed the feat/qwen3-asr branch 2 times, most recently from d2686c0 to 8afeb97 Compare September 4, 2026 11:26

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add Qwen3 - ASR to hub

1 participant