sdk/python/responsible-ai/mlflow-deployment-with-explanations/mlflow-deployment-with-explanations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MLFlow Deployment with Explanations\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Learn how to implement scoring-time explanations for [MLflow](https://www.mlflow.org/) model. This tutorial produces (1) a new model whose `predict()` method returns both predictions and explanations and (2) a custom deployment of that model to an [online endpoint](https://docs.microsoft.com/azure/machine-learning/concept-endpoints) for that model.\n",
    "\n",
    "The \"explanations\" for this model come in the form of local feature importance values. Local feature importance measures the contribution of features for a specific prediction. This tutorial leverages Microsoft's [Responsible AI Toolbox](https://github.com/microsoft/responsible-ai-toolbox) to generate these values, which uses a Mimic explainer, also known as a global surrogate model. You can learn more in the Interpret ML Book's [chapter on global surrogates](https://christophm.github.io/interpretable-ml-book/global.html).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Example of deployment with explanations](assets/DeploymentDiabetesExample.png \"Example of deployment with explanations\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Requirements\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- An Azure account with an active subscription - [Create an account for free](https://azure.microsoft.com/free/?WT.mc_id=A261C142F)\n",
    "- An Azure ML workspace with computer cluster - [Learn about workspaces](https://docs.microsoft.com/en-us/azure/machine-learning/concept-workspace)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running this notebook relies on packages in the `requirements.txt` file. Install them or run the command below:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbac0832",
   "metadata": {},
   "source": [
    "We also need to give the name of the compute cluster we want to use in AzureML. Later in this notebook, we will create it if it does not already exist:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca2ed890",
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_name = \"rai-cluster\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7402748c",
   "metadata": {},
   "source": [
    "Finally, we need to specify a version for the data and components we will create while running this notebook. This should be unique for the workspace, but the specific value doesn't matter:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132f8caa",
   "metadata": {},
   "outputs": [],
   "source": [
    "diabetes_regression_example_version_string = \"7\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21914b99",
   "metadata": {},
   "source": [
    "## Configure workspace details and get a handle to the workspace\n",
    "\n",
    "To connect to a workspace, we need identifier parameters - a subscription, resource group and workspace name. We will use these details in the MLClient from `azure.ai.ml` to get a handle to the required Azure Machine Learning workspace.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afd018cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enter details of your AML workspace\n",
    "subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "resource_group = \"<RESOURCE_GROUP>\"\n",
    "workspace = \"<AML_WORKSPACE_NAME>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d70782e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Handle to the workspace\n",
    "import mlflow\n",
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import DefaultAzureCredential\n",
    "\n",
    "\n",
    "ml_client = MLClient(\n",
    "    DefaultAzureCredential(), subscription_id, resource_group, workspace\n",
    ")\n",
    "print(f\"ML Client: {ml_client}\")\n",
    "\n",
    "azureml_tracking_uri = ml_client.workspaces.get(\n",
    "    ml_client.workspace_name\n",
    ").mlflow_tracking_uri\n",
    "mlflow.set_tracking_uri(azureml_tracking_uri)\n",
    "print(f\"Tracking URI: {azureml_tracking_uri}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff51070e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get handle to azureml registry for the RAI built in components\n",
    "registry_name = \"azureml\"\n",
    "ml_client_registry = MLClient(\n",
    "    credential=DefaultAzureCredential(),\n",
    "    subscription_id=subscription_id,\n",
    "    resource_group_name=resource_group,\n",
    "    registry_name=registry_name,\n",
    ")\n",
    "print(ml_client_registry)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce352c01",
   "metadata": {},
   "source": [
    "## Accessing the data\n",
    "\n",
    "First, we need to obtain the dataset and upload it to our AzureML workspace:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9a93781",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_path = \"data-diabetes-regression/train/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82c9970e",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_path = \"data-diabetes-regression/test/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dd7b4a8",
   "metadata": {},
   "source": [
    "Load some data for a quick view:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69655f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import mltable\n",
    "\n",
    "tbl = mltable.load(train_data_path)\n",
    "train_df: pd.DataFrame = tbl.to_pandas_dataframe()\n",
    "\n",
    "# test dataset should have less than 5000 rows\n",
    "test_df = mltable.load(test_data_path).to_pandas_dataframe()\n",
    "assert len(test_df.index) <= 5000\n",
    "\n",
    "display(train_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b29574f0",
   "metadata": {},
   "source": [
    "We are going to create two Datasets in AzureML, one for the train and one for the test datasets.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c47a6faa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Data\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "\n",
    "input_train_data = \"diabetes_regression_train_mltable\"\n",
    "input_test_data = \"diabetes_regression_test_mltable\"\n",
    "\n",
    "try:\n",
    "    # Try getting data already registered in workspace\n",
    "    train_data = ml_client.data.get(\n",
    "        name=input_train_data, version=diabetes_regression_example_version_string\n",
    "    )\n",
    "    test_data = ml_client.data.get(\n",
    "        name=input_test_data, version=diabetes_regression_example_version_string\n",
    "    )\n",
    "except Exception as e:\n",
    "    train_data = Data(\n",
    "        path=train_data_path,\n",
    "        type=AssetTypes.MLTABLE,\n",
    "        description=\"Diabetes regression example training data\",\n",
    "        name=input_train_data,\n",
    "        version=diabetes_regression_example_version_string,\n",
    "    )\n",
    "    ml_client.data.create_or_update(train_data)\n",
    "\n",
    "    test_data = Data(\n",
    "        path=test_data_path,\n",
    "        type=AssetTypes.MLTABLE,\n",
    "        description=\"Diabetes regression example test data\",\n",
    "        name=input_test_data,\n",
    "        version=diabetes_regression_example_version_string,\n",
    "    )\n",
    "    ml_client.data.create_or_update(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "088aff35",
   "metadata": {},
   "source": [
    "## A model training pipeline\n",
    "\n",
    "To simplify the model creation process, we're going to use a pipeline. This will have two stages:\n",
    "\n",
    "1. The actual training component\n",
    "1. A model registration component\n",
    "\n",
    "We have to register the model in AzureML in order for our RAI insights components to use it.\n",
    "\n",
    "### The Training Component\n",
    "\n",
    "The training component is for this particular model. In this case, we are going to train a `RandomForestRegressor` on the input data and save it using MLFlow. We need command line arguments to specify the location of the input data, the location where MLFlow should write the output model, and the name of the target column in the dataset.\n",
    "\n",
    "We start by creating a directory to hold the component source:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "642bf7b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.makedirs(\"component_src\", exist_ok=True)\n",
    "os.makedirs(\"register_model_src\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d5de0e9",
   "metadata": {},
   "source": [
    "Next, put our training script into the directory:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f80deff7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile component_src/diabetes_regression_training_script.py\n",
    "\n",
    "import argparse\n",
    "import os\n",
    "import shutil\n",
    "import tempfile\n",
    "\n",
    "\n",
    "from azureml.core import Run\n",
    "\n",
    "import mlflow\n",
    "import mlflow.sklearn\n",
    "\n",
    "import mltable\n",
    "\n",
    "import pandas as pd\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "\n",
    "def parse_args():\n",
    "    # setup arg parser\n",
    "    parser = argparse.ArgumentParser()\n",
    "\n",
    "    # add arguments\n",
    "    parser.add_argument(\"--training_data\", type=str, help=\"Path to training data\")\n",
    "    parser.add_argument(\"--target_column_name\", type=str, help=\"Name of target column\")\n",
    "    parser.add_argument(\"--model_output\", type=str, help=\"Path of output model\")\n",
    "\n",
    "    # parse args\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    # return args\n",
    "    return args\n",
    "\n",
    "\n",
    "def main(args):\n",
    "    current_experiment = Run.get_context().experiment\n",
    "    tracking_uri = current_experiment.workspace.get_mlflow_tracking_uri()\n",
    "    print(\"tracking_uri: {0}\".format(tracking_uri))\n",
    "    mlflow.set_tracking_uri(tracking_uri)\n",
    "    mlflow.set_experiment(current_experiment.name)\n",
    "\n",
    "    # Read in data\n",
    "    print(\"Reading data\")\n",
    "    tbl = mltable.load(args.training_data)\n",
    "    all_data = tbl.to_pandas_dataframe()\n",
    "\n",
    "    print(\"Extracting X_train, y_train\")\n",
    "    print(\"all_data cols: {0}\".format(all_data.columns))\n",
    "    y_train = all_data[args.target_column_name]\n",
    "    X_train = all_data.drop(labels=args.target_column_name, axis=\"columns\")\n",
    "    print(\"X_train cols: {0}\".format(X_train.columns))\n",
    "\n",
    "    print(\"Training model\")\n",
    "    # The estimator can be changed to suit\n",
    "    model = RandomForestRegressor()\n",
    "    model.fit(X_train, y_train)\n",
    "\n",
    "    # Saving model with mlflow - leave this section unchanged\n",
    "    with tempfile.TemporaryDirectory() as td:\n",
    "        print(\"Saving model with MLFlow to temporary directory\")\n",
    "        tmp_output_dir = os.path.join(td, \"my_model_dir\")\n",
    "        mlflow.sklearn.save_model(sk_model=model, path=tmp_output_dir)\n",
    "\n",
    "        print(\"Copying MLFlow model to output path\")\n",
    "        for file_name in os.listdir(tmp_output_dir):\n",
    "            print(\"  Copying: \", file_name)\n",
    "            # As of Python 3.8, copytree will acquire dirs_exist_ok as\n",
    "            # an option, removing the need for listdir\n",
    "            shutil.copy2(src=os.path.join(tmp_output_dir, file_name), dst=os.path.join(args.model_output, file_name))\n",
    "\n",
    "\n",
    "# run script\n",
    "if __name__ == \"__main__\":\n",
    "    # add space in logs\n",
    "    print(\"*\" * 60)\n",
    "    print(\"\\n\\n\")\n",
    "\n",
    "    # parse args\n",
    "    args = parse_args()\n",
    "\n",
    "    # run main function\n",
    "    main(args)\n",
    "\n",
    "    # add space in logs\n",
    "    print(\"*\" * 60)\n",
    "    print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c9ee087",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile register_model_src/register.py\n",
    "\n",
    "# ---------------------------------------------------------\n",
    "# Copyright (c) Microsoft Corporation. All rights reserved.\n",
    "# ---------------------------------------------------------\n",
    "\n",
    "import argparse\n",
    "import json\n",
    "import os\n",
    "import time\n",
    "\n",
    "\n",
    "from azureml.core import Run\n",
    "\n",
    "import mlflow\n",
    "import mlflow.sklearn\n",
    "\n",
    "# Based on example:\n",
    "# https://docs.microsoft.com/en-us/azure/machine-learning/how-to-train-cli\n",
    "# which references\n",
    "# https://github.com/Azure/azureml-examples/tree/main/cli/jobs/train/lightgbm/iris\n",
    "\n",
    "\n",
    "def parse_args():\n",
    "    # setup arg parser\n",
    "    parser = argparse.ArgumentParser()\n",
    "\n",
    "    # add arguments\n",
    "    parser.add_argument(\"--model_input_path\", type=str, help=\"Path to input model\")\n",
    "    parser.add_argument(\n",
    "        \"--model_info_output_path\", type=str, help=\"Path to write model info JSON\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--model_base_name\", type=str, help=\"Name of the registered model\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--model_name_suffix\", type=int, help=\"Set negative to use epoch_secs\"\n",
    "    )\n",
    "\n",
    "    # parse args\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    # return args\n",
    "    return args\n",
    "\n",
    "\n",
    "def main(args):\n",
    "    current_experiment = Run.get_context().experiment\n",
    "    tracking_uri = current_experiment.workspace.get_mlflow_tracking_uri()\n",
    "    print(\"tracking_uri: {0}\".format(tracking_uri))\n",
    "    mlflow.set_tracking_uri(tracking_uri)\n",
    "    mlflow.set_experiment(current_experiment.name)\n",
    "\n",
    "    print(\"Loading model\")\n",
    "    mlflow_model = mlflow.sklearn.load_model(args.model_input_path)\n",
    "\n",
    "    if args.model_name_suffix < 0:\n",
    "        suffix = int(time.time())\n",
    "    else:\n",
    "        suffix = args.model_name_suffix\n",
    "    registered_name = \"{0}_{1}\".format(args.model_base_name, suffix)\n",
    "    print(f\"Registering model as {registered_name}\")\n",
    "\n",
    "    print(\"Registering via MLFlow\")\n",
    "    mlflow.sklearn.log_model(\n",
    "        sk_model=mlflow_model,\n",
    "        registered_model_name=registered_name,\n",
    "        artifact_path=registered_name,\n",
    "    )\n",
    "\n",
    "    print(\"Writing JSON\")\n",
    "    dict = {\"id\": \"{0}:1\".format(registered_name)}\n",
    "    output_path = os.path.join(args.model_info_output_path, \"model_info.json\")\n",
    "    with open(output_path, \"w\") as of:\n",
    "        json.dump(dict, fp=of)\n",
    "\n",
    "\n",
    "# run script\n",
    "if __name__ == \"__main__\":\n",
    "    # add space in logs\n",
    "    print(\"*\" * 60)\n",
    "    print(\"\\n\\n\")\n",
    "\n",
    "    # parse args\n",
    "    args = parse_args()\n",
    "\n",
    "    # run main function\n",
    "    main(args)\n",
    "\n",
    "    # add space in logs\n",
    "    print(\"*\" * 60)\n",
    "    print(\"\\n\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54f1e905",
   "metadata": {},
   "source": [
    "Now that the training script is saved on our local drive, we create a YAML file to describe it as a component to AzureML. This involves defining the inputs and outputs, specifing the AzureML environment which can run the script, and telling AzureML how to invoke the training script:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8b9fab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import load_component\n",
    "\n",
    "yaml_contents = (\n",
    "    f\"\"\"\n",
    "$schema: http://azureml/sdk-2-0/CommandComponent.json\n",
    "name: rai_diabetes_regression_training_component\n",
    "display_name: Diabetes regression training component for RAI example\n",
    "version: {diabetes_regression_example_version_string}\n",
    "type: command\n",
    "inputs:\n",
    "  training_data:\n",
    "    type: path\n",
    "  target_column_name:\n",
    "    type: string\n",
    "outputs:\n",
    "  model_output:\n",
    "    type: path\n",
    "code: ./component_src/\n",
    "environment: azureml://registries/azureml/environments/responsibleai-ubuntu20.04-py38-cpu/versions/49\n",
    "\"\"\"\n",
    "    + r\"\"\"\n",
    "command: >-\n",
    "  python diabetes_regression_training_script.py\n",
    "  --training_data ${{{{inputs.training_data}}}}\n",
    "  --target_column_name ${{{{inputs.target_column_name}}}}\n",
    "  --model_output ${{{{outputs.model_output}}}}\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "yaml_filename = \"DiabetesRegressionTrainingComponent.yaml\"\n",
    "\n",
    "with open(yaml_filename, \"w\") as f:\n",
    "    f.write(yaml_contents.format(yaml_contents))\n",
    "\n",
    "train_model_component = load_component(source=yaml_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0581633a",
   "metadata": {},
   "outputs": [],
   "source": [
    "yaml_contents = f\"\"\"\n",
    "$schema: http://azureml/sdk-2-0/CommandComponent.json\n",
    "name: register_model\n",
    "display_name: Register Model\n",
    "version: {diabetes_regression_example_version_string}\n",
    "type: command\n",
    "is_deterministic: False\n",
    "inputs:\n",
    "  model_input_path:\n",
    "    type: path\n",
    "  model_base_name:\n",
    "    type: string\n",
    "  model_name_suffix: # Set negative to use epoch_secs\n",
    "    type: integer\n",
    "    default: -1\n",
    "outputs:\n",
    "  model_info_output_path:\n",
    "    type: path\n",
    "code: ./register_model_src/\n",
    "environment: azureml://registries/azureml/environments/responsibleai-ubuntu20.04-py38-cpu/versions/49\n",
    "command: >-\n",
    "  python register.py\n",
    "  --model_input_path ${{{{inputs.model_input_path}}}}\n",
    "  --model_base_name ${{{{inputs.model_base_name}}}}\n",
    "  --model_name_suffix ${{{{inputs.model_name_suffix}}}}\n",
    "  --model_info_output_path ${{{{outputs.model_info_output_path}}}}\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "yaml_filename = \"register.yaml\"\n",
    "\n",
    "with open(yaml_filename, \"w\") as f:\n",
    "    f.write(yaml_contents)\n",
    "\n",
    "register_component = load_component(source=yaml_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "693f7706",
   "metadata": {},
   "source": [
    "### Running a training pipeline\n",
    "\n",
    "The component to register the model is part of the suite of RAI components, so we do not have to define it here. As such, we are now ready to run the training pipeline itself.\n",
    "\n",
    "We start by ensuring that the compute cluster named above exists:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad3b43ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import AmlCompute\n",
    "\n",
    "all_compute_names = [x.name for x in ml_client.compute.list()]\n",
    "\n",
    "if compute_name in all_compute_names:\n",
    "    print(f\"Found existing compute: {compute_name}\")\n",
    "else:\n",
    "    my_compute = AmlCompute(\n",
    "        name=compute_name,\n",
    "        size=\"Standard_D2_v2\",\n",
    "        min_instances=0,\n",
    "        max_instances=4,\n",
    "        idle_time_before_scale_down=3600,\n",
    "    )\n",
    "    ml_client.compute.begin_create_or_update(my_compute).result()\n",
    "    print(\"Initiated compute creation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7dfa1d7",
   "metadata": {},
   "source": [
    "We continue by defining the name under which we want to register the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d12c37f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "model_name_suffix = int(time.time())\n",
    "model_name = \"diabetes_decision_model\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15cfc17e",
   "metadata": {},
   "source": [
    "Next, we define the pipeline using objects from the AzureML SDKv2. As mentioned above, there are two component jobs: one to train the model, and one to register it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de827a6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import dsl, Input\n",
    "\n",
    "target_column = \"y\"\n",
    "\n",
    "diabetes_train_pq = Input(\n",
    "    type=\"mltable\",\n",
    "    path=f\"azureml:{input_train_data}:{diabetes_regression_example_version_string}\",\n",
    "    mode=\"download\",\n",
    ")\n",
    "diabetes_test_pq = Input(\n",
    "    type=\"mltable\",\n",
    "    path=f\"azureml:{input_test_data}:{diabetes_regression_example_version_string}\",\n",
    "    mode=\"download\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    compute=compute_name,\n",
    "    description=\"Register Model for Diabetes Decision Making example\",\n",
    "    experiment_name=f\"Diabetes_Decision_Example_Model_Training_{model_name_suffix}\",\n",
    ")\n",
    "def my_training_pipeline(target_column_name, training_data):\n",
    "    trained_model = train_model_component(\n",
    "        target_column_name=target_column_name, training_data=training_data\n",
    "    )\n",
    "    trained_model.set_limits(timeout=3600)\n",
    "\n",
    "    _ = register_component(\n",
    "        model_input_path=trained_model.outputs.model_output,\n",
    "        model_base_name=model_name,\n",
    "        model_name_suffix=model_name_suffix,\n",
    "    )\n",
    "\n",
    "    return {}\n",
    "\n",
    "\n",
    "model_registration_pipeline_job = my_training_pipeline(target_column, diabetes_train_pq)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8acc2a29",
   "metadata": {},
   "source": [
    "With the pipeline definition created, we can submit it to AzureML. We define a helper function to do the submission, which waits for the submitted job to complete:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5c0f883",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import PipelineJob\n",
    "from IPython.core.display import HTML\n",
    "from IPython.display import display\n",
    "\n",
    "\n",
    "def submit_and_wait(ml_client, pipeline_job) -> PipelineJob:\n",
    "    created_job = ml_client.jobs.create_or_update(pipeline_job)\n",
    "    assert created_job is not None\n",
    "\n",
    "    print(\"Pipeline job can be accessed in the following URL:\")\n",
    "    display(HTML('<a href=\"{0}\">{0}</a>'.format(created_job.studio_url)))\n",
    "\n",
    "    while created_job.status not in [\n",
    "        \"Completed\",\n",
    "        \"Failed\",\n",
    "        \"Canceled\",\n",
    "        \"NotResponding\",\n",
    "    ]:\n",
    "        time.sleep(30)\n",
    "        created_job = ml_client.jobs.get(created_job.name)\n",
    "        print(\"Latest status : {0}\".format(created_job.status))\n",
    "    assert created_job.status == \"Completed\"\n",
    "    return created_job\n",
    "\n",
    "\n",
    "# This is the actual submission\n",
    "training_job = submit_and_wait(ml_client, model_registration_pipeline_job)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8575b364",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = f\"{model_name}_{model_name_suffix}\"\n",
    "azureml_model_id = f\"azureml:{model_id}:1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Wrapper Model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial uses a wrapper model to provide explanations along with predictions of the original model. The code for this wrapper can be reviewed in `explanation_wrapper.py`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model information **(user input required)**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `wrapper_model_name` will be the name of your model with explanations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Name for the model with explanations\n",
    "wrapper_model_name = f\"diabetes_model_with_explanations_{model_name_suffix}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_model_name = f\"{model_name}_{model_name_suffix}\"\n",
    "model_uri = f\"models:/{full_model_name}/1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load baseline data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data used to create the explanation wrapper must match the data used to train the original model. Be sure to drop any columns that were dropped during training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import mltable\n",
    "\n",
    "\n",
    "data_asset = ml_client.data.get(\n",
    "    name=input_train_data, version=diabetes_regression_example_version_string\n",
    ")\n",
    "baseline_df = mltable.load(train_data_path).to_pandas_dataframe()\n",
    "# Drop any columns that were dropped when training the model. Uncomment and fill in first parameter\n",
    "# baseline_df = baseline_df.drop([], axis=\"columns\")\n",
    "\n",
    "print(\"Baseline Data (first 5 rows):\")\n",
    "baseline_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create instance of ExplanationWrapper\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from explanation_wrapper import ExplanationWrapper\n",
    "\n",
    "task_type = \"regression\"\n",
    "categorical_features = []\n",
    "explanation_wrapper = ExplanationWrapper(\n",
    "    model_uri=model_uri,\n",
    "    baseline_df=baseline_df,\n",
    "    target_column=target_column,\n",
    "    task_type=task_type,\n",
    "    categorical_features=categorical_features,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register Wrapper Model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download original model **(user input required)**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the original model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local folder\n",
    "import os\n",
    "\n",
    "local_path = \"./artifacts/original_model\"\n",
    "if not os.path.exists(local_path):\n",
    "    os.makedirs(local_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "from mlflow.tracking.client import MlflowClient\n",
    "\n",
    "# Initialize MLFlow client\n",
    "mlflow_client = MlflowClient()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download run's artifacts/outputs using the latest version of the model\n",
    "mlflow.artifacts.download_artifacts(artifact_uri=model_uri, dst_path=local_path)\n",
    "print(\"Artifacts downloaded in: {}\".format(local_path))\n",
    "print(\"Artifacts: {}\".format(os.listdir(local_path)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow_model_dir = os.path.join(local_path, full_model_name)\n",
    "model_file_path = f\"./artifacts/original_model/{full_model_name}\"\n",
    "# Show the contents of the MLFlow model folder\n",
    "os.listdir(mlflow_model_dir)\n",
    "\n",
    "# You should see a list of files such as the following:\n",
    "# ['artifacts', 'conda.yaml', 'MLmodel', 'python_env.yaml', 'python_model.pkl', 'requirements.txt']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get original model signature\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If there is one, load the model signature from Azure ML.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow.models import Model\n",
    "\n",
    "mlflow_model = Model.load(model_uri)\n",
    "model_signature = mlflow_model.signature\n",
    "print(f\"Model Signature: {model_signature}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Register wrapper model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save and then register the model to your workspace. If you are re-running this step, you may need to delete the folder containing the last version of the saved wrapper model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "artifacts = {\"model\": model_file_path, \"RAI insights\": \"./artifacts/RAI_Insights\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.pyfunc.save_model(\n",
    "    path=f\"artifacts/{wrapper_model_name}\",\n",
    "    code_path=[\"./explanation_wrapper.py\"],\n",
    "    conda_env=\"./env.yml\",\n",
    "    python_model=explanation_wrapper,\n",
    "    artifacts=artifacts,\n",
    "    signature=model_signature,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow.register_model(f\"file://artifacts/{wrapper_model_name}\", wrapper_model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Create the scoring Endpoint\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Endpoint Details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = \"diabetes-decision-endpoint\"\n",
    "\n",
    "print(f\"Endpoint name: {endpoint_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow.deployments import get_deploy_client\n",
    "\n",
    "deployment_client = get_deploy_client(azureml_tracking_uri)\n",
    "endpoint = deployment_client.create_endpoint(endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Deploy Wrapper Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get wrapper model ID\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that the model name and version are correct for the newly registered wrapper model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "version_list = list(ml_client.models.list(wrapper_model_name))\n",
    "wrapper_model_version = version_list[0].version\n",
    "wrapper_model = ml_client.models.get(wrapper_model_name, wrapper_model_version)\n",
    "print(f\"Using model name: {wrapper_model_name}, version: {wrapper_model_version}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create and deploy endpoint\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This step may take a while.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import (\n",
    "    OnlineRequestSettings,\n",
    "    ManagedOnlineDeployment,\n",
    "    ProbeSettings,\n",
    ")\n",
    "\n",
    "# Define the deployment\n",
    "deployment_name = \"mlflow-deploy-with-explanations\"\n",
    "deployment = ManagedOnlineDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=endpoint_name,\n",
    "    model=wrapper_model.id,\n",
    "    instance_count=1,\n",
    "    request_settings=OnlineRequestSettings(request_timeout_ms=90000),\n",
    "    liveness_probe=ProbeSettings(\n",
    "        failure_threshold=30,\n",
    "        success_threshold=1,\n",
    "        period=100,\n",
    "        initial_delay=500,\n",
    "    ),\n",
    "    readiness_probe=ProbeSettings(\n",
    "        failure_threshold=30,\n",
    "        success_threshold=1,\n",
    "        period=100,\n",
    "        initial_delay=500,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Trigger the deployment creation\n",
    "try:\n",
    "    ml_client.begin_create_or_update(deployment).wait()\n",
    "    print(\"\\n---Deployment created successfully---\\n\")\n",
    "except Exception as err:\n",
    "    raise RuntimeError(\n",
    "        f\"Deployment creation failed. Detailed Response:\\n{err}\"\n",
    "    ) from err"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Assign all traffic to the deployment\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the traffic configuration:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "traffic_config = {\"traffic\": {deployment_name: 100}}\n",
    "traffic_config_path = \"artifacts/traffic_config.json\"\n",
    "with open(traffic_config_path, \"w\") as outfile:\n",
    "    outfile.write(json.dumps(traffic_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update the configuration:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow.deployments import get_deploy_client\n",
    "\n",
    "deployment_client = get_deploy_client(mlflow.get_tracking_uri())\n",
    "deployment_client.update_endpoint(\n",
    "    endpoint=endpoint_name,\n",
    "    config={\"endpoint-config-file\": traffic_config_path},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Deployment\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Testing the deployment can be done through the following steps or the UI in Azure ML portal.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter the number of samples you plan to use to test the endpoint. We will keep it simple in this example and just use one.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Option 1: Import data from csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# here we're dropping the target column, which in this diabetes dataset is called \"target\" rather than \"y\" although it represents the same data\n",
    "sample = (\n",
    "    pd.read_csv(\"https://azuremlexamples.blob.core.windows.net/datasets/diabetes.csv\")\n",
    "    .sample(n)\n",
    "    .drop(columns=[\"target\"])\n",
    "    .reset_index(drop=True)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Option 2: Write data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write data directly into the dataframe.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {\n",
    "    \"age\": [0.0380759064334241],\n",
    "    \"sex\": [0.0506801187398187],\n",
    "    \"bmi\": [0.0616962065186885],\n",
    "    \"bp\": [0.0218723549949558],\n",
    "    \"s1\": [-0.0442234984244464],\n",
    "    \"s2\": [-0.0348207628376986],\n",
    "    \"s3\": [-0.0434008456520269],\n",
    "    \"s4\": [-0.00259226199818282],\n",
    "    \"s5\": [0.0199084208763183],\n",
    "    \"s6\": [-0.0176461251598052],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if n == 1:\n",
    "    sample = pd.DataFrame(data=data, index=[0])\n",
    "else:\n",
    "    sample = pd.DataFrame(data=data)\n",
    "print(f\"Sample Data: {sample.head()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoke the endpoint\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get payload of `predict()` from the endpoint:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "payload = deployment_client.predict(endpoint=endpoint_name, df=sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract predictions and explanations from payload:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "if isinstance(payload, pd.DataFrame):\n",
    "    print(\"Return type is DataFrame\")\n",
    "    predictions = payload[\"predictions\"].values\n",
    "    explanations = payload[\"explanations\"].values\n",
    "elif isinstance(payload, np.ndarray):\n",
    "    print(\"Return type is ndarray\")\n",
    "    predictions = payload.item()[\"predictions\"]\n",
    "    explanations = payload.item()[\"explanations\"]\n",
    "else:\n",
    "    print(\n",
    "        \"Return type not supported - either skip the rest of this notebook or write your own code to extract the predictions and explanations lists\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print and view:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = np.array(baseline_df.drop(columns=[target_column]).columns)\n",
    "if task_type == \"classification\":\n",
    "    classes = np.array(baseline_df[target_column].unique())\n",
    "\n",
    "for i in range(len(predictions)):\n",
    "    print(f\"For data point {i}:\")\n",
    "    print(f\"{sample.loc[i]}\\n\")\n",
    "\n",
    "    print(f\"Prediction: {predictions[i]}\\n\")\n",
    "\n",
    "    if task_type == \"classification\":\n",
    "        for j in range(len(classes)):\n",
    "            importances = np.array(explanations[i][j][0])\n",
    "            explanations_df = pd.DataFrame(\n",
    "                data={\"feature\": features, \"local importance\": importances}\n",
    "            )\n",
    "            print(f\"Feature importances for class: {classes[j]}\")\n",
    "            if predictions[i] == classes[j]:\n",
    "                print(\"This is the predicted class for this row of data\")\n",
    "            else:\n",
    "                print(\"This is NOT the predicted class for this row of data\")\n",
    "            print(f\"{explanations_df}\\n\")\n",
    "    else:\n",
    "        importances = np.array(explanations[i][0])\n",
    "        explanations_df = pd.DataFrame(\n",
    "            data={\"feature\": features, \"local importance\": importances}\n",
    "        )\n",
    "        print(f\"Feature importances: {explanations_df}\\n\")\n",
    "\n",
    "    print(\"\\n\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explore data (Optional)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import plotly\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.express as px"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Feature importances for regression\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if task_type == \"regression\":\n",
    "    for i in range(len(explanations)):\n",
    "        explanations[i] = explanations[i][0]\n",
    "\n",
    "    features = np.array(baseline_df.drop(columns=[target_column]).columns)\n",
    "    for i in range(len(predictions)):\n",
    "        print(f\"For data point {i}:\")\n",
    "        print(f\"{sample.loc[i]}\\n\")\n",
    "\n",
    "        importances = np.array(explanations[i])\n",
    "        explanations_df = pd.DataFrame(\n",
    "            data={\"feature\": features, \"local importance\": importances}\n",
    "        )\n",
    "\n",
    "        plot = px.bar(\n",
    "            data_frame=explanations_df,\n",
    "            x=\"feature\",\n",
    "            y=\"local importance\",\n",
    "            title=\"Local Feature Importance\",\n",
    "        )\n",
    "        plot.show()\n",
    "\n",
    "        print(\"\\n\\n\\n\")\n",
    "else:\n",
    "    print(\"Task type is not regression - skip this section\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Feature importances for classification\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if task_type == \"classification\":\n",
    "    features = np.array(baseline_df.drop(columns=[target_column]).columns)\n",
    "    classes = np.array(baseline_df[target_column].unique())\n",
    "    for i in range(len(predictions)):\n",
    "        print(f\"For data point {i}:\")\n",
    "        print(f\"{sample.loc[i]}\\n\")\n",
    "\n",
    "        aggregated_df = pd.DataFrame(columns=[\"class\", \"feature\", \"local importance\"])\n",
    "        for j in range(len(classes)):\n",
    "            importances = explanations[i][j][0]\n",
    "            explanations_df = pd.DataFrame(\n",
    "                data={\"feature\": features, \"local importance\": importances}\n",
    "            )\n",
    "\n",
    "            for k in range(len(features)):\n",
    "                new_row = pd.DataFrame(\n",
    "                    data=[[classes[j], features[k], explanations[i][j][0][k]]],\n",
    "                    columns=[\"class\", \"feature\", \"local importance\"],\n",
    "                )\n",
    "                aggregated_df = pd.concat([aggregated_df, new_row])\n",
    "\n",
    "            title = f\"Local Importance for Class {classes[j]}\"\n",
    "            if predictions[i] == classes[j]:\n",
    "                title += \" (Predicted Class)\"\n",
    "            else:\n",
    "                title += \" (Not the Predicted Class)\"\n",
    "            plot = px.bar(\n",
    "                data_frame=explanations_df,\n",
    "                x=\"feature\",\n",
    "                y=\"local importance\",\n",
    "                title=title,\n",
    "            )\n",
    "            plot.show()\n",
    "\n",
    "        plot = px.bar(\n",
    "            data_frame=aggregated_df,\n",
    "            x=\"feature\",\n",
    "            y=\"local importance\",\n",
    "            color=\"class\",\n",
    "            title=\"Aggregated Local Importance\",\n",
    "        )\n",
    "        plot.show()\n",
    "\n",
    "        print(\"\\n\\n\\n\")\n",
    "else:\n",
    "    print(\"Task type is not classification - skip this section\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK V2",
   "language": "python",
   "name": "python310-sdkv2"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}