yobx.torch.in_transformers.exporter#

Provides YobxOnnxExporter, a subclass of transformers.exporters.OnnxExporter that replaces the default torch.onnx.export / onnxscript backend with the yobx converter.

This lets callers use the standard transformers exporter API while benefiting from yobx’s graph-builder optimisations and operator coverage.

Usage:

from transformers.exporters import OnnxConfig
from yobx.torch.in_transformers import YobxOnnxExporter

exporter = YobxOnnxExporter()
artifact = exporter.export(model, sample_inputs, config=OnnxConfig(dynamic=True))
# artifact is an ExportArtifact – save it, inspect the proto, etc.
artifact.save("model.onnx")
class yobx.torch.in_transformers.exporter.YobxOnnxExporter(target_opset: int | None = None, **kwargs: Any)[source]#

Subclass of transformers.exporters.OnnxExporter that converts a PreTrainedModel to ONNX using the yobx graph builder instead of the default torch.onnx.export / onnxscript pipeline.

When transformers is not installed, instantiating this class raises ImportError.

The interface is identical to the upstream exporter:

from transformers.exporters import OnnxConfig
from yobx.torch.in_transformers import YobxOnnxExporter

exporter = YobxOnnxExporter()
artifact = exporter.export(model, inputs, config=OnnxConfig(dynamic=True))
artifact.save("model.onnx")

Extra keyword arguments passed to the constructor are forwarded verbatim to yobx.torch.to_onnx() during export (e.g. options, dispatcher, export_modules_as_functions, …).

Parameters:

  • target_opset – ONNX opset version to target. Overrides the value carried by config.opset_version when both are supplied.

  • kwargs – extra keyword arguments forwarded verbatim to yobx.torch.to_onnx().

export(model: PreTrainedModel, sample_inputs: MutableMapping[str, Any], config: 'OnnxConfig' | Dict[str, Any]) → ExportArtifact[source]#

Exports model to ONNX using the yobx converter.

Applies the same transformers-side preprocessing as transformers.exporters.DynamoExporter (label stripping, output-flag patching, dynamic-shape inference) and then feeds the resulting torch.export.ExportedProgram to yobx.torch.to_onnx().

Parameters:

  • model – the PreTrainedModel to export.

  • sample_inputs – forward kwargs — what you would pass to model(**sample_inputs). Labels and loss-related keys must not be present (see transformers.exporters.DynamoExporter).

  • config – an OnnxConfig (or a plain dict that will be converted to one) controlling dynamic shapes, output path, opset version, etc.

Returns:

ExportArtifact wrapping the exported ONNX proto.