File size: 686 Bytes
7de15fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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)")