1  import shap
 2  import xgboost
 3  from sklearn.model_selection import train_test_split
 4  
 5  import mlflow
 6  from mlflow.models import infer_signature
 7  
 8  # Load the UCI Adult Dataset
 9  X, y = shap.datasets.adult()
10  
11  # Split the data into training and test sets
12  X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
13  
14  # Fit an XGBoost binary classifier on the training data split
15  model = xgboost.XGBClassifier().fit(X_train, y_train)
16  
17  # Infer model signature
18  predictions = model.predict(X_train)
19  signature = infer_signature(X_train, predictions)
20  
21  # Build the Evaluation Dataset from the test set
22  eval_data = X_test
23  eval_data["label"] = y_test
24  
25  with mlflow.start_run() as run:
26      # Log the XGBoost binary classifier model to MLflow
27      model_info = mlflow.sklearn.log_model(model, name="model", signature=signature)
28  
29      # Evaluate the logged model
30      result = mlflow.evaluate(
31          model_info.model_uri,
32          eval_data,
33          targets="label",
34          model_type="classifier",
35          evaluators=["default"],
36      )
37  
38  print(f"metrics:\n{result.metrics}")
39  print(f"artifacts:\n{result.artifacts}")