Add export-onnx.py for fluent_speech_commands #1992
Conversation
WalkthroughA new script is introduced to export a speech recognition transducer model's encoder, decoder, and joiner components to ONNX format, with optional int8 quantization. The script defines an ONNX-compatible encoder wrapper, provides export functions for each component, adds metadata, and supports command-line execution with logging. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Script
participant Model
participant ONNXExporter
participant Quantizer
User->>Script: Run export-onnx.py with arguments
Script->>Model: Load and prepare model (encoder, decoder, joiner)
Script->>ONNXExporter: Export encoder to ONNX
Script->>ONNXExporter: Export decoder to ONNX
Script->>ONNXExporter: Export joiner to ONNX
Script->>Quantizer: Quantize ONNX models (encoder, decoder, joiner)
Script->>Script: Add metadata to ONNX files
Script->>User: Log progress and completion
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
egs/fluent_speech_commands/SLU/transducer/export-onnx.py (2)
11-11: Consider using absolute imports for better portability.The import from
decodeappears to be a relative import. Consider using absolute imports or making the relative nature explicit with a dot prefix to avoid import issues when the script is run from different directories.-from decode import get_params, get_parser, get_transducer_model +from .decode import get_params, get_parser, get_transducer_model
31-32: Consider using a more idiomatic approach to clear metadata.The while loop for clearing metadata could be replaced with a more direct approach if the API supports it.
- while len(model.metadata_props): - model.metadata_props.pop() + del model.metadata_props[:]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
egs/fluent_speech_commands/SLU/transducer/export-onnx.py(1 hunks)
🔇 Additional comments (2)
egs/fluent_speech_commands/SLU/transducer/export-onnx.py (2)
42-75: LGTM!The
OnnxEncoderwrapper class is well-implemented with clear documentation and proper type annotations.
77-130: LGTM!The encoder export function is well-structured with comprehensive documentation, proper dynamic axes configuration, and appropriate metadata handling.
| hidden_dim=params.hidden_dim, | ||
| vocab_size=params.vocab_size, | ||
| ) | ||
| logging.info(f"Exported decoder to {joiner_filename}") |
There was a problem hiding this comment.
Fix incorrect log message.
The log message incorrectly states "decoder" instead of "joiner".
- logging.info(f"Exported decoder to {joiner_filename}")
+ logging.info(f"Exported joiner to {joiner_filename}")📝 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.
| logging.info(f"Exported decoder to {joiner_filename}") | |
| logging.info(f"Exported joiner to {joiner_filename}") |
🤖 Prompt for AI Agents
In egs/fluent_speech_commands/SLU/transducer/export-onnx.py at line 304, the log
message incorrectly refers to "decoder" instead of "joiner". Update the
logging.info message to correctly state "Exported joiner to {joiner_filename}"
to accurately reflect the exported component.
| "vocab_size": vocab_size, | ||
| "hidden_dim": hidden_dim, | ||
| } | ||
| print("joiner meta_data", meta_data) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use logging.info() for consistency.
Replace the print statement with logging.info() to maintain consistency with the rest of the codebase.
- print("joiner meta_data", meta_data)
+ logging.info(f"joiner meta_data: {meta_data}")📝 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.
| print("joiner meta_data", meta_data) | |
| logging.info(f"joiner meta_data: {meta_data}") |
🤖 Prompt for AI Agents
In egs/fluent_speech_commands/SLU/transducer/export-onnx.py at line 230, replace
the print statement with logging.info() to maintain consistency with the rest of
the codebase. Ensure that logging is properly imported and configured if not
already done.
| "num_layers": h.shape[0], | ||
| "hidden_dim": h.shape[2], | ||
| } | ||
| print("decoder meta_data", meta_data) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use logging.info() for consistency.
Replace the print statement with logging.info() to maintain consistency with the rest of the codebase.
- print("decoder meta_data", meta_data)
+ logging.info(f"decoder meta_data: {meta_data}")📝 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.
| print("decoder meta_data", meta_data) | |
| - print("decoder meta_data", meta_data) | |
| + logging.info(f"decoder meta_data: {meta_data}") |
🤖 Prompt for AI Agents
In egs/fluent_speech_commands/SLU/transducer/export-onnx.py at line 179, replace
the print statement used for outputting "decoder meta_data" with logging.info()
to maintain consistency with the project's logging practices. Ensure logging is
properly imported and configured if not already done.
| start = params.epoch - params.avg + 1 | ||
| filenames = [] | ||
| for i in range(start, params.epoch + 1): | ||
| if start >= 0: | ||
| filenames.append(f"{params.exp_dir}/epoch-{i}.pt") | ||
| logging.info(f"averaging {filenames}") | ||
| model.load_state_dict(average_checkpoints(filenames)) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Move loop-invariant condition outside the loop.
The condition if start >= 0: doesn't depend on the loop variable and should be checked once outside the loop for better efficiency.
- start = params.epoch - params.avg + 1
- filenames = []
- for i in range(start, params.epoch + 1):
- if start >= 0:
- filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ start = params.epoch - params.avg + 1
+ filenames = []
+ if start >= 0:
+ for i in range(start, params.epoch + 1):
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")📝 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.
| start = params.epoch - params.avg + 1 | |
| filenames = [] | |
| for i in range(start, params.epoch + 1): | |
| if start >= 0: | |
| filenames.append(f"{params.exp_dir}/epoch-{i}.pt") | |
| logging.info(f"averaging {filenames}") | |
| model.load_state_dict(average_checkpoints(filenames)) | |
| start = params.epoch - params.avg + 1 | |
| filenames = [] | |
| if start >= 0: | |
| for i in range(start, params.epoch + 1): | |
| filenames.append(f"{params.exp_dir}/epoch-{i}.pt") | |
| logging.info(f"averaging {filenames}") | |
| model.load_state_dict(average_checkpoints(filenames)) |
🤖 Prompt for AI Agents
In egs/fluent_speech_commands/SLU/transducer/export-onnx.py around lines 250 to
256, the condition 'if start >= 0:' is inside the loop but does not depend on
the loop variable. Move this condition outside the loop to check it once before
iterating, and only run the loop to append filenames if 'start' is greater than
or equal to zero. This will improve efficiency by avoiding redundant checks in
each iteration.
1 participant
You can find the checkpoints at
https://hf-mirror.com/YangQiangCD/fluent_speech_commands
Summary by CodeRabbit