{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Convert a XGBoost model to CoreML\n",
    "\n",
    "In this example, we will train an XGBoost regression model and convert it to CoreML.  For more details on all the different ways to convert these models, see the [coremltools XGBoost documentation](https://coremltools.readme.io/reference/convertersxgboostconvert)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an XGBoost Model\n",
    "\n",
    "For this example, we will follow a [regression example](https://github.com/dmlc/xgboost/tree/master/demo/CLI/regression) in the XGBoost repository.  For this example, we are predicting the performance of a computer system based on some features.  This data comes from the [UCI Machine Learning Repository Compute Hardware Data Set](https://archive.ics.uci.edu/ml/datasets/Computer+Hardware)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgboost as xgb\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we load the train and test data (already split):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtrain = xgb.DMatrix('machine.txt.train')\n",
    "dtest = xgb.DMatrix('machine.txt.test')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train the model using parameters from the example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param = {\n",
    "    'objective': 'reg:squarederror',\n",
    "    'eta': '1.0',\n",
    "    'gamma': 1.0,\n",
    "    'min_child_weight': 1,\n",
    "    'max_depth': 3,\n",
    "}\n",
    "bst_model = xgb.train(param, dtrain, num_boost_round=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And plot the tree:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xgb.plot_tree(bst_model, 'machine.featmap.txt')\n",
    "fig = plt.gcf()\n",
    "fig.set_size_inches(20, 40)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we can test the accuracy of the model on the training data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = bst_model.predict(dtrain)\n",
    "actual = dtrain.get_label()\n",
    "error = prediction - actual\n",
    "print('mean error:', error.mean(), 'stdev error:', error.std())\n",
    "plt.hist(error)\n",
    "plt.xlabel('prediction - actual')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to CoreML Model\n",
    "\n",
    "Converting an XGBoost model to CoreML format is much simpler than PyTorch or TensorFlow.  However, if we want to use the proper feature names for model inputs, we need to load them and pass them to the `convert` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_names = []\n",
    "with open('machine.featmap.txt') as f:\n",
    "    for line in f.readlines():\n",
    "        feature_name = line.split()[1]\n",
    "        feature_names.append(feature_name)\n",
    "feature_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use the desired feature names and the name of the model target (the output) during model conversion:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import coremltools as ct\n",
    "\n",
    "cml_model = ct.converters.xgboost.convert(bst_model, feature_names=feature_names, target='perf', mode='regressor')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the feature names in the metadata describing this model by looking at the string representation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cml_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the `vendor` categorical input is represented in the model as a one-hot-encoded value."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can write the model to disk in the CoreML format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cml_model.save('machine.mlmodel')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the CoreML Model\n",
    "\n",
    "As with other CoreML model types, if we are on a macOS system, we can use the predict method to run the model.  We can pass our input in the form of a dictionary with the feature values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example = {\n",
    "    'MYCT': 125,\n",
    "    'MMIN': 256,\n",
    "    'MMAX': 6000,\n",
    "    'CACH': 256,\n",
    "    'CHMIN': 16,\n",
    "    'CHMAX': 128,\n",
    "}\n",
    "# Set the one-hot-encoded vendor feature\n",
    "for feature_name in feature_names:\n",
    "    if feature_name == 'vendor:ibm':\n",
    "        example[feature_name] = 1\n",
    "    elif feature_name.startswith('vendor'):\n",
    "        example[feature_name] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "IS_MACOS = sys.platform == 'darwin'\n",
    "\n",
    "if IS_MACOS:\n",
    "    loaded_model = ct.models.MLModel('machine.mlmodel')\n",
    "    prediction = loaded_model.predict(example)\n",
    "    print('prediction:', prediction)\n",
    "else:\n",
    "    prediction = 'Skipping prediction on non-macOS system'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}