Test Whisper on Ascend NPU using ACL Python API - #2986
Conversation
dosubot
Bot
added
the
size:L
|
Caution Review failedThe pull request is closed. 📝 WalkthroughWalkthroughThis PR adds Ascend-NPU support for Whisper model testing by introducing a new comprehensive test module with ONNX/InferSession-based inference, audio processing, and feature computation. A shared ONNX export delegation is added, and debug print statements are removed from existing RKNN tests. Changes
Sequence Diagram(s)sequenceDiagram
participant Audio as Audio File
participant FeatComp as Feature Computation
participant Encoder as Encoder InferSession
participant Decoder as Decoder InferSession
participant KVCache as KV Cache Manager
participant TokenMap as Token Mapper
Audio->>FeatComp: load_audio(filename)<br/>16kHz WAV → float32
FeatComp->>FeatComp: compute_features()<br/>kaldi_native_fbank → 80-dim
FeatComp->>FeatComp: pad/trim to 3000 frames<br/>transpose to (1,80,3000)
FeatComp->>Encoder: run_encoder(features)
rect rgb(200, 240, 255)
Note over Encoder: Initialize OmModel<br/>load encoder/decoder sessions
Encoder->>KVCache: get_self_cache()<br/>allocate initial KV tensors
end
Encoder-->>Decoder: return cross_kv
KVCache-->>Decoder: initialized self_kv
rect rgb(240, 220, 255)
Note over Decoder: Decoding Loop (≤100 steps or EOT)
loop Until EOT or max steps
Decoder->>Decoder: causal_mask_1d(offset, 3000)<br/>compute position mask
Decoder->>Decoder: run_decoder(tokens, self_kv,<br/>cross_kv, offset, mask)
Decoder->>KVCache: update self_kv from decoder output
Decoder->>Decoder: append new token to sequence
end
end
Decoder->>TokenMap: decode token IDs<br/>lookup → base64-decode
TokenMap-->>Audio: output text transcription
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Poem
✨ Finishing touches
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (3)
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 |
Summary of ChangesHello @csukuangfj, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly expands the project's hardware compatibility by integrating Whisper model testing capabilities for Ascend NPUs. It provides a dedicated Python script that orchestrates the entire inference pipeline, from audio input and feature computation to the final transcription using optimized encoder and decoder models. This enables efficient evaluation and deployment of Whisper models on Ascend platforms. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 adds a test script for Whisper on Ascend NPU using the ACL Python API. The changes include a new test script, a file that seems to be a misplaced symlink, and cleanup of print statements in an existing script.
My review focuses on the new script and the file structure. I've found a critical issue with a file that should likely be a symbolic link. I've also suggested improvements for code quality and maintainability in the new test script, including adding type hints, specifying file encoding, and refactoring duplicated code.
| @@ -0,0 +1 @@ | |||
| ../rknn/export_onnx.py No newline at end of file | |||
There was a problem hiding this comment.
This file appears to be intended as a symbolic link to ../rknn/export_onnx.py, but it's a regular file containing the path as its content. This will not work as a Python script. If you execute it with python, it will raise a SyntaxError. It should be a symbolic link. You can create it with ln -s ../rknn/export_onnx.py scripts/whisper/ascend-npu/export_onnx.py and commit the symlink.
| def load_tokens(filename): | ||
| tokens = dict() | ||
| with open(filename, "r") as f: | ||
| for line in f: | ||
| t, i = line.split() | ||
| tokens[int(i)] = t | ||
| return tokens |
There was a problem hiding this comment.
The load_tokens function can be improved by adding type hints for better readability and maintainability. Also, it's a good practice to explicitly specify the encoding when opening text files to avoid platform-dependent behavior. I'd suggest using utf-8.
| def load_tokens(filename): | |
| tokens = dict() | |
| with open(filename, "r") as f: | |
| for line in f: | |
| t, i = line.split() | |
| tokens[int(i)] = t | |
| return tokens | |
| def load_tokens(filename: str) -> "dict[int, str]": | |
| tokens = dict() | |
| with open(filename, "r", encoding="utf-8") as f: | |
| for line in f: | |
| t, i = line.split() | |
| tokens[int(i)] = t | |
| return tokens |
| offset = np.array([0], dtype=np.int32) | ||
| for t in model.sot_sequence: | ||
| token = np.array([[t]], dtype=np.int32) # sot | ||
| mask = causal_mask_1d(offset.item(), model.n_text_ctx) | ||
|
|
||
| out = model.run_decoder( | ||
| tokens=token, self_kv=self_kv, cross_kv=cross_kv, offset=offset, mask=mask | ||
| ) | ||
|
|
||
| for i in range(1, len(out)): | ||
| self_kv[i - 1][:, offset.item() : offset.item() + 1, :] = out[i] | ||
|
|
||
| offset += 1 | ||
|
|
||
| idx = out[0][0, 0].argmax() | ||
|
|
||
| eot = model.eot | ||
|
|
||
| ans = [] | ||
|
|
||
| while idx != eot and offset.item() < 100: | ||
| ans.append(idx) | ||
| token = np.array([[idx]], dtype=np.int32) | ||
|
|
||
| mask = causal_mask_1d(offset.item(), model.n_text_ctx) | ||
|
|
||
| out = model.run_decoder( | ||
| tokens=token, self_kv=self_kv, cross_kv=cross_kv, offset=offset, mask=mask | ||
| ) | ||
|
|
||
| for i in range(1, len(out)): | ||
| self_kv[i - 1][:, offset.item() : offset.item() + 1, :] = out[i] | ||
|
|
||
| offset += 1 | ||
| idx = out[0][0, 0].argmax() |
There was a problem hiding this comment.
The decoding logic is duplicated in the loop for the SOT sequence and the main decoding while loop. This can be refactored into a helper function to improve code readability and maintainability.
offset = np.array([0], dtype=np.int32)
def _decode_one_step(model, token_id, self_kv, cross_kv, current_offset):
token = np.array([[token_id]], dtype=np.int32)
mask = causal_mask_1d(current_offset.item(), model.n_text_ctx)
out = model.run_decoder(
tokens=token, self_kv=self_kv, cross_kv=cross_kv, offset=current_offset, mask=mask
)
for i in range(1, len(out)):
self_kv[i - 1][:, current_offset.item() : current_offset.item() + 1, :] = out[i]
current_offset += 1
return out, current_offset
for t in model.sot_sequence:
out, offset = _decode_one_step(model, t, self_kv, cross_kv, offset)
idx = out[0][0, 0].argmax()
eot = model.eot
ans = []
while idx != eot and offset.item() < 100:
ans.append(idx)
out, offset = _decode_one_step(model, idx, self_kv, cross_kv, offset)
idx = out[0][0, 0].argmax()
1 participant
Summary by CodeRabbit
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.