wehub-resource-sync
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import shap
|
|
import xgboost
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
import mlflow
|
|
from mlflow.models import infer_signature
|
|
|
|
# Load the UCI Adult Dataset
|
|
X, y = shap.datasets.adult()
|
|
|
|
# Split the data into training and test sets
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
|
|
|
|
# Fit an XGBoost binary classifier on the training data split
|
|
model = xgboost.XGBClassifier().fit(X_train, y_train)
|
|
|
|
# Infer model signature
|
|
predictions = model.predict(X_train)
|
|
signature = infer_signature(X_train, predictions)
|
|
|
|
# Build the Evaluation Dataset from the test set
|
|
eval_data = X_test
|
|
eval_data["label"] = y_test
|
|
|
|
with mlflow.start_run() as run:
|
|
# Log the XGBoost binary classifier model to MLflow
|
|
model_info = mlflow.sklearn.log_model(
|
|
model, name="model", signature=signature, serialization_format="cloudpickle"
|
|
)
|
|
|
|
# Evaluate the logged model
|
|
result = mlflow.evaluate(
|
|
model_info.model_uri,
|
|
eval_data,
|
|
targets="label",
|
|
model_type="classifier",
|
|
evaluators=["default"],
|
|
)
|
|
|
|
print(f"metrics:\n{result.metrics}")
|
|
print(f"artifacts:\n{result.artifacts}")
|