Add onnxscript dependency and set dynamo=False for classic ONNX export

pyproject.toml

@@ -12,4 +12,5 @@ dependencies = [
     "numpy>=1.24.0",
     "onnx>=1.14.0",
     "onnxruntime>=1.15.0",
+    "onnxscript>=0.1.0",
 ]

requirements.txt

@@ -86,9 +86,15 @@ markupsafe==3.0.3
 mdurl==0.1.2 ; python_full_version >= '3.10'
     # via markdown-it-py
 ml-dtypes==0.5.4 ; python_full_version < '3.10'
-    # via onnx
+    # via
+    #   onnx
+    #   onnx-ir
+    #   onnxscript
 ml-dtypes==0.6.0 ; python_full_version >= '3.10'
-    # via onnx
+    # via
+    #   onnx
+    #   onnx-ir
+    #   onnxscript
 mpmath==1.3.0
     # via sympy
 narwhals==2.25.0 ; python_full_version >= '3.11'
@@ -104,7 +110,9 @@ numpy==2.0.2 ; python_full_version < '3.10'
     #   hatespeech
     #   ml-dtypes
     #   onnx
+    #   onnx-ir
     #   onnxruntime
+    #   onnxscript
     #   pandas
     #   scikit-learn
     #   scipy
@@ -114,7 +122,9 @@ numpy==2.2.6 ; python_full_version == '3.10.*'
     #   hatespeech
     #   ml-dtypes
     #   onnx
+    #   onnx-ir
     #   onnxruntime
+    #   onnxscript
     #   pandas
     #   scikit-learn
     #   scipy
@@ -124,7 +134,9 @@ numpy==2.4.6 ; python_full_version == '3.11.*'
     #   hatespeech
     #   ml-dtypes
     #   onnx
+    #   onnx-ir
     #   onnxruntime
+    #   onnxscript
     #   pandas
     #   scikit-learn
     #   scipy
@@ -134,7 +146,9 @@ numpy==2.5.2 ; python_full_version >= '3.12'
     #   hatespeech
     #   ml-dtypes
     #   onnx
+    #   onnx-ir
     #   onnxruntime
+    #   onnxscript
     #   pandas
     #   scikit-learn
     #   scipy
@@ -218,19 +232,30 @@ nvidia-nvtx==13.0.85 ; (python_full_version >= '3.10' and platform_machine == 'a
 nvidia-nvtx-cu12==12.8.90 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
     # via torch
 onnx==1.19.1 ; python_full_version < '3.10'
-    # via hatespeech
+    # via
+    #   hatespeech
+    #   onnx-ir
+    #   onnxscript
 onnx==1.22.0 ; python_full_version >= '3.10'
-    # via hatespeech
+    # via
+    #   hatespeech
+    #   onnx-ir
+    #   onnxscript
+onnx-ir==1.0.0
+    # via onnxscript
 onnxruntime==1.20.1 ; python_full_version < '3.10'
     # via hatespeech
 onnxruntime==1.24.3 ; python_full_version == '3.10.*'
     # via hatespeech
 onnxruntime==1.29.0 ; python_full_version >= '3.11'
     # via hatespeech
+onnxscript==0.7.1
+    # via hatespeech
 packaging==26.3
     # via
     #   huggingface-hub
     #   onnxruntime
+    #   onnxscript
     #   transformers
 pandas==2.3.3 ; python_full_version < '3.11'
     # via hatespeech
@@ -294,6 +319,7 @@ six==1.17.0
     # via python-dateutil
 sympy==1.14.0
     # via
+    #   onnx-ir
     #   onnxruntime
     #   torch
 threadpoolctl==3.6.0
@@ -326,6 +352,8 @@ typing-extensions==4.16.0
     #   exceptiongroup
     #   huggingface-hub
     #   onnx
+    #   onnx-ir
+    #   onnxscript
     #   torch
 tzdata==2026.3 ; python_full_version < '3.11' or sys_platform == 'emscripten' or sys_platform == 'win32'
     # via pandas

train.py

@@ -265,19 +265,28 @@ def export_to_onnx(
         dummy_encodings["attention_mask"],
     )
 
-    torch.onnx.export(
-        model,
-        dummy_inputs,
-        raw_onnx_path,
-        input_names=["input_ids", "attention_mask"],
-        output_names=["logits"],
-        dynamic_axes={
+    export_kwargs = {
+        "input_names": ["input_ids", "attention_mask"],
+        "output_names": ["logits"],
+        "dynamic_axes": {
             "input_ids": {0: "batch_size", 1: "sequence_length"},
             "attention_mask": {0: "batch_size", 1: "sequence_length"},
             "logits": {0: "batch_size"},
         },
-        opset_version=14,
-        do_constant_folding=True,
+        "opset_version": 14,
+        "do_constant_folding": True,
+    }
+
+    import inspect
+    sig = inspect.signature(torch.onnx.export)
+    if "dynamo" in sig.parameters:
+        export_kwargs["dynamo"] = False
+
+    torch.onnx.export(
+        model,
+        dummy_inputs,
+        raw_onnx_path,
+        **export_kwargs,
     )
 
     # Verify raw ONNX model validity

uv.lock

@@ -462,6 +462,7 @@ dependencies = [
     { name = "onnxruntime", version = "1.20.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
     { name = "onnxruntime", version = "1.24.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" },
     { name = "onnxruntime", version = "1.29.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
+    { name = "onnxscript" },
     { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "pandas", version = "3.0.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "scikit-learn", version = "1.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
@@ -478,6 +479,7 @@ requires-dist = [
     { name = "numpy", specifier = ">=1.24.0" },
     { name = "onnx", specifier = ">=1.14.0" },
     { name = "onnxruntime", specifier = ">=1.15.0" },
+    { name = "onnxscript", specifier = ">=0.1.0" },
     { name = "pandas", specifier = ">=2.0.0" },
     { name = "scikit-learn", specifier = ">=1.2.0" },
     { name = "torch", specifier = ">=2.0.0" },
@@ -1648,6 +1650,27 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/00/50/257a880384a1dd502d543b0067945074d63cd17d0840e958355bc8197da8/onnx-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:8e268cdc0547e3949799ffd4a44451dc2b9080b57d0824a2db680b6ec65506f0", size = 17231391, upload-time = "2026-06-15T12:50:03.047Z" },
 ]
 
+[[package]]
+name = "onnx-ir"
+version = "1.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
+    { name = "ml-dtypes", version = "0.6.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
+    { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" },
+    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
+    { name = "numpy", version = "2.5.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+    { name = "onnx", version = "1.19.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
+    { name = "onnx", version = "1.22.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
+    { name = "sympy" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d6/c2/61194cec0dbc5622273c0ebd592d37cc1dca0d7f1a744f02edd45ac905a3/onnx_ir-1.0.0.tar.gz", hash = "sha256:9e261f25fde8da9612ae5cb43b3b374d5ff469c04af0363cad588b2bb000b812", size = 163121, upload-time = "2026-08-11T14:49:46.895Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/91/cd/6d1637172eb59c7b18ac90ed089d1f599a11fe0e63b4db2d017f3bb38a32/onnx_ir-1.0.0-py3-none-any.whl", hash = "sha256:e578f0d608d3062866b48223616eb2d10a6d6d01f8b8faac596129034f483cc7", size = 185849, upload-time = "2026-08-11T14:49:45.524Z" },
+]
+
 [[package]]
 name = "onnxruntime"
 version = "1.20.1"
@@ -1780,6 +1803,28 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/60/21/d0c04b561b46e9bff89b5f500fb7415b8ca0669f7902204f76ab06bb0c7e/onnxruntime-1.29.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:1ea91cef3b971506e51ae9c37c16d027774ec64994a524ec1bdfb027d68a9832", size = 23138547, upload-time = "2026-08-17T22:54:37.491Z" },
 ]
 
+[[package]]
+name = "onnxscript"
+version = "0.7.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
+    { name = "ml-dtypes", version = "0.6.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
+    { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.10.*'" },
+    { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
+    { name = "numpy", version = "2.5.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
+    { name = "onnx", version = "1.19.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.10'" },
+    { name = "onnx", version = "1.22.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" },
+    { name = "onnx-ir" },
+    { name = "packaging" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c4/3a/4d79bce3f460e0df7fed54a92ce80827f25da66511da368bb00783ad8d20/onnxscript-0.7.1.tar.gz", hash = "sha256:309fb86484b11fa4ded90dba580e0d63f1a0827588e521cecaf2eeddb46d6e86", size = 618160, upload-time = "2026-06-29T23:33:21.526Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/dd/bd/a0c8e737b6afda10e42a597787d53d5b66e00268df6f59184701eeae37d9/onnxscript-0.7.1-py3-none-any.whl", hash = "sha256:544763b7fdef49940cdd9412ff5135cbae96d59ac6bc1921457f21280f40f4b7", size = 721970, upload-time = "2026-06-29T23:33:23.298Z" },
+]
+
 [[package]]
 name = "packaging"
 version = "26.3"