Skip to content

Export Parakeet TDT models to QNN - #3719

Merged
csukuangfj merged 9 commits into
k2-fsa:masterfrom
csukuangfj:export-parakeet-tdt-qnn
Jul 3, 2026
Merged

csukuangfj merged 9 commits into
k2-fsa:masterfrom
csukuangfj:export-parakeet-tdt-qnn

Conversation

@csukuangfj

@csukuangfj csukuangfj commented Jul 3, 2026 •

Copy link
Copy Markdown
Collaborator

Specifically, the following two models are supported:

Other models may also be converted to QNN using scripts in this PR.

C++ runtime will be added in a separate pull request.

Please see

Summary by CodeRabbit

  • New Features

    • Added support for exporting and packaging Parakeet-TDT model artifacts for multiple SoCs.
    • Added a new workflow to automate model export, conversion, and release preparation.
    • Added a new command-line entrypoint for running the Parakeet-TDT export flow.
  • Bug Fixes

    • Enabled generation of additional intermediate export files needed for downstream conversion and validation.
    • Clarified input expectations for one model component to improve consistency during export.

@dosubot dosubot Bot added the size:L This PR changes 100-499 lines, ignoring generated files. label Jul 3, 2026

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

Copy link
Copy Markdown

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 configuration generation and ONNX export scripts for the Parakeet TDT model, including enabling decoder quantization in test scripts. Feedback on these changes includes wrapping the pip package name in quotes within the shell script to prevent globbing issues, and using the length of the vocabulary instead of a loop variable to safely write the blank token index.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

set -ex

pip install \
nemo_toolkit['asr'] \

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

medium

Unquoted square brackets in shell scripts can trigger globbing/pattern matching in some shells (like zsh), leading to errors like no matches found. It is safer to wrap the package name in double quotes.

Suggested change
nemo_toolkit['asr'] \
"nemo_toolkit[asr]" \

with open("./tokens.txt", "w", encoding="utf-8") as f:
for i, s in enumerate(asr_model.joint.vocabulary):
f.write(f"{s} {i}\n")
f.write(f"<blk> {i+1}\n")

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

medium

Relying on the loop variable i outside of the for loop is a code smell and can raise a NameError if the vocabulary is empty. Using len(asr_model.joint.vocabulary) is safer and more explicit.

Suggested change
f.write(f"<blk> {i+1}\n")
f.write(f"<blk> {len(asr_model.joint.vocabulary)}\n")

@coderabbitai

coderabbitai Bot commented Jul 3, 2026 •

Copy link
Copy Markdown

Review Change Stack

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 703bdcbc-7010-4ca6-b9c7-37bab935b848

📥 Commits

Reviewing files that changed from the base of the PR and between e5cece2 and 3e1d9f4.

📒 Files selected for processing (7)
  • .github/scripts/export-qnn/generate_parakeet_tdt.py
  • .github/workflows/export-parakeet-tdt-qnn.yaml
  • scripts/nemo/qnn/parakeet-tdt-ctc/transducer/test_onnx.py
  • scripts/nemo/qnn/parakeet-tdt-ctc/transducer/wrapper.py
  • scripts/nemo/qnn/parakeet-tdt/run.sh
  • scripts/nemo/qnn/parakeet-tdt/test_onnx.py
  • scripts/nemo/qnn/parakeet-tdt/wrapper.py

📝 Walkthrough

Walkthrough

Adds a new Parakeet-TDT ONNX export wrapper, a run script, a symlinked test entrypoint, and a build-matrix generator script, plus a new GitHub Actions workflow that exports ONNX models and converts them to QNN artifacts across multiple SoCs. Also makes minor edits to shared transducer test/wrapper code.

Changes

Parakeet-TDT export and QNN workflow

Layer / File(s) Summary
ONNX export wrapper
scripts/nemo/qnn/parakeet-tdt/wrapper.py
New script with CLI args, EncoderWrapper/DecoderWrapper/JoinerWrapper classes, export helper functions, and main() that loads the NeMo ASR model, writes tokens.txt, and exports encoder/decoder/joiner ONNX graphs.
Run script and test entrypoint
scripts/nemo/qnn/parakeet-tdt/run.sh, scripts/nemo/qnn/parakeet-tdt/test_onnx.py
New run.sh installs dependencies and invokes wrapper.py with configurable max-len/model-id; test_onnx.py is replaced with a symlink to the shared transducer test script.
Shared transducer tweaks
scripts/nemo/qnn/parakeet-tdt-ctc/transducer/test_onnx.py, .../wrapper.py
Enables previously disabled decoder/joiner raw-input export, removes a stale comment, and clarifies a docstring for the decoder's y input tensor.
Build matrix generator
.github/scripts/export-qnn/generate_parakeet_tdt.py
New script defines a Config dataclass and generates a JSON build matrix across SoCs, model names, and durations.
CI export/convert workflow
.github/workflows/export-parakeet-tdt-qnn.yaml
New workflow with onnx, generate_build_matrix, and qnn jobs that export ONNX artifacts, generate a build matrix, and convert/quantize/build/package/release QNN artifacts per SoC.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
    participant CI as GitHub Actions Trigger
    participant OnnxJob as onnx Job
    participant MatrixJob as generate_build_matrix Job
    participant QnnJob as qnn Job
    participant QnnSdk as QNN Toolkit

    CI->>OnnxJob: Run wrapper.py export (encoder/decoder/joiner)
    OnnxJob->>OnnxJob: Package ONNX artifacts into tarball
    OnnxJob-->>CI: Upload ONNX artifact

    CI->>MatrixJob: Run generate_parakeet_tdt.py
    MatrixJob-->>CI: Export build matrix JSON

    CI->>QnnJob: Start with matrix + depends on onnx, generate_build_matrix
    QnnJob->>QnnJob: Download ONNX artifact
    QnnJob->>QnnSdk: Convert encoder/decoder/joiner via qnn-onnx-converter
    QnnSdk-->>QnnJob: Quantized model libs
    QnnJob->>QnnSdk: Build context binaries
    QnnSdk-->>QnnJob: .bin outputs
    QnnJob->>QnnJob: Package per-SOC tarballs
    QnnJob-->>CI: Upload quant JSON + optional release assets
Loading

Possibly related PRs

  • k2-fsa/sherpa-onnx#3693: Enables the same previously disabled raw decoder/joiner input export block in scripts/nemo/qnn/parakeet-tdt-ctc/transducer/test_onnx.py modified by this PR.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@csukuangfj
csukuangfj merged commit d250d04 into k2-fsa:master Jul 3, 2026
1 check was pending
@csukuangfj
csukuangfj deleted the export-parakeet-tdt-qnn branch July 3, 2026 09:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant