add json inference model serialization

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

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__)

ngboost/helpers.py

@@ -1,12 +1,17 @@
+import importlib
+import json
 import sys
 import threading
 import types as _types
 
 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
+
+
 # ---------------------------------------------------------------------------
 
 

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)
 # ---------------------------------------------------------------------------