def test_predict_feature_columns():
preprocessor = DummyPreprocessor()
predictor = XGBoostPredictor(model=model, preprocessor=preprocessor)
data_batch = np.array([[1, 2, 7], [3, 4, 8], [5, 6, 9]])
predictions = predictor.predict(data_batch, feature_columns=[0, 1])
assert len(predictions) == 3
assert hasattr(predictor.preprocessor, "_batch_transformed")
def test_predict():
preprocessor = DummyPreprocessor()
predictor = XGBoostPredictor(model=model, preprocessor=preprocessor)
data_batch = np.array([[1, 2], [3, 4], [5, 6]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert hasattr(predictor.preprocessor, "_batch_transformed")
def test_predict(batch_type):
preprocessor = DummyPreprocessor()
predictor = XGBoostPredictor(model=model, preprocessor=preprocessor)
raw_batch = pd.DataFrame([[1, 2], [3, 4], [5, 6]])
data_batch = convert_pandas_to_batch_type(raw_batch, type=TYPE_TO_ENUM[batch_type])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert hasattr(predictor.get_preprocessor(), "_batch_transformed")
def test_predict_feature_columns_pandas():
pandas_data = pd.DataFrame(dummy_data, columns=["A", "B"])
pandas_target = pd.Series(dummy_target)
pandas_model = (
xgb.XGBClassifier(n_estimators=10).fit(pandas_data, pandas_target).get_booster()
)
preprocessor = DummyPreprocessor()
predictor = XGBoostPredictor(model=pandas_model, preprocessor=preprocessor)
data_batch = pd.DataFrame(
np.array([[1, 2, 7], [3, 4, 8], [5, 6, 9]]), columns=["A", "B", "C"]
)
predictions = predictor.predict(data_batch, feature_columns=["A", "B"])
assert len(predictions) == 3
assert hasattr(predictor.preprocessor, "_batch_transformed")
def test_predict_no_preprocessor_no_training():
with tempfile.TemporaryDirectory() as tmpdir:
checkpoint = to_air_checkpoint(tmpdir, booster=model)
predictor = XGBoostPredictor.from_checkpoint(checkpoint)
data_batch = np.array([[1, 2], [3, 4], [5, 6]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
def test_init():
preprocessor = DummyPreprocessor()
preprocessor.attr = 1
predictor = XGBoostPredictor(model=model, preprocessor=preprocessor)
with tempfile.TemporaryDirectory() as tmpdir:
# This somewhat convoluted procedure is the same as in the
# Trainers. The reason for saving model to disk instead
# of directly to the dict as bytes is due to all callbacks
# following save to disk logic. GBDT models are small
# enough that IO should not be an issue.
model.save_model(os.path.join(tmpdir, MODEL_KEY))
save_preprocessor_to_dir(preprocessor, tmpdir)
checkpoint = Checkpoint.from_directory(tmpdir)
checkpoint_predictor = XGBoostPredictor.from_checkpoint(checkpoint)
assert get_num_trees(checkpoint_predictor.model) == get_num_trees(predictor.model)
assert checkpoint_predictor.preprocessor.attr == predictor.preprocessor.attr