Files
FLAML/test/nlp/test_autohf_regression.py
Gleb Levitski 3de0dc667e Add ruff sort to pre-commit and sort imports in the library (#1259)
* lint

* bump ver

* bump ver

* fixed circular import

---------

Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
2024-03-12 21:28:57 +00:00

45 lines
1.1 KiB
Python

import os
import shutil
import sys
import pytest
from utils import get_automl_settings, get_toy_data_seqregression
@pytest.mark.skipif(sys.platform == "darwin", reason="do not run on mac os")
def test_regression():
try:
import ray
if not ray.is_initialized():
ray.init()
except ImportError:
return
from flaml import AutoML
X_train, y_train, X_val, y_val = get_toy_data_seqregression()
automl = AutoML()
automl_settings = get_automl_settings()
automl_settings["task"] = "seq-regression"
automl_settings["metric"] = "pearsonr"
automl_settings["starting_points"] = {"transformer": {"num_train_epochs": 1}}
automl_settings["use_ray"] = {"local_dir": "data/output/"}
ray.shutdown()
ray.init()
automl.fit(X_train=X_train, y_train=y_train, X_val=X_val, y_val=y_val, **automl_settings)
automl.predict(X_val)
if os.path.exists("test/data/output/"):
try:
shutil.rmtree("test/data/output/")
except PermissionError:
print("PermissionError when deleting test/data/output/")
if __name__ == "__main__":
test_regression()