# ONNX export

This repo uses **split environments**:

- main env: latest training stack (`transformers` 5.x)
- ONNX env: exporter-compatible stack from `training/requirements-export.txt`

Reason: current `optimum-onnx` release still requires `transformers <4.58`, while training uses latest `transformers 5.x`.

## Why split env is acceptable

For stable architectures like DistilBERT token classification, exporting from a separately loaded saved checkpoint is normal practice. Safety comes from validating exported ONNX outputs against PyTorch outputs.

## Training

```bash
uv sync
python -m training.train_ner
```

Best checkpoint is saved under:

```text
training/output/resume-ner/distilbert/best
```

## ONNX export environment

```bash
uv venv .venv-onnx-export
uv pip install --python .venv-onnx-export/bin/python -r training/requirements-export.txt
source .venv-onnx-export/bin/activate
```

## Export command

```bash
optimum-cli export onnx \
  --model training/output/resume-ner/distilbert/best \
  --task token-classification \
  onnx/
```

This writes:

- `onnx/model.onnx`
- tokenizer/config files in `onnx/`

## Quantization

```bash
python - <<'PY'
from optimum.onnxruntime import ORTModelForTokenClassification, ORTQuantizer
from optimum.onnxruntime.configuration import AutoQuantizationConfig

model = ORTModelForTokenClassification.from_pretrained("onnx")
quantizer = ORTQuantizer.from_pretrained(model)
qconfig = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
quantizer.quantize(save_dir="onnx", quantization_config=qconfig)
PY
```

This writes:

- `onnx/model_quantized.onnx`

## Validation

Checked-in helper scripts:

- `python -m training.export_onnx`
- `python -m training.validate_onnx`
- `python -m training.quantize_onnx`
- `python -m training.benchmark_structured --model-dir .` — internal structured benchmark


Always compare PyTorch vs ONNX on same tokenized input.

Recommended checks:

1. output shape match
2. `np.allclose(..., rtol=1e-3, atol=1e-5)` for non-quantized ONNX
3. `argmax_equal=True`
4. quantized ONNX at least preserves argmax predictions on smoke test inputs

Minimal validation example:

```python
import numpy as np
import torch
from transformers import AutoModelForTokenClassification, AutoTokenizer
from optimum.onnxruntime import ORTModelForTokenClassification

model_dir = "training/output/resume-ner/distilbert/best"
text = "Rajesh Kumar worked at Infosys in Bangalore from April 2020 to Present using Python and AWS."

tokenizer = AutoTokenizer.from_pretrained(model_dir)
pt_model = AutoModelForTokenClassification.from_pretrained(model_dir)
pt_model.eval()
ort_model = ORTModelForTokenClassification.from_pretrained("onnx")

inputs_pt = {k: v for k, v in tokenizer(text, return_tensors="pt", truncation=True).items() if k in {"input_ids", "attention_mask"}}
inputs_np = {k: v for k, v in tokenizer(text, return_tensors="np", truncation=True).items() if k in {"input_ids", "attention_mask"}}

with torch.no_grad():
    pt_logits = pt_model(**inputs_pt).logits.cpu().numpy()
ort_logits = ort_model(**inputs_np).logits

print(np.allclose(pt_logits, ort_logits, rtol=1e-3, atol=1e-5))
print(np.array_equal(pt_logits.argmax(-1), ort_logits.argmax(-1)))
```

## Hugging Face artifact policy

Upload both:

### Root
- `model.safetensors`
- `config.json`
- tokenizer files
- `companies.json`
- `label_config.json`
- `resume_config.json`

### `onnx/`
- `model.onnx`
- `model_quantized.onnx`
- ONNX tokenizer/config files

This keeps repo usable for:
- Transformers users
- ONNX users
- future re-export or debugging