Add optimum exporter

Author: olokobayusuf

Changelog.md

@@ -1,5 +1,5 @@
 ## 0.0.90
-*INCOMPLETE*
++ Added `optimum` exporter in `OnnxRuntimeInferenceMetadata` to export LLMs using 🤗 Optimum.
 
 ## 0.0.89
 + Added `OnnxRuntimeInferenceSessionMetadata.external_data_path` field for specifying the path to external data files.

muna/beta/metadata/_torch.py

@@ -6,7 +6,7 @@
 from pydantic import BaseModel, BeforeValidator, ConfigDict, Field
 from typing import Annotated, Literal
 
-TorchExporter = Literal["none", "dynamo", "torchscript"]
+TorchExporter = Literal["none", "dynamo", "torchscript", "optimum"]
 
 def _validate_torch_module(module: "torch.nn.Module") -> "torch.nn.Module": # type: ignore
     try:
@@ -17,7 +17,9 @@ def _validate_torch_module(module: "torch.nn.Module") -> "torch.nn.Module": # ty
     except ImportError:
         raise ImportError("PyTorch is required to create this metadata but it is not installed.")
 
-def _validate_torch_tensor_args(args: list) -> list:
+def _validate_torch_tensor_args(args: list) -> list | None:
+    if args is None:
+        return args
     try:
         from torch import Tensor
         for idx, arg in enumerate(args):
@@ -27,13 +29,30 @@ def _validate_torch_tensor_args(args: list) -> list:
     except ImportError:
         raise ImportError("PyTorch is required to create this metadata but it is not installed.")
 
+def _validate_optimum_exporter_config(config: object) -> object:
+    if config is None:
+        return config
+    try:
+        from optimum.exporters.base import ExporterConfig
+        if not isinstance(config, ExporterConfig):
+            raise ValueError(f"Expected `optimum.exporters.base.ExporterConfig` instance for `optimum_config` but got `{type(config).__qualname__}`")
+        return config
+    except ImportError:
+        pass
+
 class PyTorchInferenceMetadataBase(BaseModel, **ConfigDict(arbitrary_types_allowed=True, frozen=True)):
     model: Annotated[object, BeforeValidator(_validate_torch_module)] = Field(
         description="PyTorch module to apply metadata to.",
         exclude=True
     )
-    model_args: Annotated[list[object], BeforeValidator(_validate_torch_tensor_args)] = Field(
-        description="Positional inputs to the model.",
+    exporter: TorchExporter | None = Field(
+        default=None,
+        description="PyTorch exporter to use.",
+        exclude=True
+    )
+    model_args: Annotated[list[object] | None, BeforeValidator(_validate_torch_tensor_args)] = Field(
+        default=None,
+        description="Positional inputs to the model. Required except when `exporter` is `optimum`.",
         exclude=True
     )
     input_shapes: list[tuple] | None = Field(
@@ -46,8 +65,8 @@ class PyTorchInferenceMetadataBase(BaseModel, **ConfigDict(arbitrary_types_allow
         description="Model output dictionary keys. Use this if the model returns a dictionary.",
         exclude=True
     )
-    exporter: TorchExporter | None = Field(
+    optimum_config: Annotated[object | None, BeforeValidator(_validate_optimum_exporter_config)] = Field(
         default=None,
-        description="PyTorch exporter to use.",
+        description="Optimum exporter configuration. Required when `exporter` is `optimum`.",
         exclude=True
     )