mirror of
https://github.com/microsoft/FLAML.git
synced 2026-02-09 02:09:16 +08:00
3de0dc667e
* lint * bump ver * bump ver * fixed circular import --------- Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
32 lines
969 B
Python
32 lines
969 B
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
from sklearn.datasets import load_iris
|
|
|
|
from flaml import AutoML
|
|
|
|
|
|
@pytest.mark.conda
|
|
def test_package_minimum():
|
|
# Initialize an AutoML instance
|
|
automl = AutoML()
|
|
# Specify automl goal and constraint
|
|
automl_settings = {
|
|
"time_budget": 10, # in seconds
|
|
"metric": "accuracy",
|
|
"task": "classification",
|
|
"log_file_name": "iris.log",
|
|
}
|
|
X_train, y_train = load_iris(return_X_y=True)
|
|
# Train with labeled input data
|
|
automl.fit(X_train=X_train, y_train=y_train, **automl_settings)
|
|
# Check that `best_config` is created, the log was created and best model is accessible
|
|
assert hasattr(automl, "best_config")
|
|
assert Path("iris.log").exists()
|
|
assert automl.model is not None
|
|
print(automl.model)
|
|
# Predict and check that the prediction shape is as expected
|
|
preds = automl.predict_proba(X_train)
|
|
assert preds.shape == (150, 3)
|
|
print(preds)
|