RAM_Altman_GPT / src /convert_onnx.py
brunuu's picture
1.0
7de15fc verified
Raw History Blame Contribute Delete
686 Bytes
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)")