From 80279389b321466ac299eeee2825a08a68081e66 Mon Sep 17 00:00:00 2001 From: Sarp Tan Doven <119648987+sarptandoven@users.noreply.github.com> Date: Sun, 5 Jul 2026 19:40:00 -0800 Subject: [PATCH] add json inference model serialization --- README.md | 11 ++ ngboost/__init__.py | 8 +- ngboost/helpers.py | 275 +++++++++++++++++++++++++++++++++++++++++ tests/test_pickling.py | 28 ++++- 4 files changed, 320 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8772c086..c11d12f3 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,17 @@ test_NLL = -Y_dists.logpdf(Y_test).mean() print('Test NLL', test_NLL) ``` +Fitted NGBoost models with tree base learners can also be exported to JSON for +portable inference without pickle: + +```python +from ngboost import load_ngboost_model_json, save_ngboost_model_json + +save_ngboost_model_json(ngb, "ngboost-model.json") +restored = load_ngboost_model_json("ngboost-model.json") +Y_preds = restored.predict(X_test) +``` + Details on available distributions, scoring rules, learners, tuning, and model interpretation are available in our [user guide](https://stanfordmlgroup.github.io/ngboost/intro.html), which also includes numerous usage examples and information on how to add new distributions or scores to NGBoost. ## License diff --git a/ngboost/__init__.py b/ngboost/__init__.py index 67f6c548..f54c891d 100644 --- a/ngboost/__init__.py +++ b/ngboost/__init__.py @@ -7,7 +7,11 @@ from importlib_metadata import version from .api import NGBClassifier, NGBRegressor, NGBSurvival -from .helpers import load_ngboost_model +from .helpers import ( + load_ngboost_model, + load_ngboost_model_json, + save_ngboost_model_json, +) from .ngboost import NGBoost __all__ = [ @@ -16,6 +20,8 @@ "NGBSurvival", "NGBoost", "load_ngboost_model", + "load_ngboost_model_json", + "save_ngboost_model_json", ] __version__ = version(__name__) diff --git a/ngboost/helpers.py b/ngboost/helpers.py index ea3750d1..e1ea73a0 100644 --- a/ngboost/helpers.py +++ b/ngboost/helpers.py @@ -1,3 +1,5 @@ +import importlib +import json import sys import threading import types as _types @@ -5,8 +7,11 @@ import joblib import numpy as np import sklearn.tree._tree as _sklearn_tree # pylint: disable=c-extension-no-member +from sklearn.tree import DecisionTreeRegressor from sklearn.utils import check_array +from ngboost.manifold import manifold + # --------------------------------------------------------------------------- # Backward compatibility helpers (issue #389) # --------------------------------------------------------------------------- @@ -113,6 +118,276 @@ def load_ngboost_model(filepath): return model +# --------------------------------------------------------------------------- +# JSON inference serialization helpers (issue #392) +# --------------------------------------------------------------------------- + +_JSON_FORMAT = "ngboost-json-inference" +_JSON_VERSION = 1 + + +def _qualified_name(obj): + return f"{obj.__module__}.{obj.__qualname__}" + + +def _import_qualified_name(name): + module_name, _, attr_name = name.rpartition(".") + if not module_name: + raise ValueError(f"Invalid qualified name: {name!r}") + module = importlib.import_module(module_name) + obj = module + for attr in attr_name.split("."): + obj = getattr(obj, attr) + return obj + + +def _encode_distribution(dist): + if getattr(dist, "__name__", None) == "Categorical": + return {"kind": "categorical", "K": dist.n_params + 1} + return {"kind": "qualified", "name": _qualified_name(dist)} + + +def _decode_distribution(payload): + if payload["kind"] == "categorical": + from ngboost.distns import ( # pylint: disable=import-outside-toplevel + k_categorical, + ) + + return k_categorical(payload["K"]) + if payload["kind"] == "qualified": + return _import_qualified_name(payload["name"]) + raise ValueError(f"Unknown distribution payload: {payload['kind']!r}") + + +def _encode_json_value(value): + if isinstance(value, np.ndarray): + payload = { + "__ndarray__": True, + "shape": value.shape, + } + if value.dtype.names: + payload["dtype_struct"] = { + "names": list(value.dtype.names), + "formats": [ + value.dtype.fields[name][0].str for name in value.dtype.names + ], + "offsets": [value.dtype.fields[name][1] for name in value.dtype.names], + "itemsize": value.dtype.itemsize, + } + payload["fields"] = { + name: value[name].tolist() for name in value.dtype.names + } + else: + payload["dtype"] = str(value.dtype) + payload["data"] = value.tolist() + return payload + if isinstance(value, np.generic): + return value.item() + if isinstance(value, tuple): + return {"__tuple__": True, "items": [_encode_json_value(v) for v in value]} + if isinstance(value, list): + return [_encode_json_value(v) for v in value] + if isinstance(value, dict): + return {key: _encode_json_value(val) for key, val in value.items()} + return value + + +def _decode_json_value(value): + if isinstance(value, list): + return [_decode_json_value(v) for v in value] + if not isinstance(value, dict): + return value + if value.get("__ndarray__"): + dtype = np.dtype( + { + "names": value["dtype_struct"]["names"], + "formats": value["dtype_struct"]["formats"], + "offsets": value["dtype_struct"]["offsets"], + "itemsize": value["dtype_struct"]["itemsize"], + } + if "dtype_struct" in value + else value["dtype"] + ) + if "dtype_struct" in value: + array = np.zeros(value["shape"], dtype=dtype) + for name, field_values in value["fields"].items(): + array[name] = field_values + return array + array = np.array(value["data"], dtype=dtype) + return array.reshape(value["shape"]) + if value.get("__tuple__"): + return tuple(_decode_json_value(v) for v in value["items"]) + return {key: _decode_json_value(val) for key, val in value.items()} + + +def _serialize_decision_tree(estimator): + if not isinstance(estimator, DecisionTreeRegressor): + raise TypeError( + "JSON inference serialization currently supports fitted " + "DecisionTreeRegressor base learners only." + ) + if not hasattr(estimator, "tree_"): + raise ValueError("Cannot serialize an unfitted DecisionTreeRegressor.") + + attrs = {} + for name in ( + "n_features_in_", + "n_outputs_", + "max_features_", + "feature_names_in_", + ): + if hasattr(estimator, name): + attrs[name] = _encode_json_value(getattr(estimator, name)) + + return { + "class": _qualified_name(estimator.__class__), + "params": _encode_json_value(estimator.get_params(deep=False)), + "attrs": attrs, + "tree": { + "n_features": estimator.tree_.n_features, + "n_classes": _encode_json_value(estimator.tree_.n_classes), + "n_outputs": estimator.tree_.n_outputs, + "state": _encode_json_value(estimator.tree_.__getstate__()), + }, + } + + +def _deserialize_decision_tree(payload): + estimator_class = _import_qualified_name(payload["class"]) + if estimator_class is not DecisionTreeRegressor: + raise TypeError( + "Only sklearn.tree.DecisionTreeRegressor JSON payloads are supported." + ) + + estimator = estimator_class(**_decode_json_value(payload["params"])) + for name, value in payload["attrs"].items(): + setattr(estimator, name, _decode_json_value(value)) + + tree_payload = payload["tree"] + n_classes = _decode_json_value(tree_payload["n_classes"]) + if isinstance(n_classes, int): + n_classes = np.array([n_classes], dtype=np.intp) + else: + n_classes = np.asarray(n_classes, dtype=np.intp) + tree = _sklearn_tree.Tree( # pylint: disable=c-extension-no-member + tree_payload["n_features"], n_classes, tree_payload["n_outputs"] + ) + tree.__setstate__(_decode_json_value(tree_payload["state"])) + estimator.tree_ = tree + return estimator + + +def _serialize_base_config(base): + if isinstance(base, (list, tuple)): + return { + "kind": "sequence", + "sequence_type": type(base).__name__, + "items": [_serialize_base_config(item) for item in base], + } + if isinstance(base, DecisionTreeRegressor): + return { + "kind": "decision_tree_regressor", + "params": _encode_json_value(base.get_params(deep=False)), + } + raise TypeError( + "JSON inference serialization currently supports DecisionTreeRegressor " + "base learner configuration only." + ) + + +def _deserialize_base_config(payload): + if payload["kind"] == "sequence": + items = [_deserialize_base_config(item) for item in payload["items"]] + return tuple(items) if payload.get("sequence_type") == "tuple" else items + if payload["kind"] == "decision_tree_regressor": + return DecisionTreeRegressor(**_decode_json_value(payload["params"])) + raise ValueError(f"Unknown base learner payload: {payload['kind']!r}") + + +def save_ngboost_model_json(model, filepath): + """Save a fitted NGBoost model to a JSON file for inference. + + The JSON payload stores only the fitted state needed by ``predict``, + ``pred_dist``, and classifier ``predict_proba``. It avoids pickle while + preserving exact sklearn decision-tree base learner state. + """ + if not getattr(model, "base_models", None): + raise ValueError("Cannot JSON serialize an unfitted NGBoost model.") + + payload = { + "format": _JSON_FORMAT, + "version": _JSON_VERSION, + "model_class": _qualified_name(model.__class__), + "dist": _encode_distribution(model.Dist), + "score": _qualified_name(model.Score), + "base": _serialize_base_config(model.Base), + "params": { + "natural_gradient": model.natural_gradient, + "n_estimators": model.n_estimators, + "learning_rate": model.learning_rate, + "minibatch_frac": model.minibatch_frac, + "col_sample": model.col_sample, + "verbose": model.verbose, + "verbose_eval": model.verbose_eval, + "tol": model.tol, + "validation_fraction": model.validation_fraction, + "early_stopping_rounds": model.early_stopping_rounds, + }, + "state": { + "init_params": _encode_json_value(model.init_params), + "n_features": model.n_features, + "best_val_loss_itr": model.best_val_loss_itr, + "multi_output": model.multi_output, + "estimator_type": getattr(model, "_estimator_type", None), + "scalings": _encode_json_value(model.scalings), + "col_idxs": _encode_json_value(model.col_idxs), + "evals_result": _encode_json_value(getattr(model, "evals_result", {})), + "base_models": [ + [_serialize_decision_tree(estimator) for estimator in iter_models] + for iter_models in model.base_models + ], + }, + } + + with open(filepath, "w", encoding="utf-8") as f: + json.dump(payload, f, sort_keys=True, separators=(",", ":")) + + +def load_ngboost_model_json(filepath): + """Load a fitted NGBoost model saved by ``save_ngboost_model_json``.""" + with open(filepath, "r", encoding="utf-8") as f: + payload = json.load(f) + + if payload.get("format") != _JSON_FORMAT or payload.get("version") != _JSON_VERSION: + raise ValueError("Unsupported NGBoost JSON model format.") + + model_class = _import_qualified_name(payload["model_class"]) + model = model_class.__new__(model_class) + model.Dist = _decode_distribution(payload["dist"]) + model.Score = _import_qualified_name(payload["score"]) + model.Base = _deserialize_base_config(payload["base"]) + model.Manifold = manifold(model.Score, model.Dist) + model.random_state = None + + for name, value in payload["params"].items(): + setattr(model, name, value) + + state = payload["state"] + model.init_params = _decode_json_value(state["init_params"]) + model.n_features = state["n_features"] + model.best_val_loss_itr = state["best_val_loss_itr"] + model.multi_output = state["multi_output"] + model._estimator_type = state["estimator_type"] # pylint: disable=protected-access + model.scalings = _decode_json_value(state["scalings"]) + model.col_idxs = _decode_json_value(state["col_idxs"]) + model.evals_result = _decode_json_value(state["evals_result"]) + model.base_models = [ + [_deserialize_decision_tree(estimator) for estimator in iter_models] + for iter_models in state["base_models"] + ] + return model + + # --------------------------------------------------------------------------- diff --git a/tests/test_pickling.py b/tests/test_pickling.py index a21083c1..db4006fb 100644 --- a/tests/test_pickling.py +++ b/tests/test_pickling.py @@ -9,7 +9,14 @@ import sklearn.tree._tree as _sklearn_tree # pylint: disable=c-extension-no-member from sklearn.tree import DecisionTreeRegressor -from ngboost import NGBClassifier, NGBRegressor, NGBSurvival, load_ngboost_model +from ngboost import ( + NGBClassifier, + NGBRegressor, + NGBSurvival, + load_ngboost_model, + load_ngboost_model_json, + save_ngboost_model_json, +) from ngboost.distns import MultivariateNormal @@ -59,6 +66,25 @@ def test_model_save(learners_data): assert (new_preds == preds).all() +def test_json_inference_roundtrip_preserves_predictions(learners_data): + """JSON export stores enough fitted state for inference without pickle.""" + + for learner, data, preds in learners_data[:2]: # regressor and classifier paths + with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f: + tmp_path = f.name + try: + save_ngboost_model_json(learner, tmp_path) + model = load_ngboost_model_json(tmp_path) + new_preds = model.predict(data) + assert np.allclose(new_preds, preds) + if isinstance(learner, NGBClassifier): + assert np.allclose( + model.predict_proba(data), learner.predict_proba(data) + ) + finally: + os.unlink(tmp_path) + + # --------------------------------------------------------------------------- # Helpers for backward-compatibility test (issue #389) # ---------------------------------------------------------------------------