Copy_of_cd_context_20newsgroup.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/rsong0606/Getting-Things-Done-with-Pytorch/blob/master/Copy_of_cd_context_20newsgroup.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6a288462",
      "metadata": {
        "id": "6a288462"
      },
      "source": [
        "# Context-aware drift detection on news articles\n",
        "\n",
        "## Introduction\n",
        "\n",
        "In this notebook we show how to **detect drift on text data given a specific context** using the [context-aware MMD detector](https://docs.seldon.io/projects/alibi-detect/en/latest/cd/methods/contextmmddrift.html) ([Cobb and Van Looveren, 2022](https://arxiv.org/abs/2203.08644)). Consider the following simple example: the upcoming elections result in an increase of political news articles compared to other topics such as sports or science. Given the context (the elections), it is however not surprising that we observe this uptick. Moreover, assume we have a machine learning model which is trained to classify news topics, and this model performs well on political articles. So given that we fully expect this uptick to occur given the context, and that our model performs fine on the political news articles, we do not want to flag this type of drift in the data. **This setting corresponds more closely to many real-life settings than traditional drift detection where we make the assumption that both the reference and test data are i.i.d. samples from their underlying distributions.**\n",
        "\n",
        "In our news topics example, each different topic such as politics, sports or weather represents a subpopulation of the data. Our context-aware drift detector can then detect changes in the data distribution which cannot be attributed to a change in the relative prevalences of these subpopulations, which we deem permissible. As a cherry on the cake, the context-aware detector allows you to understand which subpopulations are present in both the reference and test data. This allows you to obtain deep insights into the distribution underlying the test data.\n",
        "\n",
        "Useful context (or *conditioning*) variables for the context-aware drift detector include but are not limited to:\n",
        "\n",
        "1. **Domain or application specific** contexts such as the time of day or the weather.\n",
        "2. Conditioning on the **relative prevalences of known subpopulations**, such as the frequency of political articles. It is important to note that while the relative frequency of each subpopulation might change, the distribution underlying each subpopulation cannot change.\n",
        "3. Conditioning on **model predictions**. Assume we trained a classifier which tries to figure out which news topic an article belongs to. Given our model predictions we then want to understand whether our test data follows the same underlying distribution as reference instances with similar model predictions. This conditioning would also be useful in case of trending news topics which cause the model prediction distribution to shift but not necessarily the distribution within each of the news topics.\n",
        "4. Conditioning on **model uncertainties** which would allow increases in model uncertainty due to drift into familiar regions of high aleatoric uncertainty (often fine) to be distinguished from that into unfamiliar regions of high epistemic uncertainty (often problematic).\n",
        "\n",
        "The following settings will be illustrated throughout the notebook:\n",
        "\n",
        "1. A **change in the prevalences of subpopulations** (i.e. news topics) relative to their prevalences in the training data. Contrary to traditional drift detection approaches, **the context-aware detector does not flag drift** as this change in frequency of news topics is permissible given the context provided (e.g. more political news articles around elections).\n",
        "2. A **change in the underlying distribution of one or more subpopulations** takes place. While we allow changes in the prevalence of the subpopulations accounted for by the context variable, we do not allow changes of the subpopulations themselves. Let's assume that a newspaper usually has a certain tone (e.g. more conservative) when it comes to politics. If this tone changes (to less conservative) around elections (increased frequency of political news articles), then we want to **flag it as drift** since the change cannot be attributed to the context given to the detector.\n",
        "3. A **change in the distribution as we observe a previously unseen news topic**. A newspaper might for instance add a classified ads section, which was not present in the reference data.\n",
        "\n",
        "Under setting 1. we want our detector to be **well-calibrated** (a controlled False Positive Rate (FPR) and more generally a p-value which is uniformly distributed between 0 and 1) while under settings 2. and 3. we want our detector to be **powerful** and flag the drift. Lastly, we show how the detector can help you to **understand the connection between the reference and test data distributions** better.\n",
        "\n",
        "## Data\n",
        "\n",
        "We use the [20 newsgroup dataset](https://scikit-learn.org/0.19/datasets/twenty_newsgroups.html) which contains about 18,000 newsgroups post across 20 topics, including politics, science sports or religion.\n",
        "\n",
        "## Requirements\n",
        "\n",
        "The notebook requires the `umap-learn`, `torch`, `sentence-transformers`, `statsmodels`, `seaborn` and `datasets` packages to be installed, which can be done via `pip`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9012b12f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9012b12f",
        "outputId": "351f42d0-0fce-4b5f-8bb5-49563d5cb757"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting umap-learn\n",
            "  Downloading umap-learn-0.5.2.tar.gz (86 kB)\n",
            "\u001b[K     |████████████████████████████████| 86 kB 2.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (1.10.0+cu111)\n",
            "Collecting sentence-transformers\n",
            "  Downloading sentence-transformers-2.2.0.tar.gz (79 kB)\n",
            "\u001b[K     |████████████████████████████████| 79 kB 9.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: statsmodels in /usr/local/lib/python3.7/dist-packages (0.10.2)\n",
            "Requirement already satisfied: seaborn in /usr/local/lib/python3.7/dist-packages (0.11.2)\n",
            "Collecting datasets\n",
            "  Downloading datasets-2.0.0-py3-none-any.whl (325 kB)\n",
            "\u001b[K     |████████████████████████████████| 325 kB 40.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.7/dist-packages (from umap-learn) (1.21.5)\n",
            "Requirement already satisfied: scikit-learn>=0.22 in /usr/local/lib/python3.7/dist-packages (from umap-learn) (1.0.2)\n",
            "Requirement already satisfied: scipy>=1.0 in /usr/local/lib/python3.7/dist-packages (from umap-learn) (1.4.1)\n",
            "Requirement already satisfied: numba>=0.49 in /usr/local/lib/python3.7/dist-packages (from umap-learn) (0.51.2)\n",
            "Collecting pynndescent>=0.5\n",
            "  Downloading pynndescent-0.5.6.tar.gz (1.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1 MB 41.1 MB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from umap-learn) (4.63.0)\n",
            "Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba>=0.49->umap-learn) (0.34.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from numba>=0.49->umap-learn) (57.4.0)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from pynndescent>=0.5->umap-learn) (1.1.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn>=0.22->umap-learn) (3.1.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch) (3.10.0.2)\n",
            "Collecting transformers<5.0.0,>=4.6.0\n",
            "  Downloading transformers-4.18.0-py3-none-any.whl (4.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.0 MB 43.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (0.11.1+cu111)\n",
            "Requirement already satisfied: nltk in /usr/local/lib/python3.7/dist-packages (from sentence-transformers) (3.2.5)\n",
            "Collecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.96-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.2 MB 47.8 MB/s \n",
            "\u001b[?25hCollecting huggingface-hub\n",
            "  Downloading huggingface_hub-0.5.1-py3-none-any.whl (77 kB)\n",
            "\u001b[K     |████████████████████████████████| 77 kB 3.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (21.3)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (3.6.0)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (4.11.3)\n",
            "Collecting pyyaml>=5.1\n",
            "  Downloading PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (596 kB)\n",
            "\u001b[K     |████████████████████████████████| 596 kB 45.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (2.23.0)\n",
            "Collecting sacremoses\n",
            "  Downloading sacremoses-0.0.49-py3-none-any.whl (895 kB)\n",
            "\u001b[K     |████████████████████████████████| 895 kB 17.6 MB/s \n",
            "\u001b[?25hCollecting tokenizers!=0.11.3,<0.13,>=0.11.1\n",
            "  Downloading tokenizers-0.11.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (6.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 6.5 MB 45.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.6.0->sentence-transformers) (2019.12.20)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=20.0->transformers<5.0.0,>=4.6.0->sentence-transformers) (3.0.7)\n",
            "Requirement already satisfied: pandas>=0.19 in /usr/local/lib/python3.7/dist-packages (from statsmodels) (1.3.5)\n",
            "Requirement already satisfied: patsy>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from statsmodels) (0.5.2)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=0.19->statsmodels) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=0.19->statsmodels) (2018.9)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from patsy>=0.4.0->statsmodels) (1.15.0)\n",
            "Requirement already satisfied: matplotlib>=2.2 in /usr/local/lib/python3.7/dist-packages (from seaborn) (3.2.2)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=2.2->seaborn) (0.11.0)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=2.2->seaborn) (1.4.0)\n",
            "Collecting responses<0.19\n",
            "  Downloading responses-0.18.0-py3-none-any.whl (38 kB)\n",
            "Collecting fsspec[http]>=2021.05.0\n",
            "  Downloading fsspec-2022.3.0-py3-none-any.whl (136 kB)\n",
            "\u001b[K     |████████████████████████████████| 136 kB 24.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: multiprocess in /usr/local/lib/python3.7/dist-packages (from datasets) (0.70.12.2)\n",
            "Requirement already satisfied: dill in /usr/local/lib/python3.7/dist-packages (from datasets) (0.3.4)\n",
            "Collecting xxhash\n",
            "  Downloading xxhash-3.0.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (212 kB)\n",
            "\u001b[K     |████████████████████████████████| 212 kB 50.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: pyarrow>=5.0.0 in /usr/local/lib/python3.7/dist-packages (from datasets) (6.0.1)\n",
            "Collecting aiohttp\n",
            "  Downloading aiohttp-3.8.1-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (1.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1 MB 45.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers<5.0.0,>=4.6.0->sentence-transformers) (2021.10.8)\n",
            "Collecting urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1\n",
            "  Downloading urllib3-1.25.11-py2.py3-none-any.whl (127 kB)\n",
            "\u001b[K     |████████████████████████████████| 127 kB 10.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->datasets) (2.0.12)\n",
            "Collecting aiosignal>=1.1.2\n",
            "  Downloading aiosignal-1.2.0-py3-none-any.whl (8.2 kB)\n",
            "Collecting yarl<2.0,>=1.0\n",
            "  Downloading yarl-1.7.2-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (271 kB)\n",
            "\u001b[K     |████████████████████████████████| 271 kB 53.8 MB/s \n",
            "\u001b[?25hCollecting asynctest==0.13.0\n",
            "  Downloading asynctest-0.13.0-py3-none-any.whl (26 kB)\n",
            "Collecting frozenlist>=1.1.1\n",
            "  Downloading frozenlist-1.3.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (144 kB)\n",
            "\u001b[K     |████████████████████████████████| 144 kB 56.4 MB/s \n",
            "\u001b[?25hCollecting async-timeout<5.0,>=4.0.0a3\n",
            "  Downloading async_timeout-4.0.2-py3-none-any.whl (5.8 kB)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->datasets) (21.4.0)\n",
            "Collecting multidict<7.0,>=4.5\n",
            "  Downloading multidict-6.0.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (94 kB)\n",
            "\u001b[K     |████████████████████████████████| 94 kB 3.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->transformers<5.0.0,>=4.6.0->sentence-transformers) (3.7.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers<5.0.0,>=4.6.0->sentence-transformers) (7.1.2)\n",
            "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision->sentence-transformers) (7.1.2)\n",
            "Building wheels for collected packages: umap-learn, pynndescent, sentence-transformers\n",
            "  Building wheel for umap-learn (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for umap-learn: filename=umap_learn-0.5.2-py3-none-any.whl size=82708 sha256=5f67fc8c3edd8cb93f592dc6b6d65ba190ce32697b736edcba83cb1a1adefda5\n",
            "  Stored in directory: /root/.cache/pip/wheels/84/1b/c6/aaf68a748122632967cef4dffef68224eb16798b6793257d82\n",
            "  Building wheel for pynndescent (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pynndescent: filename=pynndescent-0.5.6-py3-none-any.whl size=53943 sha256=0546742a428c6fb34b998bb22c981e21c14b04a347a514a5a2d0a673908fb31e\n",
            "  Stored in directory: /root/.cache/pip/wheels/03/f1/56/f80d72741e400345b5a5b50ec3d929aca581bf45e0225d5c50\n",
            "  Building wheel for sentence-transformers (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for sentence-transformers: filename=sentence_transformers-2.2.0-py3-none-any.whl size=120747 sha256=bbf4689cdc800c325868cc41173bcd1c89ddeb1792468dba0146b89865328079\n",
            "  Stored in directory: /root/.cache/pip/wheels/83/c0/df/b6873ab7aac3f2465aa9144b6b4c41c4391cfecc027c8b07e7\n",
            "Successfully built umap-learn pynndescent sentence-transformers\n",
            "Installing collected packages: urllib3, multidict, frozenlist, yarl, pyyaml, asynctest, async-timeout, aiosignal, tokenizers, sacremoses, huggingface-hub, fsspec, aiohttp, xxhash, transformers, sentencepiece, responses, pynndescent, umap-learn, sentence-transformers, datasets\n",
            "  Attempting uninstall: urllib3\n",
            "    Found existing installation: urllib3 1.24.3\n",
            "    Uninstalling urllib3-1.24.3:\n",
            "      Successfully uninstalled urllib3-1.24.3\n",
            "  Attempting uninstall: pyyaml\n",
            "    Found existing installation: PyYAML 3.13\n",
            "    Uninstalling PyYAML-3.13:\n",
            "      Successfully uninstalled PyYAML-3.13\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\u001b[0m\n",
            "Successfully installed aiohttp-3.8.1 aiosignal-1.2.0 async-timeout-4.0.2 asynctest-0.13.0 datasets-2.0.0 frozenlist-1.3.0 fsspec-2022.3.0 huggingface-hub-0.5.1 multidict-6.0.2 pynndescent-0.5.6 pyyaml-6.0 responses-0.18.0 sacremoses-0.0.49 sentence-transformers-2.2.0 sentencepiece-0.1.96 tokenizers-0.11.6 transformers-4.18.0 umap-learn-0.5.2 urllib3-1.25.11 xxhash-3.0.0 yarl-1.7.2\n"
          ]
        }
      ],
      "source": [
        "!pip install umap-learn torch sentence-transformers statsmodels seaborn datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7764fddf",
      "metadata": {
        "id": "7764fddf"
      },
      "source": [
        "Before we start let's fix the random seeds for reproducibility:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "009c5794",
      "metadata": {
        "id": "009c5794"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import pandas as pd\n",
        "def set_seed(seed: int) -> None:\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    np.random.seed(seed)\n",
        "\n",
        "set_seed(2022)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "450143d0",
      "metadata": {
        "id": "450143d0"
      },
      "source": [
        "## Load data\n",
        "\n",
        "First we load the data, show which classes (news topics) are present and what an instance looks like."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c95ffe02",
      "metadata": {
        "id": "c95ffe02"
      },
      "outputs": [],
      "source": [
        "from sklearn.datasets import fetch_20newsgroups\n",
        "import re\n",
        "import string\n",
        "def twenty_newsgroup_to_csv():\n",
        "    newsgroups_train = fetch_20newsgroups(subset='all', remove=('headers', 'footers', 'quotes'))\n",
        "    df = pd.DataFrame([newsgroups_train.data, newsgroups_train.target.tolist()]).T\n",
        "    df.columns = ['text', 'target']\n",
        "    targets = pd.DataFrame( newsgroups_train.target_names)\n",
        "    targets.columns=['title']\n",
        "    dataframe = pd.merge(df, targets, left_on='target', right_index=True)\n",
        "    dataframe.dropna(inplace=True)\n",
        "    return dataframe\n",
        "\n",
        "df= twenty_newsgroup_to_csv()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "848fc91e",
      "metadata": {
        "id": "848fc91e"
      },
      "source": [
        "Let's take a look at an instance from the dataset:"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1b72e248",
      "metadata": {
        "id": "1b72e248"
      },
      "source": [
        "## Define models and train a classifier\n",
        "\n",
        "We embed the news posts using [SentenceTransformers](https://www.sbert.net/index.html) pre-trained embeddings and optionally add a dimensionality reduction step with [UMAP](https://umap-learn.readthedocs.io/en/latest/). UMAP also allows to leverage reference data labels.\n",
        "\n",
        "We define respectively a generic clustering model using UMAP, a model to embed the text input using pre-trained *SentenceTransformers* embeddings, a text classifier and a utility function to place the data on the right device."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cb03109a",
      "metadata": {
        "id": "cb03109a"
      },
      "outputs": [],
      "source": [
        "from sentence_transformers import SentenceTransformer\n",
        "import torch.nn as nn\n",
        "import umap\n",
        "\n",
        "\n",
        "class UMAPModel:\n",
        "    def __init__(\n",
        "        self, \n",
        "        n_neighbors: int = 10,\n",
        "        n_components: int = 2,\n",
        "        metric: str = 'euclidean',\n",
        "        min_dist: float = .1,\n",
        "        **kwargs: dict\n",
        "    ) -> None:\n",
        "        super().__init__()\n",
        "        kwargs = kwargs if isinstance(kwargs, dict) else dict()\n",
        "        kwargs.update(\n",
        "            n_neighbors=n_neighbors,\n",
        "            n_components=n_components,\n",
        "            metric=metric,\n",
        "            min_dist=min_dist\n",
        "        )\n",
        "        self.model = umap.UMAP(**kwargs)\n",
        "    \n",
        "    def fit(self, x: np.ndarray, y: np.ndarray = None) -> None:\n",
        "        \"\"\" Fit UMAP embedding. A combination of labeled and unlabeled data\n",
        "        can be passed. Unlabeled instances are equal to -1. \"\"\"\n",
        "        self.model.fit(x, y=y)\n",
        "    \n",
        "    def predict(self, x: np.ndarray) -> np.ndarray:\n",
        "        \"\"\" Transform the input x to the embedding space. \"\"\"\n",
        "        return self.model.transform(x)\n",
        "\n",
        "\n",
        "class EmbeddingModel:\n",
        "    def __init__(\n",
        "        self,\n",
        "        model_name: str = 'paraphrase-MiniLM-L6-v2',  # https://www.sbert.net/docs/pretrained_models.html\n",
        "        max_seq_length: int = 200,\n",
        "        batch_size: int = 32,\n",
        "        device: torch.device = None\n",
        "    ) -> None:\n",
        "        if not isinstance(device, torch.device):\n",
        "            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "        self.encode_text = SentenceTransformer(model_name).to(device)\n",
        "        self.encode_text.max_seq_length = max_seq_length\n",
        "        self.batch_size = batch_size\n",
        "    \n",
        "    def __call__(self, x: np.ndarray) -> np.ndarray:\n",
        "        return self.encode_text.encode(x, convert_to_numpy=True, batch_size=self.batch_size,\n",
        "                                       show_progress_bar=False)\n",
        "\n",
        "    \n",
        "class Classifier(nn.Module):\n",
        "    def __init__(\n",
        "        self, \n",
        "        model_name: str = 'paraphrase-MiniLM-L6-v2',\n",
        "        max_seq_length: int = 200,\n",
        "        n_classes: int = 8\n",
        "    ) -> None:\n",
        "        \"\"\" Text classification model. Note that we do not train the embedding backbone. \"\"\"\n",
        "        super().__init__()\n",
        "        self.encode_text = SentenceTransformer(model_name)\n",
        "        self.encode_text.max_seq_length = max_seq_length\n",
        "        for param in self.encode_text.parameters():\n",
        "            param.requires_grad = False\n",
        "        self.head = nn.Sequential(nn.Linear(384, 256), nn.LeakyReLU(.1), nn.Dropout(.5), nn.Linear(256, 20))\n",
        "        \n",
        "    def forward(self, tokens) -> torch.Tensor:\n",
        "        return self.head(self.encode_text(tokens)['sentence_embedding'])\n",
        "\n",
        "    \n",
        "def batch_to_device(batch: dict, target_device: torch.device):\n",
        "    \"\"\" Send a pytorch batch to a device (CPU/GPU). \"\"\"\n",
        "    for key in batch:\n",
        "        if isinstance(batch[key], torch.Tensor):\n",
        "            batch[key] = batch[key].to(target_device)\n",
        "    return batch"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7e1079c6",
      "metadata": {
        "id": "7e1079c6"
      },
      "source": [
        "First we train a classifier on a small subset of the data. The aim of the classifier is to predict the news topic of each instance. Below we define a few simple training and evaluation functions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "205369d6",
      "metadata": {
        "id": "205369d6"
      },
      "outputs": [],
      "source": [
        "def train_model(model, loader, epochs=3, lr=1e-3):\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "    for epoch in range(epochs):\n",
        "        for x, y in tqdm(loader):\n",
        "            tokens, y = tokenize(x), y.to(device)\n",
        "            y_hat = clf(tokens)\n",
        "            optimizer.zero_grad()\n",
        "            loss = criterion(y_hat, y)\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            \n",
        "def eval_model(model, loader, verbose=1):\n",
        "    model.eval()\n",
        "    logits, labels = [], []\n",
        "    with torch.no_grad():\n",
        "        if verbose == 1:\n",
        "            loader = tqdm(loader)\n",
        "        for x, y in loader:\n",
        "            tokens = tokenize(x)\n",
        "            y_hat = model(tokens)\n",
        "            logits += [y_hat.cpu().numpy()]\n",
        "            labels += [y.cpu().numpy()]\n",
        "    logits = np.concatenate(logits, 0)\n",
        "    preds = np.argmax(logits, 1)\n",
        "    labels = np.concatenate(labels, 0)\n",
        "    if verbose == 1:\n",
        "        accuracy = (preds == labels).mean()\n",
        "        print(f'Accuracy: {accuracy:.3f}')\n",
        "    return logits, preds"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dic={'comp.graphics': 1,\n",
        " 'comp.os.ms-windows.misc': 2,\n",
        " 'comp.sys.ibm.pc.hardware': 3,\n",
        " 'misc.forsale': 6,\n",
        " 'rec.autos': 7,\n",
        " 'rec.motorcycles': 8,\n",
        " 'talk.politics.misc': 18,\n",
        " 'talk.politics.guns': 16}"
      ],
      "metadata": {
        "id": "s2VcprPLPFUH"
      },
      "id": "s2VcprPLPFUH",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_map =  ['comp.graphics',\n",
        " 'comp.os.ms-windows.misc',\n",
        " 'comp.sys.ibm.pc.hardware',\n",
        " 'misc.forsale',\n",
        " 'rec.autos',\n",
        " 'rec.motorcycles',\n",
        " 'talk.politics.misc',\n",
        " 'talk.politics.guns']\n",
        "train = df[df.title.isin(train_map)]\n",
        "train_map_ = mapping ={1: 0,\n",
        " 2: 1,\n",
        " 3: 2,\n",
        " 6: 3,\n",
        " 7: 4,\n",
        " 8: 5,\n",
        " 18: 6,\n",
        " 16: 7}\n",
        "\n",
        "train['target'] = train['target'].map(train_map_)\n",
        "train = train.reset_index(drop=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "22PGNKPXO5my",
        "outputId": "dba68a31-5bb0-4f88-86e0-33976cb69d97"
      },
      "id": "22PGNKPXO5my",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:19: SettingWithCopyWarning: \n",
            "A value is trying to be set on a copy of a slice from a DataFrame.\n",
            "Try using .loc[row_indexer,col_indexer] = value instead\n",
            "\n",
            "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train.text.isnull().sum()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tiQkq9CO45oe",
        "outputId": "3a7df652-da0c-482d-f6ec-3b2d7c9db7b8"
      },
      "id": "tiQkq9CO45oe",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "exp2_cats = ['sci.crypt','sci.med','sci.electronics','sci.space','talk.religion.misc','alt.atheism','soc.religion.christian','rec.sport.baseball']"
      ],
      "metadata": {
        "id": "HXy1rCBacqvl"
      },
      "id": "HXy1rCBacqvl",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        ""
      ],
      "metadata": {
        "id": "S_kgATEPc6xB"
      },
      "id": "S_kgATEPc6xB"
    },
    {
      "cell_type": "code",
      "source": [
        "exp2_cats = ['sci.crypt','sci.med','sci.electronics','sci.space','talk.religion.misc','alt.atheism','soc.religion.christian','rec.sport.baseball']\n",
        "map_ = {0:0, \n",
        "        9:1,\n",
        "        11:2,\n",
        "        12:3,\n",
        "        13:4,\n",
        "        14:5,\n",
        "        15:6,\n",
        "        19:7}\n",
        "test = df[df.title.isin(exp2_cats)]\n",
        "test['target'] = test['target'].map(map_)\n",
        "map__ = {'soc.religion.christian':'comp.graphics',\n",
        "         'sci.electronics':'talk.politics.misc',\n",
        "         'talk.religion.misc': 'rec.motorcycles',\n",
        "         'sci.med':'comp.sys.ibm.pc.hardware',\n",
        "         'rec.sport.baseball':'misc.forsale',\n",
        "         'sci.crypt':'talk.politics.guns',\n",
        "         'alt.atheism':'rec.autos',\n",
        "         'sci.space':'comp.os.ms-windows.misc'}\n",
        "test['title'] = test['title'].map(map__)\n",
        "test = test.reset_index(drop=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FI8d4bJXQhGE",
        "outputId": "aef5cfeb-388a-4c20-e775-b50f27d63a79"
      },
      "id": "FI8d4bJXQhGE",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:11: SettingWithCopyWarning: \n",
            "A value is trying to be set on a copy of a slice from a DataFrame.\n",
            "Try using .loc[row_indexer,col_indexer] = value instead\n",
            "\n",
            "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
            "  # This is added back by InteractiveShellApp.init_path()\n",
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:20: SettingWithCopyWarning: \n",
            "A value is trying to be set on a copy of a slice from a DataFrame.\n",
            "Try using .loc[row_indexer,col_indexer] = value instead\n",
            "\n",
            "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test.text.isnull().sum()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "01KAJkuVl6uK",
        "outputId": "91ab7e69-db7c-4310-e265-4b080f00fedf"
      },
      "id": "01KAJkuVl6uK",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train.target.value_counts()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ibMoeQqI0eTE",
        "outputId": "11657365-3fdd-49e0-ee0a-e4adadd8c158"
      },
      "id": "ibMoeQqI0eTE",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5    996\n",
              "4    990\n",
              "1    985\n",
              "2    982\n",
              "3    975\n",
              "0    973\n",
              "7    910\n",
              "6    775\n",
              "Name: target, dtype: int64"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        ""
      ],
      "metadata": {
        "id": "U3Fm8R6bc2y3"
      },
      "id": "U3Fm8R6bc2y3"
    },
    {
      "cell_type": "markdown",
      "id": "788fcfd4",
      "metadata": {
        "id": "788fcfd4"
      },
      "source": [
        "We now split the data in 2 sets. The first set (`x_train`) we will use to train our text classifier, and the second set (`x_drift`) is held out to test our drift detector on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8a81218b",
      "metadata": {
        "id": "8a81218b"
      },
      "outputs": [],
      "source": [
        "idx_train = list(range(0, 5001))\n",
        "idx_keep = list(range(0,2001))\n",
        "# data used for model training\n",
        "x_train, y_train = [train['text'][_] for _ in idx_train], train.target[idx_train]\n",
        "# data used for drift detection\n",
        "x_drift, y_drift = [test['text'][_] for _ in idx_keep], test.target[idx_keep]\n",
        "n_drift = len(x_drift)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "22dd64e0",
      "metadata": {
        "id": "22dd64e0"
      },
      "source": [
        "Let's train our classifier. The classifier consists of a simple MLP head on top of a pre-trained SentenceTransformer model as the backbone. The SentenceTransformer remains frozen during training and only the MLP head is finetuned."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install alibi_detect"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zfb88MnrS_mi",
        "outputId": "85caeb42-b087-4bd2-ec1d-9647e6f1abaf"
      },
      "id": "Zfb88MnrS_mi",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: alibi_detect in /usr/local/lib/python3.7/dist-packages (0.9.1)\n",
            "Requirement already satisfied: scipy<2.0.0,>=1.3.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (1.4.1)\n",
            "Requirement already satisfied: matplotlib<4.0.0,>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (3.2.2)\n",
            "Requirement already satisfied: dill<0.4.0,>=0.3.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (0.3.4)\n",
            "Requirement already satisfied: tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (2.8.0)\n",
            "Requirement already satisfied: Pillow<10.0.0,>=5.4.1 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (7.1.2)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (2.23.0)\n",
            "Requirement already satisfied: numba!=0.54.0,<0.56.0,>=0.50.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (0.51.2)\n",
            "Requirement already satisfied: pandas<2.0.0,>=0.23.3 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (1.3.5)\n",
            "Requirement already satisfied: tensorflow-probability<0.13.0,>=0.8.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (0.12.2)\n",
            "Requirement already satisfied: scikit-image!=0.17.1,<0.20,>=0.14.2 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (0.18.3)\n",
            "Requirement already satisfied: scikit-learn<1.1.0,>=0.20.2 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (1.0.2)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.16.2 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (1.21.5)\n",
            "Requirement already satisfied: transformers<5.0.0,>=4.0.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (4.18.0)\n",
            "Requirement already satisfied: opencv-python<5.0.0,>=3.2.0 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (4.1.2.30)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.28.1 in /usr/local/lib/python3.7/dist-packages (from alibi_detect) (4.63.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib<4.0.0,>=3.0.0->alibi_detect) (2.8.2)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib<4.0.0,>=3.0.0->alibi_detect) (3.0.7)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib<4.0.0,>=3.0.0->alibi_detect) (1.4.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib<4.0.0,>=3.0.0->alibi_detect) (0.11.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from kiwisolver>=1.0.1->matplotlib<4.0.0,>=3.0.0->alibi_detect) (3.10.0.2)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from numba!=0.54.0,<0.56.0,>=0.50.0->alibi_detect) (57.4.0)\n",
            "Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba!=0.54.0,<0.56.0,>=0.50.0->alibi_detect) (0.34.0)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas<2.0.0,>=0.23.3->alibi_detect) (2018.9)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.1->matplotlib<4.0.0,>=3.0.0->alibi_detect) (1.15.0)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.21.0->alibi_detect) (1.25.11)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.21.0->alibi_detect) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.21.0->alibi_detect) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.21.0->alibi_detect) (2021.10.8)\n",
            "Requirement already satisfied: PyWavelets>=1.1.1 in /usr/local/lib/python3.7/dist-packages (from scikit-image!=0.17.1,<0.20,>=0.14.2->alibi_detect) (1.3.0)\n",
            "Requirement already satisfied: imageio>=2.3.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image!=0.17.1,<0.20,>=0.14.2->alibi_detect) (2.4.1)\n",
            "Requirement already satisfied: tifffile>=2019.7.26 in /usr/local/lib/python3.7/dist-packages (from scikit-image!=0.17.1,<0.20,>=0.14.2->alibi_detect) (2021.11.2)\n",
            "Requirement already satisfied: networkx>=2.0 in /usr/local/lib/python3.7/dist-packages (from scikit-image!=0.17.1,<0.20,>=0.14.2->alibi_detect) (2.6.3)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn<1.1.0,>=0.20.2->alibi_detect) (3.1.0)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn<1.1.0,>=0.20.2->alibi_detect) (1.1.0)\n",
            "Requirement already satisfied: tf-estimator-nightly==2.8.0.dev2021122109 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (2.8.0.dev2021122109)\n",
            "Requirement already satisfied: tensorboard<2.9,>=2.8 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (2.8.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.1.0)\n",
            "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (3.17.3)\n",
            "Requirement already satisfied: libclang>=9.0.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (13.0.0)\n",
            "Requirement already satisfied: gast>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.5.3)\n",
            "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.24.0)\n",
            "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.44.0)\n",
            "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (3.1.0)\n",
            "Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.14.0)\n",
            "Requirement already satisfied: flatbuffers>=1.12 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (2.0)\n",
            "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.6.3)\n",
            "Requirement already satisfied: keras<2.9,>=2.8.0rc0 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (2.8.0)\n",
            "Requirement already satisfied: absl-py>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.0.0)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (3.3.0)\n",
            "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.2.0)\n",
            "Requirement already satisfied: keras-preprocessing>=1.1.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.1.2)\n",
            "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.7/dist-packages (from astunparse>=1.6.0->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.37.1)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py>=2.9.0->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.5.2)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.4.6)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.6.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.8.1)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.35.0)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.0.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (3.3.6)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (4.8)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (4.2.4)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.2.8)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (1.3.1)\n",
            "Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (4.11.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (3.7.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (0.4.8)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.9,>=2.8->tensorflow!=2.6.0,!=2.6.1,<2.9.0,>=2.2.0->alibi_detect) (3.2.0)\n",
            "Requirement already satisfied: cloudpickle>=1.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow-probability<0.13.0,>=0.8.0->alibi_detect) (1.3.0)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.7/dist-packages (from tensorflow-probability<0.13.0,>=0.8.0->alibi_detect) (0.1.6)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from tensorflow-probability<0.13.0,>=0.8.0->alibi_detect) (4.4.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.0.0->alibi_detect) (6.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.0.0->alibi_detect) (3.6.0)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.13,>=0.11.1 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.0.0->alibi_detect) (0.11.6)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.0.0->alibi_detect) (21.3)\n",
            "Requirement already satisfied: sacremoses in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.0.0->alibi_detect) (0.0.49)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.0.0->alibi_detect) (2019.12.20)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.1.0 in /usr/local/lib/python3.7/dist-packages (from transformers<5.0.0,>=4.0.0->alibi_detect) (0.5.1)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers<5.0.0,>=4.0.0->alibi_detect) (7.1.2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0facda4e",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0facda4e",
        "outputId": "4c19a2c2-8c92-4ff4-9de5-533410efeb75"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 157/157 [00:23<00:00,  6.57it/s]\n",
            "100%|██████████| 157/157 [00:23<00:00,  6.66it/s]\n",
            "100%|██████████| 157/157 [00:23<00:00,  6.71it/s]\n",
            "100%|██████████| 157/157 [00:23<00:00,  6.72it/s]\n",
            "100%|██████████| 157/157 [00:23<00:00,  6.64it/s]\n",
            "100%|██████████| 157/157 [00:26<00:00,  6.00it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Accuracy: 0.865\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 63/63 [00:10<00:00,  5.93it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Accuracy: 0.288\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "from alibi_detect.utils.pytorch.data import TorchDataset\n",
        "from torch.utils.data import DataLoader\n",
        "from tqdm import tqdm\n",
        "from typing import Dict, List\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "clf = Classifier().to(device)\n",
        "train_loader = DataLoader(TorchDataset(x_train, y_train), batch_size=32, shuffle=True)\n",
        "drift_loader = DataLoader(TorchDataset(x_drift, y_drift), batch_size=32, shuffle=False)\n",
        "\n",
        "def tokenize(x: List[str]) -> Dict[str, torch.Tensor]:\n",
        "    tokens = clf.encode_text.tokenize(x)\n",
        "    return batch_to_device(tokens, device)\n",
        "\n",
        "train_model(clf, train_loader, epochs=5)\n",
        "clf.eval()\n",
        "_, _ = eval_model(clf, train_loader)\n",
        "_, _ = eval_model(clf, drift_loader)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8bab6886",
      "metadata": {
        "id": "8bab6886"
      },
      "source": [
        "## Detector calibration under no change\n",
        "\n",
        "We start with an example where no drift occurs and the reference and test data are both sampled randomly from all news topics. Under this scenario, we expect no drift to be detected by either a *normal* MMD detector or by the context-aware MMD detector.\n",
        "\n",
        "First we define some helper functions. The first one visualises the clustered text data while the second function samples disjoint reference and test sets with a specified number of instances per class (i.e. per news topic)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "10a1a38b",
      "metadata": {
        "id": "10a1a38b"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "def plot_clusters(x: np.ndarray, y: np.ndarray, classes: list, title: str = None) -> None:\n",
        "    fig, ax = plt.subplots(1, figsize=(14, 10))\n",
        "    plt.scatter(*x.T, s=0.3, c=y, cmap='Spectral', alpha=1.0)\n",
        "    plt.setp(ax, xticks=[], yticks=[])\n",
        "    nc = len(classes)\n",
        "    cbar = plt.colorbar(boundaries=np.arange(nc+1)-0.5)\n",
        "    cbar.set_ticks(np.arange(nc))\n",
        "    cbar.set_ticklabels(classes)\n",
        "    if title:\n",
        "        plt.title(title);\n",
        "        \n",
        "\n",
        "def split_data(x, y, n_ref_c, n_test_c, seed=None, y2=None, return_idx=False):\n",
        "    \n",
        "    if seed:\n",
        "        np.random.seed(seed)\n",
        "    \n",
        "    # split data by class\n",
        "    n_c = len(np.unique(y))\n",
        "    idx_c = {_: np.where(y == _)[0] for _ in range(n_c)}\n",
        "    \n",
        "    # convert nb instances per class to a list if needed\n",
        "    n_ref_c = [n_ref_c] * n_c if isinstance(n_ref_c, int) else n_ref_c\n",
        "    n_test_c = [n_test_c] * n_c if isinstance(n_test_c, int) else n_test_c\n",
        "    \n",
        "    # sample reference, test and held out data\n",
        "    idx_ref, idx_test, idx_held = [], [], []\n",
        "    for _ in range(n_c):\n",
        "        idx = np.random.choice(idx_c[_], size=len(idx_c[_]), replace=False)\n",
        "        idx_ref.append(idx[:n_ref_c[_]])\n",
        "        idx_test.append(idx[n_ref_c[_]:n_ref_c[_] + n_test_c[_]])\n",
        "        idx_held.append(idx[n_ref_c[_] + n_test_c[_]:])\n",
        "    idx_ref = np.concatenate(idx_ref)\n",
        "    idx_test = np.concatenate(idx_test)\n",
        "    idx_held = np.concatenate(idx_held)\n",
        "    x_ref, y_ref = [x[_] for _ in idx_ref], y[idx_ref]\n",
        "    x_test, y_test = [x[_] for _ in idx_test], y[idx_test]\n",
        "    x_held, y_held = [x[_] for _ in idx_held], y[idx_held]\n",
        "    if y2 is not None:\n",
        "        y_ref2, y_test2, y_held2 = y2[idx_ref], y2[idx_test], y2[idx_held]\n",
        "        return (x_ref, y_ref, y_ref2), (x_test, y_test, y_test2), (x_held, y_held, y_held2)\n",
        "    elif not return_idx:\n",
        "        return (x_ref, y_ref), (x_test, y_test), (x_held, y_held)\n",
        "    else:\n",
        "        return idx_ref, idx_test, idx_held"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "id": "n6hb6gcN7fmo",
        "outputId": "1110fd0e-949b-4709-a523-292671af90b3"
      },
      "id": "n6hb6gcN7fmo",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                                text  target  \\\n",
              "0  \\n\\nBack in high school I worked as a lab assi...       3   \n",
              "1  Many thanks to those who replied to my appeal ...       3   \n",
              "2  .........\\nI, some years ago, almost became a ...       3   \n",
              "3  \\nOh dear, time for me to try to remember my c...       3   \n",
              "4  \\n\\nWHen trying to choose a resistor with a to...       3   \n",
              "\n",
              "                title  \n",
              "0  talk.politics.misc  \n",
              "1  talk.politics.misc  \n",
              "2  talk.politics.misc  \n",
              "3  talk.politics.misc  \n",
              "4  talk.politics.misc  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-a86ca595-a5d9-4647-8068-466a6cecea08\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>text</th>\n",
              "      <th>target</th>\n",
              "      <th>title</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>\\n\\nBack in high school I worked as a lab assi...</td>\n",
              "      <td>3</td>\n",
              "      <td>talk.politics.misc</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>Many thanks to those who replied to my appeal ...</td>\n",
              "      <td>3</td>\n",
              "      <td>talk.politics.misc</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>.........\\nI, some years ago, almost became a ...</td>\n",
              "      <td>3</td>\n",
              "      <td>talk.politics.misc</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>\\nOh dear, time for me to try to remember my c...</td>\n",
              "      <td>3</td>\n",
              "      <td>talk.politics.misc</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>\\n\\nWHen trying to choose a resistor with a to...</td>\n",
              "      <td>3</td>\n",
              "      <td>talk.politics.misc</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-a86ca595-a5d9-4647-8068-466a6cecea08')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-a86ca595-a5d9-4647-8068-466a6cecea08 button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-a86ca595-a5d9-4647-8068-466a6cecea08');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "id": "0h6mFmZH7hjg",
        "outputId": "622befd8-fecf-4198-f56c-228c50e3b9e7"
      },
      "id": "0h6mFmZH7hjg",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                                text  target  \\\n",
              "0  My brother is in the market for a high-perform...       2   \n",
              "1  \\nThink!\\n\\nIt's the SCSI card doing the DMA t...       2   \n",
              "2  I need some help with a multi port serial boar...       2   \n",
              "3  [discussing the use of IRQ 7]\\n\\n\\n\\nTo put it...       2   \n",
              "4  This might be a silly question but I have to a...       2   \n",
              "\n",
              "                      title  \n",
              "0  comp.sys.ibm.pc.hardware  \n",
              "1  comp.sys.ibm.pc.hardware  \n",
              "2  comp.sys.ibm.pc.hardware  \n",
              "3  comp.sys.ibm.pc.hardware  \n",
              "4  comp.sys.ibm.pc.hardware  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-cb9805e5-d1ce-47a4-b3a3-b32152ca073b\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>text</th>\n",
              "      <th>target</th>\n",
              "      <th>title</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>My brother is in the market for a high-perform...</td>\n",
              "      <td>2</td>\n",
              "      <td>comp.sys.ibm.pc.hardware</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>\\nThink!\\n\\nIt's the SCSI card doing the DMA t...</td>\n",
              "      <td>2</td>\n",
              "      <td>comp.sys.ibm.pc.hardware</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>I need some help with a multi port serial boar...</td>\n",
              "      <td>2</td>\n",
              "      <td>comp.sys.ibm.pc.hardware</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>[discussing the use of IRQ 7]\\n\\n\\n\\nTo put it...</td>\n",
              "      <td>2</td>\n",
              "      <td>comp.sys.ibm.pc.hardware</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>This might be a silly question but I have to a...</td>\n",
              "      <td>2</td>\n",
              "      <td>comp.sys.ibm.pc.hardware</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-cb9805e5-d1ce-47a4-b3a3-b32152ca073b')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-cb9805e5-d1ce-47a4-b3a3-b32152ca073b button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-cb9805e5-d1ce-47a4-b3a3-b32152ca073b');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "(x_held, y_held)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SfV8IfWzAzb2",
        "outputId": "72c7f06b-70ee-449b-ee2e-fb37dac1bf7b"
      },
      "id": "SfV8IfWzAzb2",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "([], Series([], Name: target, dtype: int64))"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "46650730",
      "metadata": {
        "id": "46650730"
      },
      "outputs": [],
      "source": [
        "# initially assume equal distribution of topics in the reference data\n",
        "n_ref, n_test = 2000, 2000\n",
        "classes = list(set(train.title))\n",
        "n_classes = len(classes)\n",
        "n_ref_c = n_ref // n_classes\n",
        "n_test_c = n_test // n_classes\n",
        "\n",
        "(x_ref, y_ref), (x_test, y_test), (x_held, y_held) = split_data(x_drift, y_drift, n_ref_c, n_test_c)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "len(x_drift), len(y_drift)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "heUj6G1Z7o08",
        "outputId": "0f75e495-95ea-48d6-d54a-7cc1ccedf13a"
      },
      "id": "heUj6G1Z7o08",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(2001, 2001)"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "set(test.target)\n",
        "set(test.title)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fLZjUHQZzFrL",
        "outputId": "17ce31c4-7b8c-487e-9595-dca066b899f1"
      },
      "id": "fLZjUHQZzFrL",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'comp.graphics',\n",
              " 'comp.os.ms-windows.misc',\n",
              " 'comp.sys.ibm.pc.hardware',\n",
              " 'misc.forsale',\n",
              " 'rec.autos',\n",
              " 'rec.motorcycles',\n",
              " 'talk.politics.guns',\n",
              " 'talk.politics.misc'}"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "598061d0",
      "metadata": {
        "id": "598061d0"
      },
      "source": [
        "We first define the embedding model using the pre-trained *SentenceTransformer* embeddings and then embed both the reference and test sets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "22e0c313",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "22e0c313",
        "outputId": "8b6329fe-fbd4-4c1e-e03b-570e253ee54e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Shape of embedded reference and test data: (250, 384) - (139, 384)\n"
          ]
        }
      ],
      "source": [
        "model = EmbeddingModel()\n",
        "emb_ref = model(x_ref)\n",
        "emb_test = model(x_test)\n",
        "print(f'Shape of embedded reference and test data: {emb_ref.shape} - {emb_test.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "463e2e2a",
      "metadata": {
        "id": "463e2e2a"
      },
      "source": [
        "By applying UMAP clustering on the *SentenceTransformer* embeddings, we can visually inspect the various news topic clusters. Note that we fit the clustering model on the held out data first, and then make predictions on the reference and test sets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "aa726b40",
      "metadata": {
        "id": "aa726b40"
      },
      "outputs": [],
      "source": [
        "umap_model = UMAPModel()\n",
        "emb_held = model(x_held)\n",
        "umap_model.fit(emb_held, y=y_held)\n",
        "cluster_ref = umap_model.predict(emb_ref)\n",
        "cluster_test = umap_model.predict(emb_test)\n",
        "plot_clusters(cluster_ref, y_ref, classes, title='Reference data: clustered news topics')\n",
        "plot_clusters(cluster_test, y_test, classes, title='Test data: clustered news topics')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c59984be",
      "metadata": {
        "id": "c59984be"
      },
      "source": [
        "We can visually see that the reference and test set are made up of similar clusters of data, grouped by news topic. As a result, we would not expect drift to be flagged. If the data distribution did not change, we can expect the p-value distribution of our statistical test to be uniformly distributed between 0 and 1. So let's see if this assumption holds.\n",
        "\n",
        "Importantly, first we need to **define our context variable** for the context-aware MMD detector. In our experiments we allow the relative prevalences of subpopulations to vary while the distributions underlying each of the subpopulations remain unchanged. To achieve this we **condition on the prediction probabilities of the classifier we trained earlier to distinguish each of the 20 different news topics**. We can do this because the prediction probabilities can account for the frequency of occurrence of each of the topics (be it imperfectly given our classifier makes the occasional mistake)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "97a1c498",
      "metadata": {
        "id": "97a1c498"
      },
      "outputs": [],
      "source": [
        "from scipy.special import softmax\n",
        "\n",
        "def context(x: List[str], y: np.ndarray):  # y only needed for the data loader\n",
        "    \"\"\" Condition on classifier prediction probabilities. \"\"\"\n",
        "    loader = DataLoader(TorchDataset(x, y), batch_size=32, shuffle=False)\n",
        "    logits = eval_model(clf.eval(), loader, verbose=0)[0]\n",
        "    return softmax(logits, -1)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "875e8e8c",
      "metadata": {
        "id": "875e8e8c"
      },
      "source": [
        "Before we set off our experiments, we embed all the instances in `x_drift` and compute all contexts `c_drift` so we don't have to call our transformer model every single pass in the for loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e3f81d3b",
      "metadata": {
        "code_folding": [],
        "id": "e3f81d3b"
      },
      "outputs": [],
      "source": [
        "emb_drift = model(x_drift)\n",
        "c_drift = context(x_drift, y_drift)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2038968f",
      "metadata": {
        "scrolled": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 398
        },
        "id": "2038968f",
        "outputId": "d02fb09a-836d-45fd-9b4f-908e8590b92f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "  0%|          | 0/50 [00:00<?, ?it/s]\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "ZeroDivisionError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mZeroDivisionError\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-31-365f51966d32>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     14\u001b[0m     \u001b[0;31m# mmd drift detector\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m     \u001b[0mdd_mmd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mMMDDrift\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0memb_ref\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mp_val\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m.05\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn_permutations\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'pytorch'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m     \u001b[0mpreds_mmd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdd_mmd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpredict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0memb_test\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m     \u001b[0mp_vals_mmd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpreds_mmd\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'data'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'p_val'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/alibi_detect/cd/mmd.py\u001b[0m in \u001b[0;36mpredict\u001b[0;34m(self, x, return_p_val, return_distance)\u001b[0m\n\u001b[1;32m    114\u001b[0m         \u001b[0;34m'data'\u001b[0m \u001b[0mcontains\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mdrift\u001b[0m \u001b[0mprediction\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0moptionally\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mp\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mthreshold\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mMMD\u001b[0m \u001b[0mmetric\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m         \"\"\"\n\u001b[0;32m--> 116\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_detector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpredict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreturn_p_val\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreturn_distance\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    117\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    118\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mscore\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTuple\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfloat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/alibi_detect/cd/base.py\u001b[0m in \u001b[0;36mpredict\u001b[0;34m(self, x, return_p_val, return_distance)\u001b[0m\n\u001b[1;32m    571\u001b[0m         \"\"\"\n\u001b[1;32m    572\u001b[0m         \u001b[0;31m# compute drift scores\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 573\u001b[0;31m         \u001b[0mp_val\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdist\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdist_permutations\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscore\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    574\u001b[0m         \u001b[0mdrift_pred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mp_val\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mp_val\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    575\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/alibi_detect/cd/pytorch/mmd.py\u001b[0m in \u001b[0;36mscore\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    125\u001b[0m         \u001b[0mkernel_mat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkernel_matrix\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_ref\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[arg-type]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    126\u001b[0m         \u001b[0mkernel_mat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mkernel_mat\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiag\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkernel_mat\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiag\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# zero diagonal\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 127\u001b[0;31m         \u001b[0mmmd2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmmd2_from_kernel_matrix\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkernel_mat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpermute\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mzero_diag\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    128\u001b[0m         mmd2_permuted = torch.Tensor(\n\u001b[1;32m    129\u001b[0m             \u001b[0;34m[\u001b[0m\u001b[0mmmd2_from_kernel_matrix\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkernel_mat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpermute\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mzero_diag\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mn_permutations\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/alibi_detect/utils/pytorch/distance.py\u001b[0m in \u001b[0;36mmmd2_from_kernel_matrix\u001b[0;34m(kernel_mat, m, permute, zero_diag)\u001b[0m\n\u001b[1;32m    122\u001b[0m         \u001b[0mkernel_mat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mkernel_mat\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    123\u001b[0m     \u001b[0mk_xx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk_yy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk_xy\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mkernel_mat\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_mat\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_mat\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 124\u001b[0;31m     \u001b[0mc_xx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mc_yy\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mm\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mm\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    125\u001b[0m     \u001b[0mmmd2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mc_xx\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mk_xx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mc_yy\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mk_yy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m2.\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mk_xy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    126\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mmmd2\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mZeroDivisionError\u001b[0m: division by zero"
          ]
        }
      ],
      "source": [
        "from alibi_detect.cd import MMDDrift, ContextMMDDrift\n",
        "\n",
        "n_runs = 50  # number of drift detection runs, each with a different reference and test sample\n",
        "\n",
        "p_vals_mmd, p_vals_cad = [], []\n",
        "for _ in tqdm(range(n_runs)):\n",
        "    \n",
        "    # sample data\n",
        "    idx = np.random.choice(n_drift, size=n_drift, replace=False)\n",
        "    idx_ref, idx_test = idx[:n_ref], idx[n_ref:n_ref+n_test]\n",
        "    emb_ref, c_ref = emb_drift[idx_ref], c_drift[idx_ref]\n",
        "    emb_test, c_test = emb_drift[idx_test], c_drift[idx_test]\n",
        "    \n",
        "    # mmd drift detector\n",
        "    dd_mmd = MMDDrift(emb_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_mmd = dd_mmd.predict(emb_test)\n",
        "    p_vals_mmd.append(preds_mmd['data']['p_val'])\n",
        "    \n",
        "    # context-aware mmd drift detector \n",
        "    dd_cad = ContextMMDDrift(emb_ref, c_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_cad = dd_cad.predict(emb_test, c_test)\n",
        "    p_vals_cad.append(preds_cad['data']['p_val'])\n",
        "    \n",
        "p_vals_mmd = np.array(p_vals_mmd)\n",
        "p_vals_cad = np.array(p_vals_cad)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e74d76ac",
      "metadata": {
        "id": "e74d76ac"
      },
      "source": [
        "The below figure of the [Q-Q (Quantile-Quantile) plots](https://en.wikipedia.org/wiki/Q%E2%80%93Q_plot) of a random sample from the uniform distribution *U[0,1]* against the obtained p-values from the vanilla and context-aware MMD detectors illustrate how well both detectors are calibrated. A perfectly calibrated detector should have a Q-Q plot which closely follows the diagonal. Only the middle plot in the grid shows the detector's p-values. The other plots correspond to *n_runs* p-values actually sampled from *U[0,1]* to contextualise how well the central plot follows the diagonal given the limited number of samples.\n",
        "\n",
        "As expected we can see that both the normal MMD and the context-aware MMD detectors are well-calibrated."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9d0a4f6d",
      "metadata": {
        "id": "9d0a4f6d"
      },
      "outputs": [],
      "source": [
        "import statsmodels.api as sm\n",
        "from scipy.stats import uniform\n",
        "\n",
        "\n",
        "def plot_p_val_qq(p_vals: np.ndarray, title: str) -> None:\n",
        "    fig, axes = plt.subplots(nrows=3, ncols=3, sharex=True, sharey=True, figsize=(12,10))\n",
        "    fig.suptitle(title)\n",
        "    n = len(p_vals)\n",
        "    for i in range(9):\n",
        "        unifs = p_vals if i==4 else np.random.rand(n)\n",
        "        sm.qqplot(unifs, uniform(), line='45', ax=axes[i//3,i%3])\n",
        "        if i//3 < 2:\n",
        "            axes[i//3,i%3].set_xlabel('')\n",
        "        if i%3 != 0:\n",
        "            axes[i//3,i%3].set_ylabel('')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "65122a3c",
      "metadata": {
        "id": "65122a3c"
      },
      "outputs": [],
      "source": [
        "plot_p_val_qq(p_vals_mmd, 'Q-Q plot MMD detector')\n",
        "plot_p_val_qq(p_vals_cad, 'Q-Q plot Context-Aware MMD detector')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "12545ff6",
      "metadata": {
        "id": "12545ff6"
      },
      "source": [
        "## Changing the relative subpopulation prevalence\n",
        "\n",
        "We now focus our attention on a more realistic problem where the relative frequency of one or more subpopulations (i.e. news topics) is changing in a way which can be attributed to external events. Importantly, the distribution underlying each subpopulation (e.g. the distribution of *hockey* news itself) remains unchanged, only its frequency changes.\n",
        "\n",
        "In our example we assume that the World Series and Stanley Cup coincide on the calendar leading to a spike in news articles on respectively baseball and hockey. Furthermore, there is not too much news on Mac or Windows since there are no new releases or products planned anytime soon."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2d153de7",
      "metadata": {
        "id": "2d153de7"
      },
      "outputs": [],
      "source": [
        "n_ref_c = 2000 // n_classes\n",
        "n_test_c = [100] * n_classes\n",
        "n_test_c[4], n_test_c[5] = 50, 50  # few stories on Mac/Windows\n",
        "n_test_c[9], n_test_c[10] = 150, 150  # more stories on baseball/hockey"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d66209a5",
      "metadata": {
        "scrolled": true,
        "id": "d66209a5"
      },
      "outputs": [],
      "source": [
        "n_runs = 50\n",
        "\n",
        "p_vals_mmd, p_vals_cad = [], []\n",
        "for _ in tqdm(range(n_runs)):\n",
        "    \n",
        "    # sample data\n",
        "    idx_ref, idx_test, _ = split_data(x_drift, y_drift, n_ref_c, n_test_c, return_idx=True)\n",
        "    emb_ref, c_ref = emb_drift[idx_ref], c_drift[idx_ref]\n",
        "    emb_test, c_test = emb_drift[idx_test], c_drift[idx_test]\n",
        "    \n",
        "    # mmd drift detector\n",
        "    dd_mmd = MMDDrift(emb_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_mmd = dd_mmd.predict(emb_test)\n",
        "    p_vals_mmd.append(preds_mmd['data']['p_val'])\n",
        "    \n",
        "    # context-aware mmd drift detector \n",
        "    dd_cad = ContextMMDDrift(emb_ref, c_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_cad = dd_cad.predict(emb_test, c_test)\n",
        "    p_vals_cad.append(preds_cad['data']['p_val'])\n",
        "    \n",
        "p_vals_mmd = np.array(p_vals_mmd)\n",
        "p_vals_cad = np.array(p_vals_cad)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c69fac36",
      "metadata": {
        "id": "c69fac36"
      },
      "source": [
        "While the **context-aware detector remains well calibrated**, the MMD detector consistently flags drift (low p-values). Note that this is the expected behaviour since the vanilla MMD detector cannot take any external context into account and correctly detects that the reference and test data do not follow the same underlying distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6aa3031e",
      "metadata": {
        "id": "6aa3031e"
      },
      "outputs": [],
      "source": [
        "plot_p_val_qq(p_vals_mmd, 'Q-Q plot MMD detector')\n",
        "plot_p_val_qq(p_vals_cad, 'Q-Q plot Context-Aware MMD detector')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "16fcb992",
      "metadata": {
        "id": "16fcb992"
      },
      "source": [
        "We can also easily see this on the plot below where the p-values of the context-aware detector are uniformly distributed while the MMD detector's p-values are consistently close to 0. Note that we limited the y-axis range to make the plot easier to read."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "32546597",
      "metadata": {
        "id": "32546597"
      },
      "outputs": [],
      "source": [
        "import seaborn as sns\n",
        "\n",
        "def plot_hist(\n",
        "    p_vals: List[np.ndarray],\n",
        "    title: str,\n",
        "    colors: List[str] = ['salmon', 'turquoise'],\n",
        "    methods: List[str] = ['MMD', 'CA-MMD']\n",
        "):\n",
        "    for p_val, method, color in zip(p_vals, methods, colors):\n",
        "        sns.distplot(p_val, color=color, norm_hist=True, kde=True, label=f'{method}', hist=True)\n",
        "        plt.legend(loc='upper right')\n",
        "    plt.xlim(0, 1)\n",
        "    plt.ylim(0, 20)\n",
        "    plt.ylabel('Density')\n",
        "    plt.xlabel('p-values')\n",
        "    plt.title(title)\n",
        "    plt.show();\n",
        "    \n",
        "p_vals = [p_vals_mmd, p_vals_cad]\n",
        "title = 'p-value distribution for a change in subpopulation prevalence'\n",
        "plot_hist(p_vals, title)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d3faa8e2",
      "metadata": {
        "id": "d3faa8e2"
      },
      "source": [
        "## Changing the subpopulation distribution\n",
        "\n",
        "In the following example we change the distribution of one or more of the underlying subpopulations. Notice that now we do want to **flag drift** since our context variable, which permits changes in relative subpopulation prevalences, can no longer explain the change in distribution. \n",
        "\n",
        "Imagine our news topic classification model is not as granular as before and instead of the 20 categories only predicts the 6 *super* classes, organised by subject matter:\n",
        "\n",
        "1. **Computers**: comp.graphics; comp.os.ms-windows.misc; comp.sys.ibm.pc.hardware; comp.sys.mac.hardware; comp.windows.x\n",
        "2. **Recreation**: rec.autos; rec.motorcycles; rec.sport.baseball; rec.sport.hockey\n",
        "3. **Science**: sci.crypt; sci.electronics; sci.med; sci.space\n",
        "4. **Miscellaneous**:  misc.forsale\n",
        "5. **Politics**: talk.politics.misc; talk.politics.guns; talk.politics.mideast\n",
        "6. **Religion**: talk.religion.misc; talk.atheism; soc.religion.christian\n",
        "\n",
        "\n",
        "What if baseball and hockey become less popular and the distribution underlying the *Recreation* class changes? We will want to detect this as the change in distributions of the subpopulations (the 6 super classes) cannot be explained anymore by the context variable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a4b1f685",
      "metadata": {
        "id": "a4b1f685"
      },
      "outputs": [],
      "source": [
        "# map the original target labels to super classes\n",
        "class_map = {\n",
        "    0: [1, 2, 3, 4, 5],\n",
        "    1: [7, 8, 9, 10],\n",
        "    2: [11, 12, 13, 14],\n",
        "    3: [6],\n",
        "    4: [16, 17, 18],\n",
        "    5: [0, 15, 19]\n",
        "}\n",
        "\n",
        "def map_to_super(y: np.ndarray):\n",
        "    y_super = np.zeros_like(y)\n",
        "    for k, v in class_map.items():\n",
        "        for _ in v:\n",
        "            idx_chg = np.where(y == _)[0]\n",
        "            y_super[idx_chg] = k\n",
        "    return y_super\n",
        "\n",
        "y_drift_super = map_to_super(y_drift)\n",
        "n_super = len(list(class_map.keys()))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a7a02fd7",
      "metadata": {
        "id": "a7a02fd7"
      },
      "source": [
        "In order to reuse our pretrained classifier for the super classes, we add the following helper function to map the predictions on the super classes and return one-hot encoded predictions over the 6 super classes. Note that our context variable now changes from a probability distribution over the 20 news topics to a one-hot encoded representation over the 6 super classes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "775e9e02",
      "metadata": {
        "id": "775e9e02"
      },
      "outputs": [],
      "source": [
        "def ohe_super_preds(x: List[str], y: np.ndarray):\n",
        "    classes = np.argmax(context(x, y), -1)  # class predictions\n",
        "    classes_super = map_to_super(classes)  # map to super classes\n",
        "    return np.eye(n_super, dtype=np.float32)[classes_super]  # return OHE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0084c32e",
      "metadata": {
        "scrolled": true,
        "id": "0084c32e"
      },
      "outputs": [],
      "source": [
        "n_ref_c, n_test_c = 1000 // n_super, 1000 // n_super\n",
        "n_runs = 50\n",
        "\n",
        "p_vals_mmd, p_vals_cad = [], []\n",
        "for _ in tqdm(range(n_runs)):\n",
        "    \n",
        "    # sample data\n",
        "    (x_ref, y_ref, y_ref2), (x_test, y_test, y_test2), (x_held, y_held, y_held2) = \\\n",
        "        split_data(x_drift, y_drift_super, n_ref_c, n_test_c, y2=y_drift)\n",
        "\n",
        "    # remove baseball and hockey from the recreation super class in the test set\n",
        "    idx_bb, idx_hock = np.where(y_test2 == 9)[0], np.where(y_test2 == 10)[0]\n",
        "    idx_remove = np.concatenate([idx_bb, idx_hock], 0)\n",
        "    x_test = [x_test[_] for _ in np.arange(len(x_test)) if _ not in idx_remove]\n",
        "    y_test = np.delete(y_test, idx_remove)\n",
        "    \n",
        "    # embed text\n",
        "    emb_ref = model(x_ref)\n",
        "    emb_test = model(x_test)\n",
        "    \n",
        "    # mmd drift detector\n",
        "    dd_mmd = MMDDrift(emb_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_mmd = dd_mmd.predict(emb_test)\n",
        "    p_vals_mmd.append(preds_mmd['data']['p_val'])\n",
        "    \n",
        "    # context-aware mmd drift detector \n",
        "    c_ref = ohe_super_preds(x_ref, y_ref)\n",
        "    c_test = ohe_super_preds(x_test, y_test)\n",
        "    dd_cad = ContextMMDDrift(emb_ref, c_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_cad = dd_cad.predict(emb_test, c_test)\n",
        "    p_vals_cad.append(preds_cad['data']['p_val'])\n",
        "    \n",
        "p_vals_mmd = np.array(p_vals_mmd)\n",
        "p_vals_cad = np.array(p_vals_cad)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5b3e724b",
      "metadata": {
        "id": "5b3e724b"
      },
      "source": [
        "We can see that the context-aware detector is powerful to detect changes in the distributions of the subpopulations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4a2cea47",
      "metadata": {
        "id": "4a2cea47"
      },
      "outputs": [],
      "source": [
        "threshold = .05\n",
        "print(f'Power at {threshold * 100}% significance level')\n",
        "print(f'MMD: {(p_vals_mmd < threshold).mean():.3f}')\n",
        "print(f'Context-aware MMD: {(p_vals_cad < threshold).mean():.3f}')\n",
        "\n",
        "p_vals = [p_vals_mmd, p_vals_cad]\n",
        "title = 'p-value distribution for a change in subpopulation distribution'\n",
        "plot_hist(p_vals, title)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8a4bace3",
      "metadata": {
        "id": "8a4bace3"
      },
      "source": [
        "## Detect unseen topics\n",
        "\n",
        "Next we illustrate the effectiveness of the context-aware detector to detect new topics which are not present in the reference data. Obviously we also want to **flag drift** in this case. As an example we introduce movie reviews in the test data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8fb9d369",
      "metadata": {
        "scrolled": true,
        "id": "8fb9d369"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"imdb\")\n",
        "x_imdb = dataset['train']['text']\n",
        "n_imdb = len(x_imdb)\n",
        "\n",
        "n_test_imdb = 100\n",
        "n_ref_c = 1000 // n_classes\n",
        "n_test_c = 1000 // n_classes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a0d4404e",
      "metadata": {
        "id": "a0d4404e"
      },
      "outputs": [],
      "source": [
        "n_runs = 50\n",
        "\n",
        "p_vals_mmd, p_vals_cad = [], []\n",
        "for _ in tqdm(range(n_runs)):\n",
        "    \n",
        "    # sample data\n",
        "    idx_ref, idx_test, _ = split_data(x_drift, y_drift, n_ref_c, n_test_c, return_idx=True)\n",
        "    emb_ref, c_ref = emb_drift[idx_ref], c_drift[idx_ref]\n",
        "    emb_test, c_test = emb_drift[idx_test], c_drift[idx_test]\n",
        "    \n",
        "    # add random imdb reviews to the test data\n",
        "    idx_imdb = np.random.choice(n_imdb, n_test_imdb, replace=False)\n",
        "    x_imdb_sample = [x_imdb[_] for _ in idx_imdb]\n",
        "    emb_imdb = model(x_imdb_sample)\n",
        "    c_imdb = context(x_imdb_sample, np.zeros(len(x_imdb_sample)))  # value second arg does not matter\n",
        "    emb_test = np.concatenate([emb_test, emb_imdb], 0)\n",
        "    c_test = np.concatenate([c_test, c_imdb], 0)\n",
        "    \n",
        "    # mmd drift detector\n",
        "    dd_mmd = MMDDrift(emb_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_mmd = dd_mmd.predict(emb_test)\n",
        "    p_vals_mmd.append(preds_mmd['data']['p_val'])\n",
        "    \n",
        "    # context-aware mmd drift detector\n",
        "    dd_cad = ContextMMDDrift(emb_ref, c_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_cad = dd_cad.predict(emb_test, c_test)\n",
        "    p_vals_cad.append(preds_cad['data']['p_val'])\n",
        "    \n",
        "p_vals_mmd = np.array(p_vals_mmd)\n",
        "p_vals_cad = np.array(p_vals_cad)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "970f12c8",
      "metadata": {
        "id": "970f12c8"
      },
      "outputs": [],
      "source": [
        "threshold = .05\n",
        "print(f'Power at {threshold * 100}% significance level')\n",
        "print(f'MMD: {(p_vals_mmd < threshold).mean():.3f}')\n",
        "print(f'Context-aware MMD: {(p_vals_cad < threshold).mean():.3f}')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "29967207",
      "metadata": {
        "id": "29967207"
      },
      "source": [
        "\n",
        "## Changing the context variable\n",
        "\n",
        "So far we have conditioned the context-aware detector on the model predictions. There are however many other useful contexts possible. One such example would be to condition on the predictions of an unsupervised clustering algorithm. To facilitate this, we first apply kernel PCA on the embedding vectors, followed by a Gaussian mixture model which clusters the data into 6 classes (same as the super classes). We will test both the calibration under the null hypothesis (no distribution change) as well as the power when a new topic (movie reviews) is injected."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "569aa1be",
      "metadata": {
        "id": "569aa1be"
      },
      "outputs": [],
      "source": [
        "from sklearn.decomposition import KernelPCA\n",
        "from sklearn.mixture import GaussianMixture\n",
        "\n",
        "# embed training data\n",
        "emb_train = model(x_train)\n",
        "\n",
        "# apply kernel PCA to reduce dimensionality\n",
        "kernel_pca = KernelPCA(n_components=10, kernel='linear')\n",
        "kernel_pca.fit(emb_train)\n",
        "emb_train_pca = kernel_pca.transform(emb_train)\n",
        "emb_drift_pca = kernel_pca.transform(emb_drift)\n",
        "\n",
        "# cluster the data\n",
        "y_train_super = map_to_super(y_train)\n",
        "n_clusters = len(np.unique(y_train_super))\n",
        "gmm = GaussianMixture(n_components=n_clusters, covariance_type='full', random_state=2022)\n",
        "gmm.fit(emb_train_pca)\n",
        "c_all_proba = gmm.predict_proba(emb_drift_pca)\n",
        "c_all_class = gmm.predict(emb_drift_pca)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cab00cbb",
      "metadata": {
        "id": "cab00cbb"
      },
      "source": [
        "Next we change the number of instances in each cluster between the reference and test sets. Note that we do not alter the underlying distribution of each of the clusters, just the frequency. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8e1be898",
      "metadata": {
        "id": "8e1be898"
      },
      "outputs": [],
      "source": [
        "# determine cluster proportions for the reference and test samples\n",
        "n_ref_c = [100, 100, 100, 100, 100, 100]\n",
        "n_test_c = [50, 50, 100, 25, 75, 25]\n",
        "\n",
        "def sample_from_clusters():\n",
        "    idx_ref, idx_test = [], []\n",
        "    for _, (i_ref, i_test) in enumerate(zip(n_ref_c, n_test_c)):\n",
        "        idx_c = np.where(c_all_class == _)[0]\n",
        "        idx_shuffle = np.random.choice(idx_c, size=len(idx_c), replace=False)\n",
        "        idx_ref.append(idx_shuffle[:i_ref])\n",
        "        idx_test.append(idx_shuffle[i_ref:i_ref+i_test])\n",
        "    idx_ref = np.concatenate(idx_ref, 0)\n",
        "    idx_test = np.concatenate(idx_test, 0)\n",
        "    c_ref = c_all_proba[idx_ref]\n",
        "    c_test = c_all_proba[idx_test]\n",
        "    emb_ref = emb_drift[idx_ref]\n",
        "    emb_test = emb_drift[idx_test]\n",
        "    return c_ref, c_test, emb_ref, emb_test"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "171dc56f",
      "metadata": {
        "id": "171dc56f"
      },
      "source": [
        "Now we run the experiment and show the context-aware detector's calibration when changing the cluster frequencies. We also show how the usual MMD detector will consistently flag drift. Furthermore, we inject instances from the movie reviews dataset and illustrate that the context-aware detector remains powerful when the underlying cluster distribution changes (by including a previously unseen topic)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bbe8d23e",
      "metadata": {
        "scrolled": true,
        "id": "bbe8d23e"
      },
      "outputs": [],
      "source": [
        "n_test_imdb = 100  # number of imdb instances for each run\n",
        "\n",
        "n_runs = 50\n",
        "\n",
        "p_vals_null, p_vals_alt, p_vals_mmd = [], [], []\n",
        "for _ in tqdm(range(n_runs)):\n",
        "    \n",
        "    # sample data\n",
        "    c_ref, c_test_null, emb_ref, emb_test_null = sample_from_clusters()\n",
        "    \n",
        "    # sample random imdb reviews\n",
        "    idx_imdb = np.random.choice(n_imdb, n_test_imdb, replace=False)\n",
        "    x_imdb_sample = [x_imdb[_] for _ in idx_imdb]\n",
        "    emb_imdb = model(x_imdb_sample)\n",
        "    c_imdb = gmm.predict_proba(kernel_pca.transform(emb_imdb))\n",
        "    \n",
        "    # now we mix in-distribution instances with the imdb reviews\n",
        "    emb_alt = np.concatenate([emb_test_null[:n_test_imdb], emb_imdb], 0)\n",
        "    c_alt = np.concatenate([c_test_null[:n_test_imdb], c_imdb], 0)\n",
        "    \n",
        "    # mmd drift detector\n",
        "    dd_mmd = MMDDrift(emb_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_mmd = dd_mmd.predict(emb_test_null)\n",
        "    p_vals_mmd.append(preds_mmd['data']['p_val'])\n",
        "\n",
        "    # context-aware mmd drift detector\n",
        "    dd = ContextMMDDrift(emb_ref, c_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "    preds_null = dd.predict(emb_test_null, c_test_null)\n",
        "    preds_alt = dd.predict(emb_alt, c_alt)\n",
        "    p_vals_null.append(preds_null['data']['p_val'])\n",
        "    p_vals_alt.append(preds_alt['data']['p_val'])\n",
        "\n",
        "p_vals_null = np.array(p_vals_null)\n",
        "p_vals_alt = np.array(p_vals_alt)\n",
        "p_vals_mmd = np.array(p_vals_mmd)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ded16daf",
      "metadata": {
        "id": "ded16daf"
      },
      "outputs": [],
      "source": [
        "print(f'Power at {threshold * 100}% significance level')\n",
        "print(f'Context-aware MMD: {(p_vals_alt < threshold).mean():.3f}')\n",
        "plot_p_val_qq(p_vals_mmd, 'Q-Q plot MMD detector when changing the cluster frequencies')\n",
        "plot_p_val_qq(p_vals_null, 'Q-Q plot Context-Aware MMD detector when changing the cluster frequencies')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1b0d2bac",
      "metadata": {
        "id": "1b0d2bac"
      },
      "source": [
        "## Interpretability of the context-aware detector\n",
        "\n",
        "The test statistic $\\hat{t}$ of the context-aware MMD detector can be formulated as follows: $\\hat{t} = \\langle K_{0,0}, W_{0,0} \\rangle + \\langle K_{1,1}, W_{1,1} \\rangle -2\\langle K_{0,1}, W_{0,1}\\rangle$ where $0$ refers to the reference data, $1$ to the test data, and $W_{.,.}$ and $K_{.,.}$ are the weight and kernel matrices, respectively. The weight matrices $W_{.,.}$ allow us to focus on the distribution's subpopulations of interest. Reference instances which have similar contexts as the test data will have higher values for their entries in $W_{0,1}$ than instances with dissimilar contexts. We can therefore interpret $W_{0,1}$ as the coupling matrix between instances in the reference and the test sets. This allows us to investigate which subpopulations from the reference set are present and which are missing in the test data. If we also have a good understanding of the model performance on various subpopulations of the reference data, we could even try and use this coupling matrix to roughly proxy model performance on the unlabeled test instances. Note that in this case we would require labels from the reference data and make sure the reference instances come from the validation, not the training set. \n",
        "\n",
        "In the following example we only pick 2 classes to be present in the test set while all 20 are present in the reference set. We can then investigate via the coupling matrix whether the test statistic $\\hat{t}$ focused on the right classes in the reference data via $W_{0,1}$. More concretely, we can sum over the columns (the test instances) of $W_{0,1}$ and check which reference instances obtained the highest weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9ba93452",
      "metadata": {
        "id": "9ba93452"
      },
      "outputs": [],
      "source": [
        "n_ref_c = 2000 // n_classes\n",
        "n_test_c = [0] * n_classes\n",
        "n_test_c[9], n_test_c[10] = 200, 200  # only stories on baseball/hockey\n",
        "(x_ref, y_ref), (x_test, y_test), _ = split_data(x_drift, y_drift, n_ref_c, n_test_c)\n",
        "\n",
        "# embed data\n",
        "emb_ref = model(x_ref)\n",
        "emb_test = model(x_test)\n",
        "\n",
        "# condition using the classifier predictions\n",
        "c_ref = context(x_ref, y_ref)\n",
        "c_test = context(x_test, y_test)\n",
        "\n",
        "# initialise detector and make predictions\n",
        "dd = ContextMMDDrift(emb_ref, c_ref, p_val=.05, n_permutations=100, backend='pytorch')\n",
        "preds = dd.predict(emb_test, c_test, return_coupling=True)\n",
        "\n",
        "# no drift is detected since the distribution of \n",
        "# the subpopulations in the test set remain the same\n",
        "print(f'p-value: {preds[\"data\"][\"p_val\"]:.3f}')\n",
        "\n",
        "# extract coupling matrix between reference and test data\n",
        "W_01 = preds['data']['coupling_xy']\n",
        "\n",
        "# sum over test instances\n",
        "w_ref = W_01.sum(1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8a75e7f1",
      "metadata": {
        "id": "8a75e7f1"
      },
      "outputs": [],
      "source": [
        "# Map the top assigned reference weights to the associated instance labels\n",
        "# and select top 2 * n_ref_c. This tells us what the labels were of the reference \n",
        "# instances with the highest weights in the coupling matrix W_01.\n",
        "# Ideally this would correspond to instances from the baseball and hockey \n",
        "# classes in the reference set (labels 9 and 10).\n",
        "inds_ref_sort = np.argsort(w_ref)[::-1]\n",
        "y_sort = y_ref[inds_ref_sort][:2 * n_ref_c]\n",
        "\n",
        "# And indeed, we can see that we mainly matched with the correct reference instances!\n",
        "correct_match = np.array([y in [9, 10] for y in y_sort]).mean()\n",
        "print(f'The top {100 * correct_match:.2f}% couplings from the top coupled {2 * n_ref_c} instances '\n",
        "      'come from the baseball and hockey classes!')\n",
        "\n",
        "# We can also easily see from the sorted coupling weights that the test statistic \n",
        "# focuses on just the baseball and hockey classes in the reference set and then\n",
        "# the weights in the coupling matrix W_01 fall of a cliff.\n",
        "plt.plot(w_ref[inds_ref_sort]);\n",
        "plt.title('Sorted reference weights from the coupling matrix W_01');\n",
        "plt.ylabel('Reference instance weight in W_01');\n",
        "plt.xlabel('Instances sorted by weight in W_01');\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "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.7.10"
    },
    "colab": {
      "name": "Copy of cd_context_20newsgroup.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}