#!/usr/bin/env python3
"""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()