wehub-resource-sync
64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
import json
|
|
|
|
import numpy as np
|
|
from pmdarima import auto_arima, model_selection
|
|
from pmdarima.datasets import load_wineind
|
|
|
|
import mlflow
|
|
from mlflow.models import infer_signature
|
|
|
|
ARTIFACT_PATH = "model"
|
|
|
|
|
|
def calculate_cv_metrics(model, endog, metric, cv):
|
|
cv_metric = model_selection.cross_val_score(model, endog, cv=cv, scoring=metric, verbose=0)
|
|
return cv_metric[~np.isnan(cv_metric)].mean()
|
|
|
|
|
|
with mlflow.start_run():
|
|
data = load_wineind()
|
|
|
|
train, test = model_selection.train_test_split(data, train_size=150)
|
|
|
|
print("Training AutoARIMA model...")
|
|
arima = auto_arima(
|
|
train,
|
|
error_action="ignore",
|
|
trace=False,
|
|
suppress_warnings=True,
|
|
maxiter=5,
|
|
seasonal=True,
|
|
m=12,
|
|
)
|
|
|
|
print("Model trained. \nExtracting parameters...")
|
|
parameters = arima.get_params(deep=True)
|
|
|
|
metrics = {x: getattr(arima, x)() for x in ["aicc", "aic", "bic", "hqic", "oob"]}
|
|
|
|
# Cross validation backtesting
|
|
cross_validator = model_selection.RollingForecastCV(h=10, step=20, initial=60)
|
|
|
|
for x in ["smape", "mean_absolute_error", "mean_squared_error"]:
|
|
metrics[x] = calculate_cv_metrics(arima, data, x, cross_validator)
|
|
|
|
print(f"Metrics: \n{json.dumps(metrics, indent=2)}")
|
|
print(f"Parameters: \n{json.dumps(parameters, indent=2)}")
|
|
|
|
predictions = arima.predict(n_periods=30, return_conf_int=False)
|
|
signature = infer_signature(train, predictions)
|
|
|
|
model_info = mlflow.pmdarima.log_model(
|
|
pmdarima_model=arima, name=ARTIFACT_PATH, signature=signature
|
|
)
|
|
mlflow.log_params(parameters)
|
|
mlflow.log_metrics(metrics)
|
|
|
|
print(f"Model artifact logged to: {model_info.model_uri}")
|
|
|
|
loaded_model = mlflow.pmdarima.load_model(model_info.model_uri)
|
|
|
|
forecast = loaded_model.predict(30)
|
|
|
|
print(f"Forecast: \n{forecast}")
|