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