sdk/python/responsible-ai/tabular/responsibleaidashboard-healthcare-covid-classification/responsibleaidashboard-healthcare-covid-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bae671b4-9dde-4218-83dc-e8662a8bffc5",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "# COVID hospitalization prediction (classification) RAI dashboard\n",
    "This notebook demonstrates the use of the `responsibleai` API to assess a classification model trained on Fabricated COVID hospitalization dataset. The model predicts if a patient was able to **avoid hospitalization** or not based on the independent features:\n",
    "- Age\n",
    "- Gender\n",
    "- Cardiovascular-disease\n",
    "- Immuno_Compromised\n",
    "- Asthma\n",
    "- Smoking\n",
    "- Temperature\n",
    "- Blood_Sugar-mg/dL\n",
    "- Blood_Oxygen-SpO2\n",
    "- BP_Systolic\n",
    "- BP_Diastolic\n",
    "- Social_activity\n",
    "\n",
    "The Data Dictionary can be accessed through the following link: [Data_dictionary_healthcare](link-URL)\n",
    "\n",
    "The Notebook walks through the API calls necessary to create a widget with model analysis insights, then guides a visual analysis of the model."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "75b333a2",
   "metadata": {},
   "source": [
    "## **Installation**  \n",
    "\n",
    "If you are **running the notebook for the first time**, you need to follow a few of steps for smooth execution of notebook:\n",
    "\n",
    "1. Un-comment the below cell.\n",
    "2. Run the cell.\n",
    "3. After execution of this cell, comment the cell.\n",
    "4. Re-start the kernel\n",
    "5. Continue with running of all cells.\n",
    "\n",
    "**Reminder** -- Be sure to set your kernel to \"Python 3.8 - AzureML,\" via the drop-down menu at the right end of the taskbar. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2cee10f",
   "metadata": {},
   "source": [
    "### Install required dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8199f182",
   "metadata": {},
   "source": [
    "**Make sure it comment the below cell while executing the notebook more than once**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "101ee01c-79f8-4a65-863a-309fc4dedfa2",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "%pip install azure-ai-ml\n",
    "%pip install sklearn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16382f36-82c6-4bbf-8ff4-a31693b873c0",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## **User Configuration**  \n",
    "Confirm the compute name listed here is the same that was created using the included ARM template.  If not, change this name so they match. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9616d51e-f224-4d17-b0d4-b01315610389",
   "metadata": {
    "gather": {
     "logged": 1680807365263
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "# Pass the name of your compute instance (See step 6 below for it's use)\n",
    "compute_name = \"raitextcluster\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e889c9d-10f6-49dd-bdd6-2c1c661afc64",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## **After changing the above cell click on Run All.**\n",
    "**The notebook will follow the below steps and complete execution in 15-30 minutes depending upon compute configurations**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf691e2-182d-49f5-b0a6-18997b260535",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Automated Notebook steps:\n",
    "\n",
    "**Step 1:** Loading the Data.\n",
    "\n",
    "**Step 2:** Pre-processing.\n",
    "\n",
    "**Step 3:** Splitting into Train Test datasets.\n",
    "\n",
    "**Step 4:** Registering the datasets as data assets in AML.\n",
    "\n",
    "**Step 5:** Define training and registering scripts for use in Training Pipeline.\n",
    "\n",
    "**Step 6:** Create compute instance (if compute instance name not passed).\n",
    "\n",
    "**Step 7:** Executing Model Training pipeline.\n",
    "\n",
    "**Step 8:** Define components for Responsible AI Dashboard Generation Pipeline (The components are explained in later parts).\n",
    "\n",
    "**Step 9:** Execute Dashboard Generation Pipeline (generate scorecard and save in directory).\n",
    "\n",
    "**Step 10:** Click on the link at the end of the notebook to access the dashboard generated."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf0d6ecb-8904-485d-859f-01025dcafa58",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Loading required modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb59e9ef-ea11-4107-a2cd-82c20e5fe85a",
   "metadata": {
    "gather": {
     "logged": 1680807368025
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "import sklearn\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.preprocessing import StandardScaler, OneHotEncoder, OrdinalEncoder\n",
    "from sklearn.compose import ColumnTransformer\n",
    "\n",
    "import zipfile\n",
    "from io import BytesIO\n",
    "import requests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca89a3ce-b4f5-4cda-b5dd-516403cd9def",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Accessing the Data\n",
    "\n",
    "The following section examines the code necessary to create datasets and a model using components in AzureML."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29bb4891-c33b-4413-a84b-353cce46f1ce",
   "metadata": {
    "gather": {
     "logged": 1680807368530
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "def get_data(data_location, independent_features, target_feature, drop_col=None):\n",
    "    \"\"\"\n",
    "    Function to read data in Pandas dataframe\n",
    "    [TODO: Add any preprocessing steps within this function]\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    data_location: string\n",
    "        Path of the Dataset\n",
    "    independent_features: list\n",
    "        List of names of the independent features\n",
    "    target_feature: string\n",
    "        Name of the target/dependent features\n",
    "    drop_col: list\n",
    "        List of column names to drop\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    df: Pandas DataFrame\n",
    "        Pandas dataframe containing the dataset with the names passed\n",
    "    \"\"\"\n",
    "    column_names = independent_features + [target_feature]\n",
    "\n",
    "    # Download the blob data from the provided URL\n",
    "    response = requests.get(data_location)\n",
    "    blob_content = response.content\n",
    "\n",
    "    with zipfile.ZipFile(BytesIO(blob_content), \"r\") as zip_ref:\n",
    "        file_list = zip_ref.namelist()\n",
    "        if len(file_list) > 0:\n",
    "            # Assume the first file in the zip contains the data\n",
    "            inner_blob_name = file_list[0]\n",
    "            inner_blob_content = zip_ref.read(inner_blob_name)\n",
    "            df = pd.read_csv(BytesIO(inner_blob_content))\n",
    "\n",
    "    # df = pd.read_csv(data_location)\n",
    "    l = list(df.columns)\n",
    "    l.remove(target_feature)\n",
    "    df = df[l + [target_feature]]\n",
    "    df.columns = column_names\n",
    "    if drop_col is not None:\n",
    "        df.drop(drop_col, axis=1, inplace=True)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe308fb9",
   "metadata": {},
   "source": [
    "### Reading & Encoding the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "707c74ea-e489-4853-8078-987b7caa64f1",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "We load the data from github Repo directly and do basic pre-processing steps."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f12cea32-6a2a-448f-9848-a1e2f5c3c29c",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "**Categorical Codes for \"inHospital\":**\n",
    "\n",
    " **Hospitalized: Patient failed to avoid hospitalization or Hospitalized due to covid**\n",
    " \n",
    " **Not_Hospitalized: Patient avoided hospitalization or Not hospitalized due to covid**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e5d6f72-41c9-48f3-b8ea-37dc2a6fffda",
   "metadata": {
    "gather": {
     "logged": 1680807371681
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "data_df = get_data(\n",
    "    # data_location=\"./Fabricated_covid_data.csv\",\n",
    "    data_location=\"https://publictestdatasets.blob.core.windows.net/data/RAI_fabricated_covid_data.zip\",\n",
    "    target_feature=\"InHospital\",\n",
    "    independent_features=[\n",
    "        \"Age-years\",\n",
    "        \"Gender\",\n",
    "        \"Cardiovascular-disease\",\n",
    "        \"Immuno_Compromised\",\n",
    "        \"Asthma\",\n",
    "        \"Smoking\",\n",
    "        \"Temperature-F\",\n",
    "        \"Blood_Sugar-mg/dL\",\n",
    "        \"Blood_Oxygen-SpO2\",\n",
    "        \"BP_Systolic\",\n",
    "        \"BP_Diastolic\",\n",
    "        \"Social_activity\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "data_encoded = data_df.copy()\n",
    "\n",
    "hospitalized_encoding = {\n",
    "    True: \"Hospitalized\",\n",
    "    False: \"Not_Hospitalized\",\n",
    "}\n",
    "\n",
    "data_encoded.replace({\"InHospital\": hospitalized_encoding}, inplace=True)\n",
    "data_encoded"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f449db05-acab-4451-ae38-685e631d9397",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "### Splitting the Data into training and test datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3d9b848-7a43-4ef8-b27f-a069e1000b55",
   "metadata": {
    "gather": {
     "logged": 1680807372630
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "data_train, data_test = train_test_split(\n",
    "    data_encoded, test_size=0.25, random_state=7, stratify=data_encoded[\"InHospital\"]\n",
    ")\n",
    "\n",
    "data_test = (\n",
    "    data_encoded.groupby(\"InHospital\")\n",
    "    .apply(lambda x: x.sample(n=500))\n",
    "    .droplevel(level=0)\n",
    ")\n",
    "\n",
    "# Download the blob data from the provided URL\n",
    "data_location = \"https://publictestdatasets.blob.core.windows.net/data/RAI_fabricated_covid_data.zip\"\n",
    "response = requests.get(data_location)\n",
    "blob_content = response.content\n",
    "\n",
    "with zipfile.ZipFile(BytesIO(blob_content), \"r\") as zip_ref:\n",
    "    file_list = zip_ref.namelist()\n",
    "    if len(file_list) > 0:\n",
    "        # Assume the first file in the zip contains the data\n",
    "        inner_blob_name = file_list[1]\n",
    "        inner_blob_content = zip_ref.read(inner_blob_name)\n",
    "        noise_df_test = pd.read_csv(BytesIO(inner_blob_content))\n",
    "\n",
    "# noise_df_test = pd.read_csv(\"./Fabricated_covid_noise_data.csv\")\n",
    "noise_df_test.replace({\"InHospital\": hospitalized_encoding}, inplace=True)\n",
    "data_test = pd.concat([data_test, noise_df_test])\n",
    "\n",
    "\n",
    "if len(data_test) <= 5000:\n",
    "    print(\"Proceed with the analysis\")\n",
    "else:\n",
    "    print(\"Reduce your test data size\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "432fa629",
   "metadata": {},
   "source": [
    "### Get the Data to AzureML\n",
    "\n",
    "With the data now split into 'train' and 'test' DataFrames, we save them out to files in preparation for upload into AzureML:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9983401-5bb9-4191-8c73-36e1782f2439",
   "metadata": {
    "gather": {
     "logged": 1680807372972
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "train_data_path = \"./data_covid_classification/train/\"\n",
    "test_data_path = \"./data_covid_classification/test/\"\n",
    "\n",
    "os.makedirs(train_data_path, exist_ok=True)\n",
    "os.makedirs(test_data_path, exist_ok=True)\n",
    "\n",
    "train_filename = train_data_path + \"covid_classification_train.parquet\"\n",
    "test_filename = test_data_path + \"covid_classification_test.parquet\"\n",
    "\n",
    "data_train.to_parquet(train_filename, index=False)\n",
    "data_test.to_parquet(test_filename, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e872771-a967-41f1-a774-b26231d5d269",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "We are going to create two Datasets in AzureML, one for the train and one for the test datasets. The first step is to create an `MLClient` to perform the upload. The method we use assumes that there is a `config.json` file (downloadable from the Azure or AzureML portals) present in the same directory as this notebook file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85bb6c3d",
   "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": "0b6e1ca6-9ef9-43c2-95f8-a98a3b93dbb4",
   "metadata": {
    "gather": {
     "logged": 1680807373951
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "# Handle to the workspace\n",
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import DefaultAzureCredential\n",
    "\n",
    "credential = DefaultAzureCredential()\n",
    "ml_client = MLClient(\n",
    "    credential=credential,\n",
    "    subscription_id=subscription_id,\n",
    "    resource_group_name=resource_group,\n",
    "    workspace_name=workspace,\n",
    ")\n",
    "print(ml_client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a9f86f2-cbe2-4b8d-8e51-31e2c018bdbc",
   "metadata": {
    "gather": {
     "logged": 1680807374370
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "# Define Version string (optional)\n",
    "rai_covid_classification_example_version_string = \"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2c758e6-a5cc-47fa-9714-16f22c8bf031",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "### Create an asset MLtable (or URI file) to register the Data into workspace\n",
    "This is essential,  as the dashboard recognizes only registered assets. \n",
    "\n",
    "Reference:\n",
    "https://learn.microsoft.com/en-us/azure/machine-learning/how-to-create-data-assets?tabs=Python-SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6629e7f2-fe9f-44f5-8244-d9d04c670d04",
   "metadata": {
    "gather": {
     "logged": 1680807374739
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Data\n",
    "from azure.ai.ml.constants import AssetTypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3f3065f-cc91-489e-8aae-d7d9539c1de4",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "#### Change the asset name of the below file if the train/test data has changed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2b4360e-d977-43c5-a543-0a634532c5f4",
   "metadata": {
    "gather": {
     "logged": 1680807375515
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "input_train_data = \"train_covid_classification\"\n",
    "\n",
    "try:\n",
    "    # Try getting data already registered in workspace\n",
    "    train_data = ml_client.data.get(\n",
    "        name=input_train_data,\n",
    "        version=rai_covid_classification_example_version_string,\n",
    "    )\n",
    "\n",
    "except Exception as e:\n",
    "    train_data = Data(\n",
    "        path=train_filename,\n",
    "        type=AssetTypes.URI_FILE,\n",
    "        description=\"RAI covid classification example training data\",\n",
    "        name=input_train_data,\n",
    "        version=rai_covid_classification_example_version_string,\n",
    "    )\n",
    "    ml_client.data.create_or_update(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6e9317e-d755-4128-9700-733e16a9bec3",
   "metadata": {
    "gather": {
     "logged": 1680807376118
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "# Change the asset name of the below file if the test data has changed\n",
    "\n",
    "input_test_data = \"test_covid_classification\"\n",
    "\n",
    "try:\n",
    "    # Try getting data already registered in workspace\n",
    "    test_data = ml_client.data.get(\n",
    "        name=input_test_data,\n",
    "        version=rai_covid_classification_example_version_string,\n",
    "    )\n",
    "\n",
    "except Exception as e:\n",
    "    test_data = Data(\n",
    "        path=test_filename,\n",
    "        type=AssetTypes.URI_FILE,\n",
    "        description=\"RAI covid classification example test data\",\n",
    "        name=input_test_data,\n",
    "        version=rai_covid_classification_example_version_string,\n",
    "    )\n",
    "    ml_client.data.create_or_update(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f66aeab-a4bd-4ca8-9e96-fc141f87d2ca",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "# Create a model training pipeline\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",
    "2. 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",
    "The training component is for this particular model. In this case, we are going to train an `Gaussian Naive Bayes classifier` 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",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c0f9bde",
   "metadata": {
    "gather": {
     "logged": 1680807376485
    }
   },
   "outputs": [],
   "source": [
    "os.makedirs(\"./component_src\", exist_ok=True)\n",
    "os.makedirs(\"./register_model_src\", exist_ok=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ab759d4c-d62e-4e3d-8436-8828490632ee",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "**Create the training script**  \n",
    "This cell creates a machine learning pipeline that trains a Gaussian Naive Bayes classifier using labeled data and then saves the trained model to a specified output path using MLFlow. \n",
    "- The code reads in the training data as a pandas dataframe from a specified path, extracts the target column name, and separates the target column from the feature columns. \n",
    "- Feature columns are then preprocessed using both a standard scaler for numeric data and a one-hot encoder for categorical data. \n",
    "- Preprocessed feature columns and target column are then fed into the Gaussian Naive Bayes classifier. \n",
    "- The trained model is saved to a temporary directory and then copied to the specified output path. \n",
    "- The code takes command-line arguments for the paths of the training data, the output model, and the name of the target column. \n",
    "- The code also uses the Azure Machine Learning (AML) Python SDK to log the model and tracking information with MLFlow. \n",
    "- Additional comments in the code provide details on each section of the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb4e1af",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile component_src/classification_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",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.preprocessing import StandardScaler, OneHotEncoder, OrdinalEncoder\n",
    "from sklearn.compose import ColumnTransformer\n",
    "\n",
    "import pandas as pd\n",
    "from sklearn.naive_bayes import GaussianNB\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",
    "    all_data = pd.read_parquet(args.training_data)\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(\"Executing Model Training pipeline\")\n",
    "    # We create the preprocessing pipelines for both numeric and categorical data.\n",
    "    numeric_transformer = Pipeline(steps=[\n",
    "        ('scaler', StandardScaler())])\n",
    "\n",
    "    categorical_transformer = Pipeline(steps=[\n",
    "        ('onehot', OneHotEncoder(handle_unknown='ignore'))])\n",
    "\n",
    "    continuous_features_names = ['Age-years', 'BP_Diastolic', 'BP_Systolic', 'Blood_Oxygen-SpO2',\n",
    "                                'Blood_Sugar-mg/dL', 'Temperature-F']\n",
    "    categorical_features_names = ['Gender', 'Cardiovascular-disease', 'Immuno_Compromised',\n",
    "       'Asthma', 'Smoking', 'Social_activity']\n",
    "    transformations = ColumnTransformer(\n",
    "        transformers=[\n",
    "            ('num', numeric_transformer, continuous_features_names),\n",
    "            ('cat', categorical_transformer, categorical_features_names)])\n",
    "\n",
    "    # Append classifier to preprocessing pipeline.\n",
    "    # Now we have a full prediction pipeline.\n",
    "    # The estimator can be changed to suit\n",
    "    model = Pipeline(steps=[('preprocessor', transformations),\n",
    "                          ('classifier', GaussianNB())])\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": "markdown",
   "id": "c702f24b-ad0e-49b9-8895-94d48a9a2f9c",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "**Define the YAML file**\n",
    "\n",
    "This code snippet defines an Azure Machine Learning Command Component for training a classification model on a dataset. It starts by defining a YAML configuration file that specifies the inputs and outputs of the component, the command to run, and the environment to use. The YAML file is then saved to disk.\n",
    "\n",
    "Next, the code uses the Azure ML Python SDK to load the Command Component from the YAML file. The resulting object can be used to run the component on a dataset, passing in the input paths and output paths as arguments.\n",
    "\n",
    "Overall, this code provides a simple and reusable way to define and run machine learning training components in Azure ML."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a0ce700",
   "metadata": {
    "gather": {
     "logged": 1680807378226
    }
   },
   "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_classification_training_component\n",
    "display_name: Classification training component for RAI example\n",
    "version: {rai_covid_classification_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-tabular/versions/4\n",
    "\"\"\"\n",
    "    + r\"\"\"\n",
    "command: >-\n",
    "  python classification_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 = \"RAIcovidTrainingComponent.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": "markdown",
   "id": "4c978257-8be5-4d7e-9354-28ab21282453",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "This script loads a trained model, registers it via MLFlow, and saves the registered model information to a JSON file. Users need to provide the necessary arguments to register the model, including the path to the input model, path to the output model info JSON file, base name of the registered model, and an optional suffix for the registered model name.\n",
    "\n",
    "To use this script, the following arguments must be defined: \n",
    "- model_input_path: Path to the input model  \n",
    "- model_info_output_path: Path to write the model info JSON  \n",
    "- model_base_name: Name of the registered model  \n",
    "- model_name_suffix: An integer value to add as a suffix to the registered model name. If this is negative, the epoch time is used as the suffix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "280531f8",
   "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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b7e1b2",
   "metadata": {},
   "source": [
    "Now that the model registration 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 model registration script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ad3244d-096f-4e74-9836-54d7e84c5c7a",
   "metadata": {
    "gather": {
     "logged": 1680807378999
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "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: {rai_covid_classification_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-tabular/versions/4\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": "1e65fc2e-7979-4316-995f-90c6f419aec9",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "We will create a new compute instance to run the jobs if it does not already exist by the name passed in the beginning of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d04d6675-9de0-4e8e-9d9d-f92a2a2d2259",
   "metadata": {
    "gather": {
     "logged": 1680807379589
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "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_DS4_V2\",\n",
    "        min_instances=0,\n",
    "        max_instances=1,\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": "693f7706",
   "metadata": {},
   "source": [
    "### Running the training pipeline\n",
    "\n",
    "The 2 YAML files (\"RAIcovidTrainingComponent.yaml\" & \"register.yaml\") are used to define the 2 components in the model training pipeline\n",
    "\n",
    "We start by ensuring that the compute cluster named above exists:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0917f58e-5171-426d-80f8-a27d793b4aa9",
   "metadata": {
    "gather": {
     "logged": 1680807380089
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "model_name_suffix = int(time.time())\n",
    "model_name = \"rai_covid_classsification_model\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c8dbc93-5a39-462d-903c-8e5ded80098d",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "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": "9dcfd57a-03b5-446c-bf87-af2db9aadac8",
   "metadata": {
    "gather": {
     "logged": 1680807380534
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "from azure.ai.ml import dsl, Input\n",
    "\n",
    "target_feature = \"InHospital\"\n",
    "categorical_features = [\n",
    "    \"Gender\",\n",
    "    \"Cardiovascular-disease\",\n",
    "    \"Immuno_Compromised\",\n",
    "    \"Asthma\",\n",
    "    \"Smoking\",\n",
    "    \"Social_activity\",\n",
    "]\n",
    "\n",
    "\n",
    "covid_train_pq = Input(\n",
    "    type=\"uri_file\",\n",
    "    path=f\"azureml:{input_train_data}:{rai_covid_classification_example_version_string}\",\n",
    "    mode=\"download\",\n",
    ")\n",
    "covid_test_pq = Input(\n",
    "    type=\"uri_file\",\n",
    "    path=f\"azureml:{input_test_data}:{rai_covid_classification_example_version_string}\",\n",
    "    mode=\"download\",\n",
    ")\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    compute=compute_name,\n",
    "    description=\"Register Model for RAI covid classification example\",\n",
    "    experiment_name=f\"RAI_classification_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=1200)\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_feature, covid_train_pq)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d7f7771-f4df-4493-8fbc-1bad044474d2",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "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": "ea80ebc2-099d-4a22-959a-3ceb1fc087d7",
   "metadata": {
    "gather": {
     "logged": 1680807473024
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "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": "markdown",
   "id": "c9889100",
   "metadata": {},
   "source": [
    "##  Creating the RAI Insights\n",
    "\n",
    "We have a registered model, and can now run a pipeline to create the RAI insights. First off, compute the name of the model we registered:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c65da4a0-81da-4d30-b8e0-65be2a2caaee",
   "metadata": {
    "gather": {
     "logged": 1680807473381
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "expected_model_id = f\"{model_name}_{model_name_suffix}:1\"\n",
    "azureml_model_id = f\"azureml:{expected_model_id}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4bad8a5",
   "metadata": {},
   "source": [
    "Now, we create the RAI pipeline itself. There are four 'component stages' in this pipeline:\n",
    "\n",
    "1. Construct an empty `RAIInsights` object\n",
    "1. Run the RAI tool components\n",
    "1. Gather the tool outputs into a single `RAIInsights` object\n",
    "1. (Optional) Generate a score card in pdf format summarizing model performance, and key aspects from the rai tool components\n",
    "\n",
    "We start by loading the RAI component definitions for use in our pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f52201d1-6962-4c3a-b27a-29c3a2547a98",
   "metadata": {
    "gather": {
     "logged": 1680807474102
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "# Get handle to azureml registry for the RAI built in components\n",
    "registry_name = \"azureml\"\n",
    "ml_client_registry = MLClient(\n",
    "    credential=credential,\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": "32e1594d-0bba-43cc-bc20-44f761881350",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "## Add different components of ResponsibleAI dashboard to the Pipeline\n",
    "\n",
    "Next, we retrive various RAI (Responsible AI) built-in components from the ml_client_registry.components object and assigns them to different variables. The label variable is used to specify which version of the rai_constructor_component to use, and the version variable is used to ensure that all other components are of the same version as the rai_constructor_component. Finally, the current version of the RAI built-in components is printed to the console.\n",
    "\n",
    "\n",
    "Reference:\n",
    "https://learn.microsoft.com/en-us/azure/machine-learning/how-to-responsible-ai-insights-sdk-cli?tabs=python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1fe7afc-2d6c-4a7d-9108-edd78eb019ae",
   "metadata": {
    "gather": {
     "logged": 1680807474379
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "label = \"latest\"\n",
    "\n",
    "rai_constructor_component = ml_client_registry.components.get(\n",
    "    name=\"microsoft_azureml_rai_tabular_insight_constructor\", label=label\n",
    ")\n",
    "\n",
    "# We get latest version and use the same version for all components\n",
    "version = rai_constructor_component.version\n",
    "print(\"The current version of RAI built-in components is: \" + version)\n",
    "\n",
    "rai_causal_component = ml_client_registry.components.get(\n",
    "    name=\"microsoft_azureml_rai_tabular_causal\", version=version\n",
    ")\n",
    "\n",
    "rai_counterfactual_component = ml_client_registry.components.get(\n",
    "    name=\"microsoft_azureml_rai_tabular_counterfactual\", version=version\n",
    ")\n",
    "rai_erroranalysis_component = ml_client_registry.components.get(\n",
    "    name=\"microsoft_azureml_rai_tabular_erroranalysis\", version=version\n",
    ")\n",
    "\n",
    "rai_explanation_component = ml_client_registry.components.get(\n",
    "    name=\"microsoft_azureml_rai_tabular_explanation\", version=version\n",
    ")\n",
    "\n",
    "rai_gather_component = ml_client_registry.components.get(\n",
    "    name=\"microsoft_azureml_rai_tabular_insight_gather\", version=version\n",
    ")\n",
    "\n",
    "rai_scorecard_component = ml_client_registry.components.get(\n",
    "    name=\"microsoft_azureml_rai_tabular_score_card\", version=version\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64a22471",
   "metadata": {},
   "source": [
    "## Score card generation config\n",
    "For score card generation, we need some additional configuration in a separate json file. Here we configure the following model performance metrics for reporting:\n",
    "- accuracy\n",
    "- precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f740ae68-8379-4d04-8445-ce5111dbc9c1",
   "metadata": {
    "gather": {
     "logged": 1680807474708
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "score_card_config_dict = {\n",
    "    \"Model\": {\n",
    "        \"ModelName\": \"covid classification\",\n",
    "        \"ModelType\": \"Classification\",\n",
    "        \"ModelSummary\": \"<model summary>\",\n",
    "    },\n",
    "    \"Metrics\": {\"accuracy_score\": {\"threshold\": \">=0.5\"}, \"precision_score\": {}},\n",
    "}\n",
    "\n",
    "score_card_config_filename = \"rai_covid_classification_score_card_config.json\"\n",
    "\n",
    "with open(score_card_config_filename, \"w\") as f:\n",
    "    json.dump(score_card_config_dict, f)\n",
    "\n",
    "score_card_config_path = Input(\n",
    "    type=\"uri_file\", path=score_card_config_filename, mode=\"download\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cecaea36",
   "metadata": {},
   "source": [
    "Now define the pipeline itself. This creates an empty `RAIInsights` object, adds the analyses, and then gathers everything into the final `RAIInsights` output. Where complex objects need to be passed (such as a list of treatment feature names), they must be encoded as JSON strings.\n",
    "\n",
    "The included components are:\n",
    "- Set limits for each component job within the pipeline\n",
    "- Initiate RAI Insights\n",
    "- Add explanation to RAI Insights\n",
    "- Add causal analysis to RAI Insights\n",
    "- Add counterfactual analysis to RAI Insights\n",
    "- Add error analysis to RAI Insights\n",
    "- Combine all the RAI insights together\n",
    "- Generate a scorecard in PDF format for a summary report on model performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86692d9a-351e-404f-9ba0-1051d60be7f2",
   "metadata": {
    "gather": {
     "logged": 1680807474973
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from azure.ai.ml import Input\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "\n",
    "classes_in_target = json.dumps([\"Hospitalized\", \"Not_Hospitalized\"])\n",
    "treatment_features = json.dumps(\n",
    "    [\n",
    "        \"BP_Systolic\",\n",
    "        \"Blood_Oxygen-SpO2\",\n",
    "        \"Blood_Sugar-mg/dL\",\n",
    "        \"Temperature-F\",\n",
    "        \"Social_activity\",\n",
    "    ]\n",
    ")\n",
    "\n",
    "timeout = 2400\n",
    "\n",
    "\n",
    "@dsl.pipeline(\n",
    "    compute=compute_name,\n",
    "    description=\"Example RAI computation on covid Classification\",\n",
    "    experiment_name=f\"RAI_covid_Classification_Example_RAIInsights_Computation_{model_name_suffix}\",\n",
    ")\n",
    "def rai_classification_pipeline(\n",
    "    target_column_name,\n",
    "    train_data,\n",
    "    test_data,\n",
    "    score_card_config_path,\n",
    "):\n",
    "    # Initiate the RAIInsights\n",
    "    create_rai_job = rai_constructor_component(\n",
    "        title=\"RAI Dashboard Example\",\n",
    "        task_type=\"classification\",\n",
    "        model_info=expected_model_id,\n",
    "        model_input=Input(type=AssetTypes.MLFLOW_MODEL, path=azureml_model_id),\n",
    "        train_dataset=train_data,\n",
    "        test_dataset=test_data,\n",
    "        target_column_name=target_column_name,\n",
    "        categorical_column_names=json.dumps(categorical_features),\n",
    "        classes=classes_in_target,\n",
    "    )\n",
    "    create_rai_job.set_limits(timeout=timeout)\n",
    "\n",
    "    # Add an explanation\n",
    "    explain_job = rai_explanation_component(\n",
    "        comment=\"Explanation for the classification dataset\",\n",
    "        rai_insights_dashboard=create_rai_job.outputs.rai_insights_dashboard,\n",
    "    )\n",
    "    explain_job.set_limits(timeout=timeout)\n",
    "\n",
    "    # Add causal analysis\n",
    "    causal_job = rai_causal_component(\n",
    "        rai_insights_dashboard=create_rai_job.outputs.rai_insights_dashboard,\n",
    "        treatment_features=treatment_features,\n",
    "    )\n",
    "    causal_job.set_limits(timeout=timeout)\n",
    "\n",
    "    # Add counterfactual analysis\n",
    "    counterfactual_job = rai_counterfactual_component(\n",
    "        rai_insights_dashboard=create_rai_job.outputs.rai_insights_dashboard,\n",
    "        total_cfs=10,\n",
    "        desired_class=\"opposite\",\n",
    "    )\n",
    "    counterfactual_job.set_limits(timeout=timeout)\n",
    "\n",
    "    # Add error analysis\n",
    "    erroranalysis_job = rai_erroranalysis_component(\n",
    "        rai_insights_dashboard=create_rai_job.outputs.rai_insights_dashboard,\n",
    "    )\n",
    "    erroranalysis_job.set_limits(timeout=timeout)\n",
    "\n",
    "    # Combine everything\n",
    "    rai_gather_job = rai_gather_component(\n",
    "        constructor=create_rai_job.outputs.rai_insights_dashboard,\n",
    "        insight_1=explain_job.outputs.explanation,\n",
    "        insight_2=causal_job.outputs.causal,\n",
    "        insight_3=counterfactual_job.outputs.counterfactual,\n",
    "        insight_4=erroranalysis_job.outputs.error_analysis,\n",
    "    )\n",
    "    rai_gather_job.set_limits(timeout=timeout)\n",
    "\n",
    "    rai_gather_job.outputs.dashboard.mode = \"upload\"\n",
    "    rai_gather_job.outputs.ux_json.mode = \"upload\"\n",
    "\n",
    "    # Generate score card in pdf format for a summary report on model performance,\n",
    "    # and observe distrbution of error between prediction vs ground truth.\n",
    "    rai_scorecard_job = rai_scorecard_component(\n",
    "        dashboard=rai_gather_job.outputs.dashboard,\n",
    "        pdf_generation_config=score_card_config_path,\n",
    "    )\n",
    "\n",
    "    return {\n",
    "        \"dashboard\": rai_gather_job.outputs.dashboard,\n",
    "        \"ux_json\": rai_gather_job.outputs.ux_json,\n",
    "        \"scorecard\": rai_scorecard_job.outputs.scorecard,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d73be27-7ae1-4bb6-a1ba-7a75b7f57190",
   "metadata": {},
   "source": [
    "Next, we define the pipeline object itself, and ensure that the outputs will be available for download:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b877890-c3d8-45d8-b91f-906968b7b053",
   "metadata": {
    "gather": {
     "logged": 1680807475288
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "import uuid\n",
    "from azure.ai.ml import Output\n",
    "\n",
    "# Pipeline to construct the RAI Insights\n",
    "insights_pipeline_job = rai_classification_pipeline(\n",
    "    target_column_name=target_feature,\n",
    "    train_data=covid_train_pq,\n",
    "    test_data=covid_test_pq,\n",
    "    score_card_config_path=score_card_config_path,\n",
    ")\n",
    "\n",
    "# Workaround to enable the download\n",
    "rand_path = str(uuid.uuid4())\n",
    "insights_pipeline_job.outputs.dashboard = Output(\n",
    "    path=f\"azureml://datastores/workspaceblobstore/paths/{rand_path}/dashboard/\",\n",
    "    mode=\"upload\",\n",
    "    type=\"uri_folder\",\n",
    ")\n",
    "insights_pipeline_job.outputs.ux_json = Output(\n",
    "    path=f\"azureml://datastores/workspaceblobstore/paths/{rand_path}/ux_json/\",\n",
    "    mode=\"upload\",\n",
    "    type=\"uri_folder\",\n",
    ")\n",
    "insights_pipeline_job.outputs.scorecard = Output(\n",
    "    path=f\"azureml://datastores/workspaceblobstore/paths/{rand_path}/scorecard/\",\n",
    "    mode=\"upload\",\n",
    "    type=\"uri_folder\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0c05a8a",
   "metadata": {},
   "source": [
    "And submit the pipeline to AzureML for execution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b806764e-f0d3-401f-98cb-744f6cff159e",
   "metadata": {
    "gather": {
     "logged": 1680809200997
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "insights_job = submit_and_wait(ml_client, insights_pipeline_job)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69721140",
   "metadata": {},
   "source": [
    "## Downloading the Scorecard PDF\n",
    "\n",
    "We can download the scorecard PDF from our pipeline as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b2a0db-3efc-444c-aec1-8a2a00d70bdb",
   "metadata": {
    "gather": {
     "logged": 1680809201871
    }
   },
   "outputs": [],
   "source": [
    "target_directory = \".\"\n",
    "\n",
    "ml_client.jobs.download(\n",
    "    insights_job.name, download_path=target_directory, output_name=\"scorecard\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9156f22",
   "metadata": {},
   "source": [
    "## To Access the Dashboard follow the link below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "970faf39-3ead-4488-a903-f03b2dc97cf5",
   "metadata": {
    "gather": {
     "logged": 1680809201451
    },
    "jupyter": {
     "outputs_hidden": false,
     "source_hidden": false
    },
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "outputs": [],
   "source": [
    "sub_id = ml_client._operation_scope.subscription_id\n",
    "rg_name = ml_client._operation_scope.resource_group_name\n",
    "ws_name = ml_client.workspace_name\n",
    "\n",
    "expected_uri = f\"https://ml.azure.com/model/{expected_model_id}/model_analysis?wsid=/subscriptions/{sub_id}/resourcegroups/{rg_name}/workspaces/{ws_name}\"\n",
    "\n",
    "print(f\"Please visit {expected_uri} to see your analysis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83d4b5d0-1704-49fb-a65c-8b30c58e7b0a",
   "metadata": {
    "nteract": {
     "transient": {
      "deleting": false
     }
    }
   },
   "source": [
    "Once this is complete, we can go to the Registered Models view in the AzureML portal, and find the model we have just registered. On the 'Model Details' page, there is a \"Responsible AI dashboard\" tab where we can view the insights which we have just uploaded."
   ]
  }
 ],
 "metadata": {
  "categories": [
   "SDK v2",
   "sdk",
   "python",
   "responsible-ai"
  ],
  "kernel_info": {
   "name": "python310-sdkv2"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  },
  "microsoft": {
   "host": {
    "AzureML": {
     "notebookHasBeenCompleted": true
    }
   },
   "ms_spell_check": {
    "ms_ignore_dictionary": [
     "covid"
    ],
    "ms_spell_check_language": "en"
   }
  },
  "nteract": {
   "version": "nteract-front-end@1.0.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "8fd340b5477ca1a0b454d48a3973beff39fee032ada47a04f6f3725b469a8988"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}