API Reference / utils /

serialization.py

Persist trained models to disk, and export them to other runtimes.

Reach for this module whenever a fitted TuiML model (or any picklable Python object) needs to survive past the current process: hand it save_model / load_model for a quick round trip, ModelSerializer when you want to reuse one format/protocol configuration across many saves, ModelCheckpoint to snapshot a model periodically during a training loop, or export_to_onnx to hand the model to a non-Python runtime.

Overview

  1. Formats: every format ultimately stores a pickle-serialized
dict of {model, metadata, saved_at, model_class, model_module, params}. "pickle" writes it directly; "compressed" gzips the pickle stream; "joblib" defers to the optional joblib package, which is more efficient for models holding large NumPy arrays.
  1. Metadata tracking: save (and
save_model) automatically capture the save timestamp, the model's class/module, and its constructor parameters (via get_params() if available) alongside any caller-supplied metadata dict. load_model_info gives convenient access to this record; note it still unpickles the whole saved bundle (model included) internally, it just returns a lighter dict.
  1. Checkpointing: ModelCheckpoint wraps save_model to
snapshot a model during long-running training, optionally keeping only the best-scoring checkpoint and pruning old ones.

Security

Every format here is pickle-based (joblib also pickles under the hood). Unpickling executes arbitrary code embedded in the file. Only load files you created yourself or otherwise trust — never call load_model, load, or load_latest/load_best on a file from an untrusted source.

Classes

ModelSerializer

class utils.serialization.ModelSerializer

Unified interface for machine learning model serialization.

Provides a consistent API for persistent storage and retrieval of trained models using several backend formats. It handles automatic metadata tracking (e.g., save time, model class, hyperparameters).
Constructor 
__init__(
    self,
    format: str = 'pickle',
    protocol: int = 4,
)

Overview

The serializer supports the following backend formats:

  • pickle: Python's native serialization format. Best for general use.
  • joblib: Optimized for models containing large NumPy arrays.
  • compressed: Gzip-compressed pickle for significant disk space savings.

Parameters

format
{"pickle", "joblib", "compressed"} = "pickle"
The serialization strategy to employ.
protocol
int = 4
The pickle protocol version to use (typically 0-5). Higher versions are more efficient but less compatible with older Python versions.

Attributes

format
str
The active serialization format.
protocol
int
The active pickle protocol version.

Raises

ValueError
If format is not one of SUPPORTED_FORMATS.

Notes

Every format is pickle-based under the hood (joblib also pickles). Only call load on files you trust — see the module-level "Security" note.

Save and reload a fitted model together with custom metadata:

python 
>>> import numpy as np
>>> from tuiml.algorithms.bayesian import NaiveBayesClassifier
>>> from tuiml.utils.serialization import ModelSerializer
>>> X = np.array([[0, 0], [1, 1], [2, 2], [3, 3]])
>>> y = np.array([0, 0, 1, 1])
>>> my_model = NaiveBayesClassifier().fit(X, y)
>>> serializer = ModelSerializer(format="compressed")
>>> meta = {"accuracy": 0.98, "dataset": "iris_v1"}
>>> serializer.save(my_model, "model.pkl.gz", metadata=meta)  # doctest: +SKIP
'model.pkl.gz'
>>> loaded_model = serializer.load("model.pkl.gz")  # doctest: +SKIP

Methods

get_parameter_schema (cls) -> dict

Return JSON Schema for parameters.

save (self, model: Any, filepath: Union[str, Path], metadata: Optional[Dict]=None) -> str

Save a trained model to disk in this instance's format.

Parameters
model
Any
Trained model object to save. Must be picklable.
filepath
str or Path
Path to save the model to. For format="compressed", a .gz suffix is appended to the file actually written on disk if filepath doesn't already end in .gz.
metadata
dict
Additional metadata to store alongside the model.
Returns
filepath
str
The filepath argument, unchanged, as a string.
Raises
ImportError
If format="joblib" and the optional joblib package is not installed.
load (self, filepath: Union[str, Path]) -> Any

Load a model from disk, ignoring any stored metadata.

Parameters
filepath
str or Path
Path to the saved model.
Returns
model
Any
Loaded model object.
Raises
FileNotFoundError
If filepath does not exist.
ImportError
If the file requires joblib and it is not installed.
load_with_metadata (self, filepath: Union[str, Path]) -> Dict[str, Any]

Load a model together with everything save recorded about it.

Parameters
filepath
str or Path
Path to the saved model.
Returns
data
dict
Dictionary with keys "model", "metadata", "params", "saved_at", "format", "model_class", and "model_module".
Raises
FileNotFoundError
If filepath does not exist (raised by the underlying file open, unlike load this method does not pre-check existence).
ImportError
If the file requires joblib and it is not installed.

ModelCheckpoint

class utils.serialization.ModelCheckpoint

Manage model persistence during iterative training processes.

Automatically handles saving model snapshots at regular intervals, tracking the 'best' model based on a chosen metric, and cleaning up disk space by keeping only a limited number of recent checkpoints.
Constructor 
__init__(
    self,
    directory: Union[str, Path],
    prefix: str = 'checkpoint',
    max_to_keep: int = 5,
    save_best_only: bool = False,
    monitor: str = 'loss',
    mode: str = 'min',
)

Parameters

directory
str or Path
The directory where checkpoint files will be written.
prefix
str = "checkpoint"
The filename prefix for all generated snapshots.
max_to_keep
int = 5
The rolling limit of recent checkpoints to retain.
save_best_only
bool = False
If True, a new checkpoint is only committed if the monitored metric improves.
monitor
str = "loss"
The key name of the metric to track (passed to save).
mode
{"min", "max"} = "min"
Whether to look for the minimum (e.g., loss) or maximum (e.g., accuracy).

Attributes

best_value
float
The best recorded value of the monitored metric so far (inf / -inf until the first checkpoint when save_best_only=True).
checkpoints
list of Path
A chronologically ordered list of active checkpoint file paths (only files written by this instance; not populated from disk).

Notes

Checkpoints are always written with save_model's default format="pickle". Unpickling runs arbitrary code embedded in the file — only call load_latest/load_best on a directory you trust.

Keep only the 3 best checkpoints by F1-score:

python 
>>> import tempfile
>>> import numpy as np
>>> from tuiml.algorithms.bayesian import NaiveBayesClassifier
>>> from tuiml.utils.serialization import ModelCheckpoint
>>> X = np.array([[0, 0], [1, 1], [2, 2], [3, 3]])
>>> y = np.array([0, 0, 1, 1])
>>> my_model = NaiveBayesClassifier().fit(X, y)
>>> checkpoint = ModelCheckpoint(
...     directory=tempfile.mkdtemp(),
...     save_best_only=True,
...     monitor="f1",
...     mode="max",
...     max_to_keep=3,
... )  # doctest: +SKIP
>>> for i, score in enumerate([0.5, 0.6, 0.55, 0.7, 0.65]):
...     checkpoint.save(my_model, epoch=i, f1=score)  # doctest: +SKIP
>>> round(checkpoint.best_value, 2)  # doctest: +SKIP
0.7
>>> len(checkpoint.checkpoints)  # doctest: +SKIP
3

Methods

get_parameter_schema (cls) -> dict

Return JSON Schema for parameters.

save (self, model: Any, epoch: int=None, **metrics) -> Optional[str]

Save a checkpoint, unless save_best_only says to skip it.

Parameters
model
Any
Model to checkpoint.
epoch
int
Current epoch number, included in the checkpoint filename.
**metrics
Metric values (e.g., accuracy=0.95, loss=0.1), stored in the checkpoint's metadata. Must include the key named by self.monitor when save_best_only=True.
Returns
filepath
str or None
Path to the saved checkpoint, or None if save_best_only=True and self.monitor's value did not improve on best_value (so nothing was written).
Raises
ValueError
If save_best_only=True and metrics does not contain the key named by self.monitor.
load_latest (self) -> Any

Load the most recently saved checkpoint.

Returns
Any
The loaded model object.
Raises
FileNotFoundError
If checkpoints is empty and no matching files are found in directory.
load_best (self) -> Any

Load the best checkpoint saved so far.

Returns
Any
The loaded model object.
Raises
ValueError
If save_best_only=False.
FileNotFoundError
If no checkpoint has been saved yet.

Functions

Func

save_model

Line 376
save_model(model: Any, filepath: Union[str, Path], format: str='pickle', metadata: Optional[Dict]=None, compress: bool=False) -> str

Save a trained model to disk with automatic metadata tracking.

Thin functional wrapper around ModelSerializer: builds one with the resolved format and delegates to save.

Parameters

model
Any
The trained model instance (or any picklable object) to save.
filepath
str or Path
The destination file path.
format
{"pickle", "joblib", "compressed"} = "pickle"
The underlying serialization format. Ignored if compress=True (which forces "compressed").
metadata
dict
A dictionary of custom metadata (e.g., metrics, data version) to store alongside the model.
compress
bool = False
If True, applies gzip compression to the output file (shorthand for format="compressed").

Returns

str
The filepath argument, unchanged, as a string. Note: if compress=True (or format="compressed") and filepath does not already end in .gz, the file written to disk has .gz appended, but this return value does not reflect that — pass a path already ending in .gz to avoid the mismatch.

Raises

ImportError
If format="joblib" and the optional joblib package is not installed.

See Also

load_model

The inverse operation.

Save a fitted model with gzip compression:

python 
>>> import numpy as np
>>> from tuiml.algorithms.bayesian import NaiveBayesClassifier
>>> from tuiml.utils.serialization import save_model
>>> X = np.array([[0, 0], [1, 1], [2, 2], [3, 3]])
>>> y = np.array([0, 0, 1, 1])
>>> my_clf = NaiveBayesClassifier().fit(X, y)
>>> save_model(my_clf, "classifier.pkl.gz", compress=True)  # doctest: +SKIP
'classifier.pkl.gz'
Func

load_model

Line 444
load_model(filepath: Union[str, Path]) -> Any

Load a model from disk, discarding any stored metadata.

Detects the format from filepath's extension: .gz is loaded as gzip-compressed pickle, .joblib/.jbl via joblib, and anything else as a plain pickle.

Parameters

filepath
str or Path
Path to the saved model file.

Returns

model
Any
The reconstructed model object.

Raises

FileNotFoundError
If filepath does not exist.
ImportError
If the file requires joblib and it is not installed.

Notes

Unpickling runs arbitrary code embedded in the file. Only load files you trust (see the module-level "Security" note).
python 
>>> import numpy as np
>>> from tuiml.algorithms.bayesian import NaiveBayesClassifier
>>> from tuiml.utils.serialization import save_model, load_model
>>> X = np.array([[0, 0], [1, 1], [2, 2], [3, 3]])
>>> y = np.array([0, 0, 1, 1])
>>> save_model(NaiveBayesClassifier().fit(X, y), "classifier.pkl")  # doctest: +SKIP
>>> model = load_model("classifier.pkl")  # doctest: +SKIP
>>> model.predict(X[:1])  # doctest: +SKIP
array([0])
Func

load_model_info

Line 488
load_model_info(filepath: Union[str, Path]) -> Dict[str, Any]

Load a saved model's metadata and hyper-parameters, without the model itself.

Useful for inspecting model properties, comparing performance, or retrieving training timestamps when only the recorded metadata is needed, so the caller's code doesn't have to unpack load_with_metadata's dict by hand.

Parameters

filepath
str or Path
Path to the saved model file.

Returns

dict

A dictionary with keys:

  • •"model_class": Name of the Python class.
  • •"model_module": Fully-qualified module the class lives in.
  • •"params": Configuration parameters/hyper-parameters (from

model.get_params() if available, else public __dict__ attributes).

  • •"metadata": The custom metadata dict provided during save.
  • •"saved_at": ISO timestamp of the save operation.
  • •"format": The serialization format used ("pickle",
"joblib", or "compressed").

Raises

FileNotFoundError
If filepath does not exist.
ImportError
If the file requires joblib and it is not installed.

Notes

This still unpickles the whole saved bundle (model included) — it only returns a lighter dict. It offers no memory savings over load_model, but is convenient when metadata and the model object are both wanted, or when only the metadata is of interest and reading it is simpler this way.

Check the accuracy of a saved model:

python 
>>> import numpy as np
>>> from tuiml.algorithms.bayesian import NaiveBayesClassifier
>>> from tuiml.utils.serialization import save_model, load_model_info
>>> X = np.array([[0, 0], [1, 1], [2, 2], [3, 3]])
>>> y = np.array([0, 0, 1, 1])
>>> save_model(NaiveBayesClassifier().fit(X, y), "model.pkl",
...            metadata={"accuracy": 0.98})  # doctest: +SKIP
>>> info = load_model_info("model.pkl")  # doctest: +SKIP
>>> print(f"Model: {info['model_class']} | Acc: {info['metadata']['accuracy']}")  # doctest: +SKIP
Model: NaiveBayesClassifier | Acc: 0.98
Func

export_to_onnx

Line 828
export_to_onnx(model: Any, filepath: Union[str, Path], input_shape: tuple, input_names: list=None, output_names: list=None) -> str

Export a trained model to ONNX format for cross-platform interoperability.

Overview

ONNX (Open Neural Network Exchange) allows models trained in Python to be executed in high-performance runtimes like ONNX Runtime, C++, or mobile devices.

Parameters

model
Any
A Scikit-learn compatible model or pipeline to be converted.
filepath
str or Path
The output path for the .onnx file.
input_shape
tuple
The expected input dimensions (e.g., (None, 20) for a batch size of 20 features). Only the last element (the feature count) is currently used.
input_names
list of str
Custom names for the input nodes. Not currently applied — the input node is always named "input" regardless of this argument; see Notes.
output_names
list of str
Custom names for the output nodes. Not currently applied — this argument is accepted but not passed through to the underlying converter; see Notes.

Returns

str
The path to the successfully exported ONNX model, with a .onnx suffix appended if filepath didn't already have one.

Raises

ImportError
If the optional skl2onnx package is not installed.

Notes

This utility requires the skl2onnx package. Install via: pip install skl2onnx

input_names renames the single input node of the exported graph. output_names is not supported yet and raises NotImplementedError rather than being silently ignored, so a caller never ends up with a graph whose node names differ from what it asked for.