def test_xgboost_save_load(booster_params, metadata,
save_proc): # noqa # pylint: disable
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
bentomodel = save_proc(booster_params,
metadata,
labels=labels,
custom_objects={"func": custom_f})
assert bentomodel.info.metadata is not None
assert_have_file_extension(bentomodel.path, ".json")
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
xgb_loaded = bentoml.xgboost.load(bentomodel.tag,
booster_params=booster_params)
config = json.loads(xgb_loaded.save_config())
if not booster_params:
assert config["learner"]["generic_param"]["nthread"] == str(
psutil.cpu_count())
else:
assert config["learner"]["generic_param"]["nthread"] == str(2)
assert isinstance(xgb_loaded, xgb.Booster)
assert predict_df(xgb_loaded, test_df) == 1
def test_statsmodels_save_load(metadata, holt_model): # noqa # pylint: disable
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
tag = bentoml.statsmodels.save(
TEST_MODEL_NAME,
holt_model,
metadata=metadata,
labels=labels,
custom_objects={"func": custom_f},
)
bentomodel = bentoml.models.get(tag)
assert bentomodel.info.metadata is not None
assert_have_file_extension(bentomodel.path, ".pkl")
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
statsmodels_loaded = bentoml.statsmodels.load(tag)
assert isinstance(
statsmodels_loaded,
statsmodels.tsa.holtwinters.results.HoltWintersResultsWrapper,
)
np.testing.assert_array_equal(holt_model.predict(),
statsmodels_loaded.predict())
def test_tensorflow_v1_save_load(tf1_model_path: Callable[[],
Generator[str, None,
None]]):
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
tag = bentoml.tensorflow_v1.save(
"tensorflow_test",
tf1_model_path,
labels=labels,
custom_objects={"func": custom_f},
)
bentomodel = bentoml.models.get(tag)
assert_have_file_extension(bentomodel.path, ".pb")
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
tf1_loaded = bentoml.tensorflow_v1.load("tensorflow_test")
with tf.get_default_graph().as_default():
tf.global_variables_initializer()
prediction = _model_dunder_call(tf1_loaded, test_tensor)
assert prediction.shape == (1, )
def test_tensorflow_v2_save_load(
mcls: "tf_ext.Module",
tensor: "tf_ext.TensorLike",
predict_fn: Callable[
["tf_ext.AutoTrackable", "tf_ext.TensorLike"], "tf_ext.TensorLike"
],
is_ragged: bool,
):
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
tag = bentoml.tensorflow.save(
MODEL_NAME, mcls, labels=labels, custom_objects={"func": custom_f}
)
bentomodel = bentoml.models.get(tag)
assert_have_file_extension(bentomodel.path, ".pb")
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
model = bentoml.tensorflow.load(MODEL_NAME)
output = predict_fn(model, tensor)
if is_ragged:
assert all(output.numpy() == np.array([[15.0]] * 3))
else:
assert all(output.numpy() == np.array([[15.0]]))
def test_pl_save_load():
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
model: "pl.LightningModule" = AdditionModel()
tag = bentoml.pytorch_lightning.save(
"pytorch_lightning_test",
model,
labels=labels,
custom_objects={"func": custom_f},
)
bentomodel = bentoml.models.get(tag)
assert_have_file_extension(bentomodel.path, ".pt")
assert bentomodel.info.context.get(
"model_format") == "pytorch_lightning:v1"
pl_loaded: "pl.LightningModule" = bentoml.pytorch_lightning.load(tag)
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
assert predict_df(pl_loaded, test_df) == [[6, 5, 4, 3]]
def test_catboost_save_load(
model_params: t.Dict[str, t.Any],
metadata: t.Dict[str, t.Any],
) -> None:
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
tag = save_procedure(
model_params,
metadata,
labels=labels,
custom_objects={"func": custom_f},
)
_model = bentoml.models.get(tag)
assert _model.info.metadata is not None
assert_have_file_extension(_model.path, ".cbm")
cbt_loaded = bentoml.catboost.load(_model.tag, model_params=model_params)
assert isinstance(cbt_loaded, CatBoostClassifier)
assert cbt_loaded.predict(test_df) == np.array([1])
for k in labels.keys():
assert labels[k] == _model.info.labels[k]
assert _model.custom_objects["func"](3) == custom_f(3)
def test_lightgbm_sklearn_save_load(save_sklearn_proc):
bentomodel = save_sklearn_proc(None)
assert_have_file_extension(bentomodel.path, ".pkl")
sklearn_loaded = bentoml.lightgbm.load(bentomodel.tag, )
assert isinstance(sklearn_loaded, lgb.LGBMClassifier)
assert sklearn_loaded.predict(np.array([[0] * 10] * 10)).any() == np.array(
[0])
def test_paddle_save_load(train_paddle_model, input_spec):
tag = bentoml.paddle.save(
"linear_model",
train_paddle_model,
input_spec=input_spec,
)
info = bentoml.models.get(tag)
assert_have_file_extension(info.path, ".pdmodel")
loaded = bentoml.paddle.load(tag)
res1 = predict_df(loaded, test_df)
res2 = np.array([[0.9003858]], dtype=np.float32)
assert np.isclose(res1, res2).all()
def test_mlflow_save_load():
(model, data) = sklearn_model_data()
uri = Path(current_file, "sklearn_clf")
if not uri.exists():
mlflow.sklearn.save_model(model, uri.resolve())
tag = bentoml.mlflow.import_from_uri(MODEL_NAME, str(uri.resolve()))
model_info = bentoml.models.get(tag)
assert_have_file_extension(os.path.join(model_info.path, "sklearn_clf"),
".pkl")
loaded = bentoml.mlflow.load(tag)
np.testing.assert_array_equal(loaded.predict(data), res_arr) # noqa
def test_keras_save_load(
model: "keras.Model",
kwargs: t.Dict[str, t.Any],
) -> None:
tag = bentoml.keras.save(MODEL_NAME, model, **kwargs)
model_info = bentoml.models.get(tag)
if kwargs.get("custom_objects") is not None:
assert_have_file_extension(model_info.path, ".pkl")
if kwargs["store_as_json_and_weights"]:
assert_have_file_extension(model_info.path, ".json")
if kwargs["save_format"] == "h5":
assert_have_file_extension(model_info.path, ".hdf5")
else:
if kwargs["save_format"] == "h5":
assert_have_file_extension(model_info.path, ".h5")
if not TF2:
session = bentoml.keras.get_session()
# Initialize variables in the graph/model
session.run(tf.global_variables_initializer())
with session.as_default():
loaded = bentoml.keras.load(tag)
predict_assert_equal(loaded)
else:
loaded = bentoml.keras.load(tag)
predict_assert_equal(loaded)
def test_pytorch_save_load(models):
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
tag = models(labels=labels, custom_objects={"func": custom_f})
bentomodel = bentoml.models.get(tag)
assert_have_file_extension(bentomodel.path, ".pt")
assert bentomodel.info.context.get("model_format") == "torch.save:v1"
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
pytorch_loaded: nn.Module = bentoml.pytorch.load(tag)
assert predict_df(pytorch_loaded, test_df) == 5.0
def test_lightgbm_save_load(metadata, save_proc):
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
bentomodel = save_proc(metadata,
labels=labels,
custom_objects={"func": custom_f})
assert bentomodel.info.metadata is not None
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
assert_have_file_extension(bentomodel.path, ".txt")
lgb_loaded = bentoml.lightgbm.load(bentomodel.tag, )
assert isinstance(lgb_loaded, lgb.basic.Booster)
assert lgb_loaded.predict(np.array([[0]])) == np.array([0.0])
def test_sklearn_save_load(metadata: t.Dict[str, t.Any]) -> None:
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
_, data = sklearn_model_data(clf=RandomForestClassifier)
tag = save_procedure(metadata, labels=labels, custom_objects={"func": custom_f})
bentomodel = bentoml.models.get(tag)
assert bentomodel.info.metadata is not None
assert_have_file_extension(bentomodel.path, ".pkl")
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
loaded = bentoml.sklearn.load(bentomodel.tag)
assert isinstance(loaded, RandomForestClassifier)
np.testing.assert_array_equal(loaded.predict(data), res_arr)
def test_pycaret_save_load(
get_pycaret_data, metadata, save_proc
): # noqa # pylint: disable
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
_, test_data = get_pycaret_data
bentomodel = save_proc(metadata, labels=labels, custom_objects={"func": custom_f})
assert bentomodel.info.metadata is not None
assert_have_file_extension(bentomodel.path, ".pkl")
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
pycaret_loaded = bentoml.pycaret.load(
bentomodel.tag,
)
assert isinstance(pycaret_loaded, sklearn.pipeline.Pipeline)
assert predict_model(pycaret_loaded, data=test_data)["Score"][0] == 0.7609
def test_onnxmlir_save_load(
compile_model,
tmpdir,
): # noqa
labels = {"stage": "dev"}
def custom_f(x: int) -> int:
return x + 1
model = os.path.join(tmpdir, "model.so")
tag = bentoml.onnxmlir.save("onnx_model_tests",
model,
labels=labels,
custom_objects={"func": custom_f})
bentomodel = bentoml.models.get(tag)
assert "compiled_path" in bentomodel.info.options
assert_have_file_extension(str(bentomodel.path), ".so")
for k in labels.keys():
assert labels[k] == bentomodel.info.labels[k]
assert bentomodel.custom_objects["func"](3) == custom_f(3)
session = bentoml.onnxmlir.load(tag)
assert predict_df(session, test_df)[0] == np.array([[15.0]])
def test_fastai_save_pack(tmpdir):
pack_models(tmpdir)
assert_have_file_extension(tmpdir, ".pkl")
loaded_fastai: "Learner" = FastAIModel.load(tmpdir)
assert predict_df(loaded_fastai, test_df) == 5.0