import shutil
import sys
import tempfile
from pathlib import Path
from config import OUTPUT_DIR
from optimum.exporters.onnx import base, main_export

MODEL = Path(OUTPUT_DIR)
OUT = MODEL / "onnx" / "model.onnx"

if not (MODEL / "model.safetensors").exists():
    sys.exit("NO MODEL FOUND!")


# optimum only looks for a package named "onnxruntime", so it misses onnxruntime-directml
base.is_onnxruntime_available = lambda: True

with tempfile.TemporaryDirectory() as tmp:
    main_export(str(MODEL), output=tmp, task="text-generation-with-past")
    OUT.parent.mkdir(exist_ok=True)
    shutil.move(Path(tmp) / "model.onnx", OUT)

print(f"done: {OUT} ({OUT.stat().st_size // 1024} KB)")