ONNX Export#

Export a fitted Pipeline (or single transformer) to ONNX and run inference with ONNX Runtime.

Every converter follows the same tensor naming convention: inputs are named {col}__in and outputs use the resulting Gators column name. Dtypes are derived from each transformer’s _input_dtypes/_output_dtypes (Polars String -> ONNX STRING, Boolean -> BOOL, integer/temporal dtypes -> INT64, floats -> FLOAT/DOUBLE per float_datatype).

Graph Export#

gators.onnx_converters.to_onnx_graph(transformer: gators.transformer._base_transformer._BaseTransformer, errors: str = 'raise', float_datatype: str = 'float64') → onnx.onnx_ml_pb2.ModelProto[source]#

Export a single fitted Gators transformer to a self-contained ONNX ModelProto.

All columns flow through the graph. Columns not in the transformer’s subset pass through via Identity nodes. The graph is validated by the ONNX checker before being returned.

Input tensor naming convention : {col}__in Output tensor naming convention: {col} (original column names, inplace semantics)

Parameters:

  • transformer (gators.transformer._base_transformer._BaseTransformer) – A fitted Gators transformer instance.

  • errors ({'coerce', 'raise'}, default 'coerce') – How to handle unsupported transformers or strategies. ‘raise’ – raise OnnxNotSupportedError immediately. ‘coerce’ – emit Identity pass-through nodes so the graph stays structurally valid.

Returns:

Validated ONNX model.

Return type:

onnx.ModelProto

Examples

>>> import polars as pl
>>> import numpy as np
>>> from gators.imputers import NumericImputer
>>> from gators.onnx import to_onnx_graph
>>>
>>> X = pl.DataFrame({"A": [1.0, None, 3.0], "B": [4.0, 5.0, None]})
>>> imputer = NumericImputer(strategy="median")
>>> imputer.fit(X)
>>> model = to_onnx_graph(imputer)
gators.onnx_converters.pipeline_to_onnx(pipeline: gators.pipeline.pipeline.Pipeline, errors: str = 'raise', float_datatype: str = 'float64') → onnx.onnx_ml_pb2.ModelProto[source]#

Export a fitted Gators Pipeline to a single ONNX ModelProto.

Each pipeline step contributes its nodes to the same graph, chained via intermediate tensor names. Column additions (IsNull, MathFeatures, …), column drops (DropColumns, …), and renames (RenameColumns) are all supported.

Steps with inplace=False and drop_columns=True are rejected (errors='raise') or skipped with Identity pass-throughs (errors='coerce') because they would silently break the column ordering.

Input tensor naming convention : {col}__in Output tensor naming convention: {col}

Parameters:

  • pipeline (fitted gators.Pipeline)

  • errors ({'raise', 'coerce'}, default 'raise') – How to handle unsupported steps or strategies. 'raise' – raise OnnxNotSupportedError immediately. 'coerce' – emit Identity pass-through nodes so the graph stays valid.

Returns:

Validated ONNX model.

Return type:

onnx.ModelProto

Examples

>>> import polars as pl
>>> from gators.pipeline import Pipeline
>>> from gators.imputers import NumericImputer
>>> from gators.scalers import StandardScaler
>>> from gators.onnx import pipeline_to_onnx
>>>
>>> X = pl.DataFrame({"A": [1.0, None, 3.0], "B": [4.0, 5.0, None]})
>>> pipe = Pipeline(steps=[
...     ("imputer", NumericImputer(strategy="median")),
...     ("scaler", StandardScaler()),
... ])
>>> pipe.fit(X)
>>> model = pipeline_to_onnx(pipe)
gators.onnx_converters.pipeline_to_scoring_onnx(pipeline: gators.pipeline.pipeline.Pipeline, model_onnx: onnx.onnx_ml_pb2.ModelProto, feature_columns: list[str], errors: str = 'raise', float_datatype: str = 'float64') → onnx.onnx_ml_pb2.ModelProto[source]#

Chain a Gators preprocessing Pipeline with an ONNX ML model (XGBoost/LightGBM).

The preprocessing graph outputs one 1D float[N] tensor per column. This function inserts a reshape bridge that Unsqueezes each selected column to float[N, 1] and Concatenates them into the float[N, n_features] matrix the ML model expects, then stitches everything into a single graph.

Parameters:

  • pipeline (fitted gators.Pipeline) – Preprocessing pipeline (same one used when training the ML model).

  • model_onnx (onnx.ModelProto) – The ML model exported to ONNX (e.g. via onnxmltools.convert_xgboost or onnxmltools.convert_lightgbm).

  • feature_columns (list[str]) – Ordered list of preprocessing output column names to stack into the feature matrix. Must exactly match the column order seen by the model at training time.

  • errors ({'coerce', 'raise'}, default 'coerce') – Forwarded to pipeline_to_onnx.

Returns:

Single validated ONNX model: preprocessing → reshape bridge → scoring. Inputs : same as the preprocessing pipeline ({col}__in tensors). Outputs : same as the ML model (labels / probabilities).

Return type:

onnx.ModelProto

Examples

>>> import lightgbm as lgb
>>> from onnxmltools import convert_lightgbm
>>> from onnxmltools.utils import float_array_type
>>> from gators.onnx_converters import pipeline_to_scoring_onnx
>>>
>>> lgb_model = lgb.train(...)
>>> n_features = len(feature_columns)
>>> model_onnx = convert_lightgbm(lgb_model, initial_types=[
...     ("input", float_array_type([None, n_features]))
... ])
>>> scoring_model = pipeline_to_scoring_onnx(pipe, model_onnx, feature_columns)

Compatibility Check#

gators.onnx_converters.check_pipeline_onnx_compatibility(pipeline: gators.pipeline.pipeline.Pipeline) → dict[str, bool][source]#

Return the ONNX compatibility status of every step in a fitted Gators Pipeline.

Parameters:

pipeline (gators.Pipeline) – A pipeline whose steps attribute is a list of (name, transformer) tuples.

Returns:

{step_name: True} when a converter is registered for the transformer, {step_name: False} when only the fallback (Identity pass-through) would be used.

Return type:

dict[str, bool]

Examples

>>> from gators.pipeline import Pipeline
>>> from gators.imputers import NumericImputer
>>> from gators.scalers import StandardScaler
>>> pipe = Pipeline(steps=[
...     ("imputer", NumericImputer(strategy="mean")),
...     ("scaler",  StandardScaler()),
... ])
>>> from gators.onnx_converters import check_pipeline_onnx_compatibility
>>> check_pipeline_onnx_compatibility(pipe)
{'imputer': True, 'scaler': True}

Inference#

gators.onnx_converters.create_session(model: onnx.ModelProto | str | bytes, providers: list[str] | None = None, intra_op_num_threads: int = 0, inter_op_num_threads: int = 1, optimized_model_path: str | None = None) → ort.InferenceSession[source]#

Create an optimised OnnxRuntime InferenceSession from a Gators ONNX model.

Applies ORT_ENABLE_ALL graph optimisations (constant folding, op fusion, layout transformations) and optionally saves the pre-optimised graph to disk so subsequent loads skip the optimisation pass entirely.

Parameters:

  • model (onnx.ModelProto | str | bytes) – The ONNX model — a ModelProto object, a serialised byte string, or a file path to a .onnx file.

  • providers (list[str] or None, default None) – Execution provider priority list. None delegates to ORT’s default (["CPUExecutionProvider"]). Pass ["CoreMLExecutionProvider", "CPUExecutionProvider"] on macOS or ["CUDAExecutionProvider", "CPUExecutionProvider"] on CUDA hosts for hardware acceleration.

  • intra_op_num_threads (int, default 0) – Thread count within a single operator (0 = use all available cores).

  • inter_op_num_threads (int, default 1) – Thread count for executing operators in parallel (1 = sequential, best for single-request latency).

  • optimized_model_path (str or None, default None) – When set, ORT writes the pre-optimised ONNX graph to this path. On subsequent runs you can load that file directly and skip re-optimisation. Saved graphs are capped at ORT_ENABLE_EXTENDED (skipping the hardware-specific NCHWc layout transformer) so the file stays portable across machines; in-memory-only sessions (no path given) still get the full ORT_ENABLE_ALL optimisations.

Returns:

A ready-to-use session with all graph optimisations applied.

Return type:

ort.InferenceSession

Examples

>>> from gators.onnx_converters import pipeline_to_onnx, create_session, run_session
>>> model = pipeline_to_onnx(fitted_pipeline)
>>> sess = create_session(model, optimized_model_path="pipeline_opt.onnx")
>>> result = run_session(sess, X_test, batch_size=10_000)
gators.onnx_converters.run_session(session: ort.InferenceSession, X: pl.DataFrame, batch_size: int = 0) → pl.DataFrame[source]#

Run an ORT session on a Polars DataFrame, with optional batching.

Automatically converts each Polars column to the tensor type declared by the session’s input metadata, runs inference (in chunks of batch_size rows when requested), and returns a Polars DataFrame matching the session’s output column names.

Parameters:

  • session (ort.InferenceSession) – A session created by create_session() or any ort.InferenceSession.

  • X (pl.DataFrame) – Input DataFrame. Column names must match {col}__in for each input tensor declared by the session (the Gators ONNX naming convention).

  • batch_size (int, default 0) – Number of rows per inference call. 0 (or any value ≥ len(X)) runs the whole DataFrame in a single call — best when memory allows. Use a positive value to cap peak memory usage on very large datasets.

Returns:

DataFrame whose columns correspond to the session’s output tensors.

Return type:

pl.DataFrame

Examples

>>> sess  = create_session(pipeline_to_onnx(pipe))
>>> result = run_session(sess, X_test)                    # single pass
>>> result = run_session(sess, X_test, batch_size=50_000) # batched

Exceptions#

class gators.onnx_converters.OnnxNotSupportedError[source]#

Bases: Exception

Raised when a transformer or strategy cannot be expressed in a standard ONNX graph.