Export Whisper models to QNN - #3697
Merged
Merged
Conversation
dosubot
Bot
added
the
size:M
|
Caution Review failedAn error occurred during the review process. Please try again later. 📝 WalkthroughWalkthroughAdds a new GitHub Actions workflow ( ChangesWhisper QNN Export Pipeline
Minor Workflow Fixes
Sequence Diagram(s)sequenceDiagram
participant branch as export-whisper-qnn branch
participant onnx_job as onnx job (macOS)
participant matrix_job as generate_build_matrix (Ubuntu)
participant qnn_job as qnn job (Ubuntu)
participant qnn_sdk as QNN SDK tools
participant release as sherpa-onnx releases
branch->>onnx_job: trigger per model_name matrix
onnx_job->>onnx_job: install PyTorch/Whisper, download pretrained sources
onnx_job->>onnx_job: export_onnx.py → test_onnx.py
onnx_job->>onnx_job: upload encoder/decoder/tokens/WAV artifact
branch->>matrix_job: trigger
matrix_job->>matrix_job: generate_whisper.py → JSON matrix
matrix_job-->>qnn_job: matrix (soc × model_name)
qnn_job->>onnx_job: download ONNX artifact
qnn_job->>qnn_sdk: qnn-onnx-converter (encoder + decoder)
qnn_sdk-->>qnn_job: C/graph configs
qnn_job->>qnn_sdk: qnn-model-lib-generator
qnn_sdk-->>qnn_job: .so model libraries
qnn_job->>qnn_sdk: qnn-context-binary-generator
qnn_sdk-->>qnn_job: context binaries
qnn_job->>qnn_job: package tarballs, upload JSON configs
qnn_job->>release: publish tar.bz2 (conditional)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Whisper model export and testing scripts for QNN and RKNN, including a configuration generator and updates to ONNX export and testing utilities. Key changes include exposing the argument parser in export_onnx.py, adding support for custom WAV files and distilled/special Whisper models in test_onnx.py, and fixing the KV cache update order. Feedback focuses on resolving a potential ModuleNotFoundError in generate_whisper.py by correctly setting up the system path for importing device_info, and refactoring the repetitive model-loading logic in test_onnx.py using a dictionary-based lookup to improve maintainability.
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.
Comment on lines
+4
to
+8
| import json | ||
|
|
||
| from device_info import soc_info_dict | ||
| from dataclasses import asdict, dataclass | ||
| import itertools |
There was a problem hiding this comment.
The script imports device_info directly, but device_info.py is located in scripts/qnn/device_info.py. When running this script, it will fail with a ModuleNotFoundError because scripts/qnn is not in the Python search path. Additionally, itertools is imported but never used.
We can resolve this by dynamically adding the scripts/qnn directory to sys.path before importing device_info, and removing the unused itertools import.
Suggested change
| import json | |
| from device_info import soc_info_dict | |
| from dataclasses import asdict, dataclass | |
| import itertools | |
| import json | |
| import sys | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| # Add scripts/qnn to sys.path to import device_info | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "scripts" / "qnn")) | |
| from device_info import soc_info_dict |
Comment on lines
+127
to
+207
| name = args.model | ||
| if name == "distil-medium.en": | ||
| filename = "./distil-medium-en-original-model.bin" | ||
| if not Path(filename).is_file(): | ||
| raise ValueError( | ||
| """ | ||
| Please go to https://huggingface.co/distil-whisper/distil-medium.en | ||
| to download original-model.bin | ||
| You can use the following command to do that: | ||
|
|
||
| wget -O distil-medium-en-original-model.bin https://huggingface.co/distil-whisper/distil-medium.en/resolve/main/original-model.bin | ||
| """ | ||
| ) | ||
| torch_model = whisper.load_model(filename) | ||
| elif name == "distil-large-v2": | ||
| filename = "./distil-large-v2-original-model.bin" | ||
| if not Path(filename).is_file(): | ||
| raise ValueError( | ||
| """ | ||
| Please go to https://huggingface.co/distil-whisper/distil-large-v2 | ||
| to download original-model.bin | ||
| You can use the following command to do that: | ||
|
|
||
| wget -O distil-large-v2-original-model.bin https://huggingface.co/distil-whisper/distil-large-v2/resolve/main/original-model.bin | ||
| """ | ||
| ) | ||
| torch_model = whisper.load_model(filename) | ||
| elif name == "distil-large-v3": | ||
| filename = "./distil-large-v3-original-model.bin" | ||
| if not Path(filename).is_file(): | ||
| raise ValueError( | ||
| """ | ||
| Please go to https://huggingface.co/distil-whisper/distil-large-v3-openai | ||
| to download model.bin | ||
| You can use the following command to do that: | ||
|
|
||
| wget -O distil-large-v3-original-model.bin https://huggingface.co/distil-whisper/distil-large-v3-openai/resolve/main/model.bin | ||
| """ | ||
| ) | ||
| torch_model = whisper.load_model(filename) | ||
| elif name == "distil-large-v3.5": | ||
| filename = "./distil-large-v3.5-original-model.bin" | ||
| if not Path(filename).is_file(): | ||
| raise ValueError( | ||
| """ | ||
| Please go to https://huggingface.co/distil-whisper/distil-large-v3.5-openai/ | ||
| to download model.bin | ||
| You can use the following command to do that: | ||
|
|
||
| wget -O distil-large-v3.5-original-model.bin https://huggingface.co/distil-whisper/distil-large-v3.5-openai/resolve/main/model.bin | ||
| """ | ||
| ) | ||
| torch_model = whisper.load_model(filename) | ||
| elif name == "distil-small.en": | ||
| filename = "./distil-small-en-original-model.bin" | ||
| if not Path(filename).is_file(): | ||
| raise ValueError( | ||
| """ | ||
| Please go to https://huggingface.co/distil-whisper/distil-small.en | ||
| to download original-model.bin | ||
| You can use the following command to do that: | ||
|
|
||
| wget -O distil-small-en-original-model.bin https://huggingface.co/distil-whisper/distil-small.en/resolve/main/original-model.bin | ||
| """ | ||
| ) | ||
| torch_model = whisper.load_model(filename) | ||
| elif name == "medium-aishell": | ||
| filename = "./medium-aishell.pt" | ||
| if not Path(filename).is_file(): | ||
| raise ValueError( | ||
| """ | ||
| Please go to https://huggingface.co/yuekai/icefall_asr_aishell_whisper/tree/main/exp_medium | ||
| to download whisper-medium-aishell1-epoch-10-avg-4.pt | ||
| You can use the following command to do that: | ||
|
|
||
| wget -O medium-aishell.pt https://huggingface.co/yuekai/icefall_asr_aishell_whisper/resolve/main/exp_medium/whisper-medium-aishell1-epoch-10-avg-4.pt | ||
| """ | ||
| ) | ||
| torch_model = whisper.load_model(filename) | ||
| else: | ||
| torch_model = whisper.load_model(name) |
There was a problem hiding this comment.
The if/elif/else block for loading special/distilled models is highly repetitive and contains duplicated multi-line error messages. This can be simplified significantly by using a dictionary mapping model names to their metadata (local filename, download URL, repository, and download name). This improves readability, reduces boilerplate, and makes it much easier to add new models in the future.
name = args.model
special_models = {
"distil-medium.en": {
"filename": "./distil-medium-en-original-model.bin",
"url": "https://huggingface.co/distil-whisper/distil-medium.en/resolve/main/original-model.bin",
"repo": "https://huggingface.co/distil-whisper/distil-medium.en",
"download_name": "original-model.bin",
},
"distil-large-v2": {
"filename": "./distil-large-v2-original-model.bin",
"url": "https://huggingface.co/distil-whisper/distil-large-v2/resolve/main/original-model.bin",
"repo": "https://huggingface.co/distil-whisper/distil-large-v2",
"download_name": "original-model.bin",
},
"distil-large-v3": {
"filename": "./distil-large-v3-original-model.bin",
"url": "https://huggingface.co/distil-whisper/distil-large-v3-openai/resolve/main/model.bin",
"repo": "https://huggingface.co/distil-whisper/distil-large-v3-openai",
"download_name": "model.bin",
},
"distil-large-v3.5": {
"filename": "./distil-large-v3.5-original-model.bin",
"url": "https://huggingface.co/distil-whisper/distil-large-v3.5-openai/resolve/main/model.bin",
"repo": "https://huggingface.co/distil-whisper/distil-large-v3.5-openai/",
"download_name": "model.bin",
},
"distil-small.en": {
"filename": "./distil-small-en-original-model.bin",
"url": "https://huggingface.co/distil-whisper/distil-small.en/resolve/main/original-model.bin",
"repo": "https://huggingface.co/distil-whisper/distil-small.en",
"download_name": "original-model.bin",
},
"medium-aishell": {
"filename": "./medium-aishell.pt",
"url": "https://huggingface.co/yuekai/icefall_asr_aishell_whisper/resolve/main/exp_medium/whisper-medium-aishell1-epoch-10-avg-4.pt",
"repo": "https://huggingface.co/yuekai/icefall_asr_aishell_whisper/tree/main/exp_medium",
"download_name": "whisper-medium-aishell1-epoch-10-avg-4.pt",
},
}
if name in special_models:
info = special_models[name]
filename = info["filename"]
if not Path(filename).is_file():
raise ValueError(
f"Please go to {info['repo']} to download {info['download_name']}\n"
f"You can use the following command to do that:\n\n"
f"wget -O {filename} {info['url']}"
)
torch_model = whisper.load_model(filename)
else:
torch_model = whisper.load_model(name)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
C++ runtime will be added in a separate PR.
Summary by CodeRabbit
New Features
Improvements