Repository navigation
Export SenseVoice ASR models to Ascend NPU 910B - #2707
Conversation
WalkthroughAdds a GitHub Actions workflow for automated SenseVoice model export to Ascend NPU, along with supporting Python scripts for ONNX export, model wrapping, and OM format inference testing. Supports multiple frameworks and OS configurations via matrix builds. Changes
Sequence Diagram(s)sequenceDiagram
participant GHA as GitHub Actions
participant Checkout
participant Env as Environment Setup
participant Deps as Dependencies
participant Model as SenseVoice Model
participant Export as ONNX Export
participant ATC as atc Converter
participant Package as Packaging
participant Release as GitHub Releases
GHA->>Checkout: Trigger on push/dispatch
Checkout->>Env: Setup Python & Toolchain
Env->>Deps: Install numpy, onnx, torch, libs
Deps->>Model: Download pre-trained model
Model->>Export: Load & wrap model (export_onnx.py)
Export->>Export: Generate ONNX with CMVN
Export->>ATC: Export model.onnx
ATC->>ATC: Convert ONNX → Ascend OM
ATC->>Package: Create OM model artifact
Package->>Package: Compress to tar.bz2
alt Repository owner match
Package->>Release: Upload to GitHub Releases
end
Release->>GHA: Workflow complete
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes The PR introduces heterogeneous changes across CI/CD workflow and multiple Python scripts with moderate logic density (model wrapping, ONNX export, Kaldi feature extraction, inference pipeline). The workflow includes complex matrix builds and conditional logic; the scripts introduce new model handling and audio processing logic. Changes are primarily new code rather than refactoring, and scope spans several files with distinct purposes. Possibly related PRs
Suggested labels
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ 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 |
dosubot
Bot
added
the
size:L
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (10)
.github/workflows/export-sense-voice-to-ascend-npu.yaml (2)
55-61: CANN env sourcing duplicated; consider a small helper stepYou source
set_env.shand overrideLD_LIBRARY_PATHtwice. Extract into a single reusable step or composite action to reduce drift.Also applies to: 99-110, 163-174
78-87: Add SHA256 checksum verification to hardened supply chain securityThe review comment is accurate: all URLs use
/resolve/main/, which points to the repository branch HEAD and can change when the repo is updated. HuggingFace files have SHA256 checksums available via the metadata API, making verification implementable.Affected sections:
- Lines 78–80: FunASR model downloads (am.mvn, model.pt, chn_jpn_yue_eng_ko_spectok.bpe.model)
- Lines 82–86, 90–91: FunASR test files and metadata
- Lines 141–149: WSYue-ASR model downloads (sensevoice_small_yue/model.pt, am.mvn, chn_jpn_yue_eng_ko_spectok.bpe.model)
- Lines 150–153: WSYue-ASR test file loop
- Lines 156–158: WSYue-ASR metadata downloads
Recommended approach: retrieve checksums from HuggingFace Hub API or document them in a separate manifest, then add inline
sha256sum -cverification for each critical file (especially model files).scripts/sense-voice/ascend-npu/export_onnx.py (3)
37-49: ModelWrapper is defined but unused; simplify or export the wrapper intentionallyYou build
ModelWrapperreturning(part1, part2)but exportmodel.mwith a singlelogitsoutput. This leaves dead code and can confuse readers. Either export the wrapper (two outputs) or drop the wrapper and export the base model. Given test_om consumes a single logits tensor, prefer removing the wrapper.Apply this diff to simplify:
-class ModelWrapper(torch.nn.Module): - def __init__(self, m): - super().__init__() - self.m = m - - def forward(self, x, prompt): - logits = self.m(x[None], prompt)[0] - part1 = logits[:4] - part2 = logits[4:] - part1 = part1.reshape(4, 25055) - part2 = part2.reshape(x.size(0), 25055) - return part1, part2 +# Model exports a single 'logits' tensor; wrapper not needed. @@ - model = ModelWrapper(model) - model.eval() - x = torch.randn(1, 93, 560, dtype=torch.float32) @@ - torch.onnx.export( - model.m, + torch.onnx.export( + model, (x, prompt), filename, opset_version=opset_version, input_names=["x", "prompt"], - output_names=["logits"], + output_names=["logits"], dynamic_axes={ "x": {0: "N", 1: "T"}, "logits": {0: "N", 1: "T_4"}, }, )Also applies to: 83-95
59-67: State dict normalization is good; add strict=False guard for resilienceThird‑party checkpoints sometimes omit buffers or have prefix drift. Using
strict=Falseavoids brittle failures while still logging mismatches.- model.load_state_dict(state_dict) + missing, unexpected = model.load_state_dict(state_dict, strict=False) + if missing or unexpected: + print("load_state_dict warnings:", {"missing": missing, "unexpected": unexpected})
51-56: SentencePiece load error pathIf the
.bpe.modelfails to download,sp.loadreturns False without raising. Add an assert to fail fast.- sp.load("./chn_jpn_yue_eng_ko_spectok.bpe.model") + ok = sp.load("./chn_jpn_yue_eng_ko_spectok.bpe.model") + assert ok, "Failed to load SentencePiece model"scripts/sense-voice/ascend-npu/test_om.py (5)
111-116: Minor cleanups: remove unused constants; fix typo (ja) and commentsSeveral
language_*constants andwithout_itnare unused; drop or group into a dict. Alsolanguage_yashould belanguage_ja. Ensure commented lines start with#.Apply this diff:
- language_auto = 0 - language_zh = 3 - language_en = 4 - language_yue = 7 - language_ya = 11 - language_ko = 12 - language_nospeech = 13 + language_auto, language_zh, language_en = 0, 3, 4 + language_yue, language_ja, language_ko, language_nospeech = 7, 11, 12, 13 @@ - with_itn = 14 - without_itn = 15 + with_itn, without_itn = 14, 15 # keep if you plan to use later; otherwise remove @@ - # language = np.array([language], dtype=np.int32) - # text_norm = np.array([text_norm], dtype=np.int32) + # language = np.array([language], dtype=np.int32) + # text_norm = np.array([text_norm], dtype=np.int32)Or simply remove the unused variables entirely. As per static analysis hints.
Also applies to: 132-137
81-84: Fbank scaling: avoid Python list conversion in hot path
accept_waveformcan take a NumPy array; converting to list adds overhead. Prefernp.asarray(samples * 32768, dtype=np.float32).- online_fbank.accept_waveform(sample_rate, (samples * 32768).tolist()) + online_fbank.accept_waveform(sample_rate, np.asarray(samples * 32768, dtype=np.float32))
14-26: Device/config introspection is useful; add a short check that model.om existsPrints are helpful; add an existence check to fail fast with a clear message if users forget to run atc.
def __init__(self): - self.model = InferSession(device_id=0, model_path="./model.om", debug=False) + import os + assert os.path.exists("./model.om"), "model.om not found; run atc conversion first" + self.model = InferSession(device_id=0, model_path="./model.om", debug=False)
101-110: Optional: add CLI args for wav path and device idHardcoded zh.wav and device 0 are fine for the README, but a simple argparse improves reusability.
I can send a small argparse patch if you want.
158-163: Token mapping assumes contiguous IDs; handle explicit ids from file
load_tokensignores the numeric id in tokens.txt and uses line index. Safer to parse both fields and use the file’s id column.- tokens = load_tokens("./tokens.txt") - text = "".join([tokens[i] for i in ids]) + tokens = {} + with open("./tokens.txt", encoding="utf-8") as f: + for line in f: + piece, tok_id = line.rstrip("\n").rsplit(" ", 1) + tokens[int(tok_id)] = piece + text = "".join([tokens.get(i, "") for i in ids])
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
.github/workflows/export-sense-voice-to-ascend-npu.yaml(1 hunks)scripts/sense-voice/ascend-npu/export_onnx.py(1 hunks)scripts/sense-voice/ascend-npu/test_om.py(1 hunks)scripts/sense-voice/ascend-npu/torch_model.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
scripts/sense-voice/ascend-npu/export_onnx.py (1)
scripts/sense-voice/ascend-npu/torch_model.py (8)
SenseVoiceSmall(541-602)forward(41-46)forward(68-70)forward(201-220)forward(250-329)forward(336-344)forward(452-478)forward(592-602)
scripts/sense-voice/ascend-npu/test_om.py (3)
sherpa-onnx/csrc/ten-vad-model.cc (1)
frame_opts(237-255)scripts/sense-voice/ascend-npu/export_onnx.py (1)
main(52-96)scripts/sense-voice/ascend-npu/torch_model.py (1)
argmax(527-538)
🪛 Flake8 (7.3.0)
scripts/sense-voice/ascend-npu/test_om.py
[error] 4-4: 'argparse' imported but unused
(F401)
[error] 10-10: 'torch' imported but unused
(F401)
[error] 118-118: local variable 'language_zh' is assigned to but never used
(F841)
[error] 119-119: local variable 'language_en' is assigned to but never used
(F841)
[error] 120-120: local variable 'language_yue' is assigned to but never used
(F841)
[error] 121-121: local variable 'language_ya' is assigned to but never used
(F841)
[error] 122-122: local variable 'language_ko' is assigned to but never used
(F841)
[error] 123-123: local variable 'language_nospeech' is assigned to but never used
(F841)
[error] 128-128: local variable 'without_itn' is assigned to but never used
(F841)
[error] 142-142: block comment should start with '# '
(E265)
🪛 Ruff (0.14.1)
scripts/sense-voice/ascend-npu/test_om.py
118-118: Local variable language_zh is assigned to but never used
Remove assignment to unused variable language_zh
(F841)
119-119: Local variable language_en is assigned to but never used
Remove assignment to unused variable language_en
(F841)
120-120: Local variable language_yue is assigned to but never used
Remove assignment to unused variable language_yue
(F841)
121-121: Local variable language_ya is assigned to but never used
Remove assignment to unused variable language_ya
(F841)
122-122: Local variable language_ko is assigned to but never used
Remove assignment to unused variable language_ko
(F841)
123-123: Local variable language_nospeech is assigned to but never used
Remove assignment to unused variable language_nospeech
(F841)
128-128: Local variable without_itn is assigned to but never used
Remove assignment to unused variable without_itn
(F841)
🔇 Additional comments (3)
.github/workflows/export-sense-voice-to-ascend-npu.yaml (2)
198-207: Release token scope and tag reuseReleasing to k2-fsa with tag
asr-modelsoverwrites artifacts on every run. That’s intended? If not, append date/model name to the tag or userelease_namewith run-id. Ensure${{ secrets.UPLOAD_GH_SHERPA_ONNX_TOKEN }}is repo‑scoped withcontents:write.Also applies to: 209-216
33-37: Fix PyTorch wheel source URL to ensure install succeeds on CIThe PyTorch wheel URL
-f https://download.pytorch.org/whl/torchneeds to use the stable page. Change line 67 to:torch==2.0.0+cpu -f https://download.pytorch.org/whl/torch_stable.html \Also add a comment explaining the
numpy<2pin to document the version constraint reasoning.(Verification confirmed the
torch_stable.htmlURL resolves successfully.)Likely an incorrect or invalid review comment.
scripts/sense-voice/ascend-npu/torch_model.py (1)
1-1: Symlink is valid and working on Linux, but Windows compatibility concern is legitimateVerification confirms the symlink is correct:
- Target
../rknn/torch_model.pyproperly resolves to./scripts/sense-voice/rknn/torch_model.py- Target file exists and git tracking is valid (mode 120000)
- Works correctly on Linux CI
The concern about Windows users with
core.symlinks=falseis valid—they would see a 1-line text file instead of the actual module, causing import failures. The recommendation to either provide a real module during packaging or document the symlink requirement is sound for cross-platform support.
| return self.model.infer([x, prompt], mode="static", custom_sizes=10000000)[0][0] | ||
| logits_part1, logits_part2 = self.model.infer( | ||
| [x, prompt], mode="dymshape", custom_sizes=10000000 | ||
| ) | ||
| logits = np.concatenate([logits_part1, logits_part2], axis=0) | ||
| return logits | ||
|
|
There was a problem hiding this comment.
Unreachable code and stale two‑output merge path
Early returns make the merge path dead. Remove the static‑mode branch and the (part1, part2) concat logic to avoid confusion.
Apply this diff:
- return self.model.infer([x, prompt], mode="static", custom_sizes=10000000)[0][0]
- logits_part1, logits_part2 = self.model.infer(
- [x, prompt], mode="dymshape", custom_sizes=10000000
- )
- logits = np.concatenate([logits_part1, logits_part2], axis=0)
- return logits
+ # Single-output logits path onlyCommittable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In scripts/sense-voice/ascend-npu/test_om.py around lines 39 to 45, the early
return using mode="static" makes the subsequent dynamic-shape two-output path
dead and leaves stale concat logic; remove the static-mode return and the
logits_part1/logits_part2 capture and concatenation, and instead call
self.model.infer([x, prompt], mode="dymshape", custom_sizes=10000000) once and
return its single output (or the appropriate element of the returned structure)
so only the dynamic-shape path remains.
| features = np.stack( | ||
| [online_fbank.get_frame(i) for i in range(online_fbank.num_frames_ready)] | ||
| ) | ||
| assert features.data.contiguous is True | ||
| assert features.dtype == np.float32, features.dtype | ||
|
|
||
| T = (features.shape[0] - window_size) // window_shift + 1 | ||
| features = np.lib.stride_tricks.as_strided( | ||
| features, | ||
| shape=(T, features.shape[1] * window_size), | ||
| strides=((window_shift * features.shape[1]) * 4, 4), | ||
| ) |
There was a problem hiding this comment.
Incorrect NumPy contiguity check; potential negative T for short audio
features.data.contiguousis not a NumPy API and will raise. Usefeatures.flags.c_contiguous.- Guard
Tfrom going negative when audio is shorter thanwindow_size.
Apply this diff:
- assert features.data.contiguous is True
- assert features.dtype == np.float32, features.dtype
+ assert features.flags.c_contiguous, "features must be C-contiguous"
+ assert features.dtype == np.float32, features.dtype
@@
- T = (features.shape[0] - window_size) // window_shift + 1
+ num_frames = features.shape[0]
+ if num_frames < window_size:
+ return np.zeros((0, features.shape[1] * window_size), dtype=np.float32)
+ T = (num_frames - window_size) // window_shift + 1📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| features = np.stack( | |
| [online_fbank.get_frame(i) for i in range(online_fbank.num_frames_ready)] | |
| ) | |
| assert features.data.contiguous is True | |
| assert features.dtype == np.float32, features.dtype | |
| T = (features.shape[0] - window_size) // window_shift + 1 | |
| features = np.lib.stride_tricks.as_strided( | |
| features, | |
| shape=(T, features.shape[1] * window_size), | |
| strides=((window_shift * features.shape[1]) * 4, 4), | |
| ) | |
| features = np.stack( | |
| [online_fbank.get_frame(i) for i in range(online_fbank.num_frames_ready)] | |
| ) | |
| assert features.flags.c_contiguous, "features must be C-contiguous" | |
| assert features.dtype == np.float32, features.dtype | |
| num_frames = features.shape[0] | |
| if num_frames < window_size: | |
| return np.zeros((0, features.shape[1] * window_size), dtype=np.float32) | |
| T = (num_frames - window_size) // window_shift + 1 | |
| features = np.lib.stride_tricks.as_strided( | |
| features, | |
| shape=(T, features.shape[1] * window_size), | |
| strides=((window_shift * features.shape[1]) * 4, 4), | |
| ) |
1 participant
You can find the exported models at
https://github.com/k2-fsa/sherpa-onnx/releases/tag/asr-models
To test them, just download, unzip and run the
test_om.pyextracted from them.Summary by CodeRabbit