Refactor ONNX export script for nemotron-3.5-asr-streaming-0.6b - #3734
Conversation
dosubot
Bot
added
the
size:L
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThis PR simplifies the Nemotron-3.5 ASR export script by replacing SentencePiece-based token extraction with vocabulary-based extraction, removing several ONNX graph validation assertions, adjusting the export opset version and feature dimension derivation, and adds CI steps to install sherpa-onnx and validate exported ONNX models across multiple chunk sizes. ChangesExport script and CI test updates
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
✨ 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 simplifies the ONNX export script for Nemotron ASR streaming by removing several validation checks, refactoring token saving to extract vocabulary directly from the model's joint vocabulary, dynamically retrieving the feature dimension from the model configuration, and downgrading the ONNX opset version from 17 to 13. Feedback highlights an incorrect return type annotation and potential UnboundLocalError in the save_tokens function, as well as a hardcoded feature dimension remaining in the metadata dictionary.
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.
| def save_tokens(asr_model, filename: str = "tokens.txt") -> int: | ||
| sp = _find_sentencepiece_processor(getattr(asr_model, "tokenizer", None)) | ||
| if sp is None: | ||
| raise RuntimeError("Could not find the SentencePiece tokenizer in the model") | ||
|
|
||
| vocab_size = sp.get_piece_size() | ||
| with open(filename, "w", encoding="utf-8") as f: | ||
| for i in range(vocab_size): | ||
| f.write(f"{sp.id_to_piece(i)} {i}\n") | ||
| f.write(f"<blk> {vocab_size}\n") | ||
|
|
||
| print(f"Saved {filename}") | ||
| return vocab_size | ||
|
|
||
|
|
||
| def assert_forward_for_export_signature(encoder): | ||
| if not hasattr(encoder, "forward_for_export"): | ||
| raise RuntimeError("Expected encoder.forward_for_export for ONNX export") | ||
|
|
||
| signature = inspect.signature(encoder.forward_for_export) | ||
| missing = [ | ||
| name for name in FORWARD_FOR_EXPORT_ARGS if name not in signature.parameters | ||
| ] | ||
| if missing: | ||
| raise RuntimeError( | ||
| "encoder.forward_for_export is missing expected argument(s): " | ||
| f"{missing}. Signature: {signature}" | ||
| ) | ||
| for i, s in enumerate(asr_model.joint.vocabulary): | ||
| f.write(f"{s} {i}\n") | ||
| f.write(f"<blk> {i+1}\n") |
There was a problem hiding this comment.
The save_tokens function has a return type annotation of -> int, but it does not return any value (implicitly returning None). Additionally, referencing the loop variable i outside the loop to write the <blk> token can raise an UnboundLocalError if the vocabulary is empty, and is generally considered a code smell in Python.
Consider updating the return type annotation to -> None and using len(vocab) instead of i+1 to determine the vocabulary size safely.
| def save_tokens(asr_model, filename: str = "tokens.txt") -> int: | |
| sp = _find_sentencepiece_processor(getattr(asr_model, "tokenizer", None)) | |
| if sp is None: | |
| raise RuntimeError("Could not find the SentencePiece tokenizer in the model") | |
| vocab_size = sp.get_piece_size() | |
| with open(filename, "w", encoding="utf-8") as f: | |
| for i in range(vocab_size): | |
| f.write(f"{sp.id_to_piece(i)} {i}\n") | |
| f.write(f"<blk> {vocab_size}\n") | |
| print(f"Saved {filename}") | |
| return vocab_size | |
| def assert_forward_for_export_signature(encoder): | |
| if not hasattr(encoder, "forward_for_export"): | |
| raise RuntimeError("Expected encoder.forward_for_export for ONNX export") | |
| signature = inspect.signature(encoder.forward_for_export) | |
| missing = [ | |
| name for name in FORWARD_FOR_EXPORT_ARGS if name not in signature.parameters | |
| ] | |
| if missing: | |
| raise RuntimeError( | |
| "encoder.forward_for_export is missing expected argument(s): " | |
| f"{missing}. Signature: {signature}" | |
| ) | |
| for i, s in enumerate(asr_model.joint.vocabulary): | |
| f.write(f"{s} {i}\n") | |
| f.write(f"<blk> {i+1}\n") | |
| def save_tokens(asr_model, filename: str = "tokens.txt") -> None: | |
| vocab = asr_model.joint.vocabulary | |
| with open(filename, "w", encoding="utf-8") as f: | |
| for i, s in enumerate(vocab): | |
| f.write(f"{s} {i}\n") | |
| f.write(f"<blk> {len(vocab)}\n") |
| feat_dim = asr_model.cfg.preprocessor.features | ||
|
|
||
| audio_signal = torch.zeros(1, feat_dim, window_size, dtype=dtype, device=device) |
There was a problem hiding this comment.
While feat_dim is now dynamically retrieved from asr_model.cfg.preprocessor.features to initialize the dummy audio_signal tensor, the metadata dictionary on line 359 still hardcodes "feat_dim": 128. To ensure consistency and avoid potential issues if the model configuration changes, the metadata should also use asr_model.cfg.preprocessor.features instead of a hardcoded value.
1 participant
Remove a lot of unncessary code.
Summary by CodeRabbit
Bug Fixes
Tests