Update ONNX inference metadata types

Author: olokobayusuf

Changelog.md

@@ -1,4 +1,7 @@
 ## 0.0.89
++ Added `OnnxRuntimeInferenceSessionMetadata.external_data_path` field for specifying the path to external data files.
++ Added `OnnxRuntimeInferenceSessionMetadata.providers` field for specifying enabled ONNXRuntime providers.
++ Added `OnnxRuntimeInferenceMetadata.providers` field for specifying enabled ONNXRuntime providers.
 + Removed `muna source` CLI command.
 
 ## 0.0.88

muna/beta/metadata/__init__.py

@@ -9,10 +9,16 @@
 from .iree import IREEInferenceBackend, IREEInferenceMetadata
 from .litert import LiteRTInferenceMetadata
 from .llama import LlamaCppBackend, LlamaCppInferenceMetadata
-from .onnx import OnnxRuntimeInferenceMetadata, OnnxRuntimeOptimizationLevel
+from .onnx import (
+    OnnxRuntimeInferenceMetadata, OnnxRuntimeExecutionProvider,
+    OnnxRuntimeOptimizationLevel
+)
 from .onnxruntime import OnnxRuntimeInferenceSessionMetadata
 from .openvino import OpenVINOInferenceMetadata
 from .qnn import QnnInferenceBackend, QnnInferenceMetadata, QnnInferenceQuantization
-from .tensorrt import CudaArchitecture, TensorRTInferenceMetadata, TensorRTHardwareCompatibility, TensorRTPrecision
+from .tensorrt import (
+    CudaArchitecture, TensorRTInferenceMetadata,
+    TensorRTHardwareCompatibility, TensorRTPrecision
+)
 from .tensorrt_rtx import TensorRTRTXInferenceMetadata
 from .tflite import TFLiteInterpreterMetadata

muna/beta/metadata/onnx.py

@@ -8,6 +8,7 @@
 
 from ._torch import PyTorchInferenceMetadataBase
 
+OnnxRuntimeExecutionProvider = Literal["cpu", "coreml", "cuda", "openvino", "xnnpack"]
 OnnxRuntimeOptimizationLevel = Literal["none", "basic", "extended"]
 
 class OnnxRuntimeInferenceMetadata(PyTorchInferenceMetadataBase):
@@ -27,4 +28,9 @@ class OnnxRuntimeInferenceMetadata(PyTorchInferenceMetadataBase):
         default="none",
         description="ONNX model optimization level. Defaults to `none`.",
         exclude=True
+    )
+    providers: list[OnnxRuntimeExecutionProvider] | None = Field(
+        default=None,
+        description="ONNXRuntime execution providers to build with.",
+        exclude=True
     )

muna/beta/metadata/onnxruntime.py

@@ -7,6 +7,8 @@
 from pydantic import BaseModel, BeforeValidator, ConfigDict, Field
 from typing import Annotated, Literal
 
+from .onnx import OnnxRuntimeExecutionProvider
+
 def _validate_ort_inference_session(session: "onnxruntime.InferenceSession") -> "onnxruntime.InferenceSession": # type: ignore
     try:
         from onnxruntime import InferenceSession
@@ -32,4 +34,14 @@ class OnnxRuntimeInferenceSessionMetadata(BaseModel, **ConfigDict(arbitrary_type
     model_path: str | Path = Field(
         description="ONNX model path. The model must exist at this path in the compiler sandbox.",
         exclude=True
+    )
+    external_data_path: str | Path | None = Field(
+        default=None,
+        description="Path to ONNX external data file (e.g. .onnx.data).",
+        exclude=True
+    )
+    providers: list[OnnxRuntimeExecutionProvider] | None = Field(
+        default=None,
+        description="ONNXRuntime execution providers to build with.",
+        exclude=True
     )

muna/cli/__init__.py

@@ -40,14 +40,14 @@
     help="Compile a Python function for deployment.",
     rich_help_panel="Compilation"
 )(compile_function)
+
+# Functions
 app.command(
     name="predict",
     help="Invoke a compiled Python function.",
     context_settings={ "allow_extra_args": True, "ignore_unknown_options": True },
-    rich_help_panel="Compilation"
+    rich_help_panel="Functions"
 )(create_prediction)
-
-# Predictors
 app.command(
     name="retrieve",
     help="Retrieve a compiled function.",