Add an ONNX export script#17
Open
AlJohri wants to merge 1 commit into
Open
Conversation
Add src/deep_impact/scripts/export_onnx.py to export a trained DeepImpact checkpoint to a single-file ONNX model, with a built-in PyTorch-vs-ONNX Runtime parity check over a padded batch of documents. optimum-cli can't be used because DeepImpact isn't registered with transformers' AutoModel machinery -- it falls back to a plain BertModel and silently drops the impact-score head -- so this exports a thin wrapper via torch.onnx.export directly. Adds onnx and onnxruntime to requirements.txt; both are used only by this script. Verified on the repo's pinned torch 2.0.1 (single-file export, parity max|diff| 5.7e-06).
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.
What
Adds
src/deep_impact/scripts/export_onnx.py, a small utility that exports a trained DeepImpact checkpoint to a single-file ONNX model and verifies it against PyTorch:python -m src.deep_impact.scripts.export_onnx \ --model_checkpoint_path soyuj/deeper-impact \ --output_path onnx/model.onnxAlso adds
onnxandonnxruntimetorequirements.txt(used only by this script).Why
There's currently no ONNX export of DeepImpact, which makes it awkward to run inference outside Python (ONNX Runtime from Rust/C++/JS, etc.).
optimum-cli export onnxcan't produce a correct export here: DeepImpact setsarchitectures: ["DeepImpact"]but isn't registered with transformers'AutoModel*, so optimum falls back to a plainBertModel(viamodel_type: "bert") and silently drops theimpact_score_encoderhead — emitting hidden states instead of impact scores, with no error. This script sidesteps that by tracing a thin wrapper around the real module withtorch.onnx.export.What gets exported
input_ids,attention_mask,token_type_ids[batch, seq]impact_scores[batch, seq]Batch and sequence axes are dynamic (one model serves any batch size / length up to 512). A term's impact is the score at its first subword token — the same indexing
DeepImpact.compute_term_impactsalready does.Parity
After export, the script runs a built-in check: it feeds a padded batch of two documents through both PyTorch and ONNX Runtime and asserts the per-term scores agree.
Verified on the repo's pinned
torch==2.0.1(Python 3.11). The output is a single self-containedmodel.onnx(~436 MB, no external-data sidecar).Companion PR
A matching PR on the Hugging Face repo uploads the exported
model.onnx(toonnx/model.onnx) and adds an ONNX section to the model card: https://huggingface.co/soyuj/deeper-impact/discussions/1