WIT_COMPAS_with_SHAP.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "WIT COMPAS with SHAP",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f1Id_0GukgNv",
        "colab_type": "text"
      },
      "source": [
        "## What-If Tool and SHAP on COMPAS keras model\n",
        "\n",
        "This notebook shows:\n",
        "- Training of a keras model on the [COMPAS](https://www.kaggle.com/danofer/compass) dataset.\n",
        "- Explanation of inference results using [SHAP](https://github.com/slundberg/shap).\n",
        "- Use of What-If Tool on the trained model, including SHAP values.\n",
        "\n",
        "For ML fairness background on COMPAS see:\n",
        "\n",
        "- https://www.propublica.org/article/machine-bias-risk-assessments-in-criminal-sentencing\n",
        "- https://www.propublica.org/article/how-we-analyzed-the-compas-recidivism-algorithm\n",
        "- http://www.crj.org/assets/2017/07/9_Machine_bias_rejoinder.pdf\n",
        "\n",
        "This notebook trains a model to mimic the behavior of the COMPAS recidivism classifier and uses the SHAP library to provide feature importance for each prediction by the model. We can then analyze our COMPAS proxy model for fairness using the What-If Tool, and explore how important each feature was to each prediction through the SHAP values.\n",
        "\n",
        "The specific binary classification task for this model is to determine if a person belongs in the \"Low\" risk class according to COMPAS (negative class), or the \"Medium\" or \"High\" risk class (positive class). We then analyze it with the What-If Tool for its ability to predict recidivism within two years of arrest.\n",
        "\n",
        "A simpler version of this notebook that doesn't make use of the SHAP explainer can be found [here](https://colab.research.google.com/github/pair-code/what-if-tool/blob/master/WIT_COMPAS.ipynb).\n",
        "\n",
        "Copyright 2019 Google LLC.\n",
        "SPDX-License-Identifier: Apache-2.0"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x1HvYDrvor2i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#@title Install What-If Tool Widget and SHAP library\n",
        "!pip install --upgrade --quiet witwidget shap"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wzEwkr3SoyLh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#@title Read training dataset from CSV {display-mode: \"form\"}\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "tf.compat.v1.disable_v2_behavior()\n",
        "import witwidget\n",
        "import os\n",
        "import pickle\n",
        "\n",
        "from tensorflow.keras.layers import Dense\n",
        "from tensorflow.keras.models import Sequential\n",
        "\n",
        "from sklearn.utils import shuffle\n",
        "\n",
        "df = pd.read_csv('https://storage.googleapis.com/what-if-tool-resources/computefest2019/cox-violent-parsed_filt.csv')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DBx6BTumpb1-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Preprocess the data\n",
        "\n",
        "# Filter out entries with no indication of recidivism or no compass score\n",
        "df = df[df['is_recid'] != -1]\n",
        "df = df[df['decile_score'] != -1]\n",
        "\n",
        "# Rename recidivism column\n",
        "df['recidivism_within_2_years'] = df['is_recid']\n",
        "\n",
        "# Make the COMPASS label column numeric (0 and 1), for use in our model\n",
        "df['COMPASS_determination'] = np.where(df['score_text'] == 'Low', 0, 1)\n",
        "\n",
        "df = pd.get_dummies(df, columns=['sex', 'race'])\n",
        "\n",
        "# Get list of all columns from the dataset we will use for model input or output.\n",
        "input_features = ['sex_Female', 'sex_Male', 'age', 'race_African-American', 'race_Caucasian', 'race_Hispanic', 'race_Native American', 'race_Other', 'priors_count', 'juv_fel_count', 'juv_misd_count', 'juv_other_count']\n",
        "\n",
        "to_keep = input_features + ['recidivism_within_2_years', 'COMPASS_determination']\n",
        "\n",
        "to_remove = [col for col in df.columns if col not in to_keep]\n",
        "df = df.drop(columns=to_remove)\n",
        "\n",
        "input_columns = df.columns.tolist()\n",
        "labels = df['COMPASS_determination']\n",
        "df.head()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U0ZfePT1rTmZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create data structures needing for training and testing.\n",
        "# The training data doesn't contain the column we are predicting,\n",
        "# 'COMPASS_determination', or the column we are using for evaluation of our\n",
        "# trained model, 'recidivism_within_2_years'.\n",
        "df_for_training = df.drop(columns=['COMPASS_determination', 'recidivism_within_2_years'])\n",
        "train_size = int(len(df_for_training) * 0.8)\n",
        "\n",
        "train_data = df_for_training[:train_size]\n",
        "train_labels = labels[:train_size]\n",
        "\n",
        "test_data_with_labels = df[train_size:]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5T2XThgosWX-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create the model\n",
        "\n",
        "# This is the size of the array we'll be feeding into our model for each example\n",
        "input_size = len(train_data.iloc[0])\n",
        "\n",
        "model = Sequential()\n",
        "model.add(Dense(200, input_shape=(input_size,), activation='relu'))\n",
        "model.add(Dense(50, activation='relu'))\n",
        "model.add(Dense(25, activation='relu'))\n",
        "model.add(Dense(1, activation='sigmoid'))\n",
        "\n",
        "model.compile(loss='mean_squared_error', optimizer='adam')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fjgNhSCDsayt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Train the model\n",
        "model.fit(train_data.values, train_labels.values, epochs=4, batch_size=32, validation_split=0.1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lI18CwYiQotq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create a SHAP explainer by passing a subset of our training data\n",
        "import shap\n",
        "explainer = shap.DeepExplainer(model, train_data.values[:200])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iywHwbJJkeYG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Explain predictions of the model on the first 5 examples from our training set\n",
        "# to test the SHAP explainer.\n",
        "shap_values = explainer.shap_values(train_data.values[:5])\n",
        "shap_values"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KF00pJvkeicT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#@title Show model results and SHAP values in WIT\n",
        "from witwidget.notebook.visualization import WitWidget, WitConfigBuilder\n",
        "num_datapoints = 1000  #@param {type: \"number\"}\n",
        "\n",
        "# Column indices to strip out from data from WIT before passing it to the model.\n",
        "columns_not_for_model_input = [\n",
        "    test_data_with_labels.columns.get_loc(\"recidivism_within_2_years\"),\n",
        "    test_data_with_labels.columns.get_loc(\"COMPASS_determination\")\n",
        "]\n",
        "\n",
        "# Return model predictions and SHAP values for each inference.\n",
        "def custom_predict_with_shap(examples_to_infer):\n",
        "  # Delete columns not used by model\n",
        "  model_inputs = np.delete(\n",
        "      np.array(examples_to_infer), columns_not_for_model_input, axis=1)\n",
        "\n",
        "  # Get the class predictions from the model.\n",
        "  preds = model.predict(model_inputs)\n",
        "  preds = [[1 - pred[0], pred[0]] for pred in preds]\n",
        "\n",
        "  # Get the SHAP values from the explainer and create a map of feature name\n",
        "  # to SHAP value for each example passed to the model.\n",
        "  shap_output = explainer.shap_values(np.array(model_inputs))[0]\n",
        "  attributions = []\n",
        "  for shap in shap_output:\n",
        "    attrs = {}\n",
        "    for i, col in enumerate(df_for_training.columns):\n",
        "      attrs[col] = shap[i]\n",
        "    attributions.append(attrs)\n",
        "  ret = {'predictions': preds, 'attributions': attributions}\n",
        "  return ret\n",
        "\n",
        "examples_for_shap_wit = test_data_with_labels.values.tolist()\n",
        "column_names = test_data_with_labels.columns.tolist()\n",
        "\n",
        "config_builder = WitConfigBuilder(\n",
        "    examples_for_shap_wit[:num_datapoints],\n",
        "    feature_names=column_names).set_custom_predict_fn(\n",
        "  custom_predict_with_shap).set_target_feature('recidivism_within_2_years')\n",
        "\n",
        "ww = WitWidget(config_builder, height=800)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AsFCKdt2HJHO",
        "colab_type": "text"
      },
      "source": [
        "#### What-If Tool exploration ideas\n",
        "\n",
        "- Organize datapoints by \"inference score\" (can do this through binning or use of scatter plot) to see points ordered by how likely they were determined to re-offend.\n",
        "  - Select a point near the boundary line (where red points turn to blue points)\n",
        "  - Find the nearest counterfactual to see a similar person with a different decision. What is different?\n",
        "  - Look at the partial dependence plots for the selected person. What changes in what features would change the decision on this person?\n",
        "- Explore the attribution values provided by SHAP.\n",
        "  - For a variety of selected datapoints, look at which features have the highest positive attribution values. These are making the model predict higher risk for this person.\n",
        "  - Look at which features have the lowest negative attribution values as well. These are making the model predict lower risk for this person.\n",
        "  - How well do these attribution scores line up with the partial dependence plots for those datapoints?\n",
        "  - Use the attribution scores in the datapoints visualizations to look for interesting patterns. As one example, you could set the scatter X-axis to \"attributions__age\" and the scatter Y-axis to \"attributions__priors_count\" with the points colored by \"Inference score\" to investigate the relationship between feature importance of those two features and how those relate to the score the model gives for each datapoint being \"High risk\".\n",
        "- In \"Performance and Fairness\" tab, slice the dataset by different features (such as race or sex)\n",
        "  - Look at the confusion matrices for each slice - How does performance compare in those slices? What from the training data may have caused the difference in performance between the slices? What root causes could exist?\n",
        "  - Use the threshold optimization buttons to optimize positive classification thresholds for each slice based on any of the possible fairness constraints - How different do the thresholds have to be to achieve that constraint? How varied are the thresholds depending on the fairness constraint chosen?\n",
        "\n",
        "- In the \"Performance + Fairness\" tab, change the cost ratio so that you can optimize the threshold based off of a non-symmetric cost of false positives vs false negatives. Then click the \"optimize threshold\" button and see the effect on the confusion matrix. \n",
        "  - Slice the dataset by a feature, such as sex or race. How has the new cost ratio affected the disparity in performance between slices? Click the different threshold optimization buttons to see how the changed cost ratio affects the disparity given different fairness constraints.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W9qxcBj72Q_E",
        "colab_type": "text"
      },
      "source": [
        "#### Further exploration ideas\n",
        "\n",
        "- Edit the training data so that race fields are not included as a feature and train a new model with this data as input (make sure to create a new explainer and a new custom prediction function that filters race out of model input and uses the right explainer and model).\n",
        "- Load the new model with set_compare_custom_predict_fn and compare it with the original model.\n",
        "  - HINT: You'll need to make edits in 3 separate code cells.\n",
        "  - Is there still a racial disparity in model results? If so, what could be the causes?\n",
        "  - How did the SHAP attributions change?"
      ]
    }
  ]
}