| |
| """Export a trained ModernBERT classifier for Jetson qualification.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
|
|
| class LogitsOnly(torch.nn.Module): |
| def __init__(self, model): |
| super().__init__() |
| self.model = model |
|
|
| def forward(self, input_ids, attention_mask): |
| return self.model(input_ids=input_ids, attention_mask=attention_mask).logits |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--opset", type=int, default=18) |
| args = parser.parse_args() |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| model = AutoModelForSequenceClassification.from_pretrained(args.model, attn_implementation="eager") |
| model.eval() |
| tokenizer = AutoTokenizer.from_pretrained(args.model) |
| sample = tokenizer("user: Send me a voice note about the clear evening.", return_tensors="pt") |
| wrapper = LogitsOnly(model) |
| torch.onnx.export( |
| wrapper, |
| (sample["input_ids"], sample["attention_mask"]), |
| str(args.output), |
| input_names=["input_ids", "attention_mask"], |
| output_names=["logits"], |
| dynamic_axes={ |
| "input_ids": {0: "batch", 1: "sequence"}, |
| "attention_mask": {0: "batch", 1: "sequence"}, |
| "logits": {0: "batch"}, |
| }, |
| opset_version=args.opset, |
| dynamo=False, |
| ) |
| import onnx |
| import onnxruntime as ort |
| onnx.checker.check_model(str(args.output)) |
| session = ort.InferenceSession(str(args.output), providers=["CPUExecutionProvider"]) |
| ort_logits = session.run(None, {name: value.numpy() for name, value in sample.items()})[0] |
| with torch.no_grad(): |
| torch_logits = wrapper(sample["input_ids"], sample["attention_mask"]).numpy() |
| worst = float(abs(ort_logits - torch_logits).max()) |
| result = {"model": str(args.model), "onnx": str(args.output), "max_abs_logit_error": worst, |
| "torch_pred": int(torch_logits.argmax()), "ort_pred": int(ort_logits.argmax())} |
| print(json.dumps(result), flush=True) |
| if result["torch_pred"] != result["ort_pred"] or worst > 0.01: |
| raise RuntimeError("ONNX export differs from PyTorch on smoke input") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
