examples/notebooks/Cell-Type-Annotation.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yXsc6JTIw0HX"
      },
      "source": [
        "# Cell type annotation prediction\n",
        "\n",
        "Run this notebook on google colab to use a free GPU!\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/helicalAI/helical/blob/main/examples/notebooks/Cell-Type-Annotation.ipynb)\n",
        "\n",
        "In this notebook, an [scGPT](https://www.nature.com/articles/s41592-024-02201-0) model is used to predict a cell type annotation with a given gene expression profile.\n",
        "\n",
        "This follows the tutorial from scGPT [here](https://github.com/bowang-lab/scGPT/blob/main/tutorials/Tutorial_Annotation.ipynb), but instead of fine-tuning the entire model, a smaller neural network is trained, using the embeddings of the gene expressions as inputs, to make a prediction.\n",
        "\n",
        "The same approach is made with the [Geneformer](https://www.nature.com/articles/s41586-023-06139-9.epdf?sharing_token=u_5LUGVkd3A8zR-f73lU59RgN0jAjWel9jnR3ZoTv0N2UB4yyXENUK50s6uqjXH69sDxh4Z3J4plYCKlVME-W2WSuRiS96vx6t5ex2-krVDS46JkoVvAvJyWtYXIyj74pDWn_DutZq1oAlDaxfvBpUfSKDdBPJ8SKlTId8uT47M%3D) model and the results are compared against each other.\n",
        "\n",
        "This approach greatly reduces time and complexity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VVhpT3bSxAML",
        "outputId": "49cc4e8e-7599-4320-c9db-d7a35eba6f1f"
      },
      "outputs": [],
      "source": [
        "# !pip install helical"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# !pip install datasets --upgrade"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aDf6LsRIZ80C",
        "outputId": "b7ef3182-a30f-4946-a2d5-72b6c9082b87"
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras.models import Sequential\n",
        "from tensorflow.keras.layers import Dense, Dropout\n",
        "from tensorflow.keras.callbacks import TensorBoard\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.utils import to_categorical\n",
        "import tensorflow as tf\n",
        "from sklearn.metrics import accuracy_score, precision_score, f1_score, recall_score\n",
        "from sklearn.preprocessing import LabelEncoder\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import datetime\n",
        "from sklearn.model_selection import train_test_split\n",
        "import anndata as ad\n",
        "import torch\n",
        "from scipy.sparse import lil_matrix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "## Numa Node Issue https://github.com/tensorflow/tensorflow/issues/42738\n",
        "gpus = tf.config.list_physical_devices('GPU')\n",
        "if gpus:\n",
        "  try:\n",
        "    # Currently, memory growth needs to be the same across GPUs\n",
        "    for gpu in gpus:\n",
        "      tf.config.experimental.set_memory_growth(gpu, True)\n",
        "    logical_gpus = tf.config.list_logical_devices('GPU')\n",
        "    print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n",
        "  except RuntimeError as e:\n",
        "    # Memory growth must be set before GPUs have been initialized\n",
        "    print(e)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from helical.models.scgpt.model import scGPT, scGPTConfig\n",
        "from helical.models.geneformer.model import Geneformer, GeneformerConfig"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We load the our dataset from the EMBL [website](https://www.ebi.ac.uk/biostudies/arrayexpress/studies/E-MTAB-11673). To download the dataset, we are leveraging Huggingface's optimized library!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "ds = load_dataset(\"helical-ai/yolksac_human\",trust_remote_code=True, split=\"train[:65%]\",download_mode=\"reuse_cache_if_exists\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TpvpV2yZpZBO",
        "outputId": "0caaa74e-f0ff-4122-bd27-32af6d85d155"
      },
      "outputs": [],
      "source": [
        "observation_columns = [obs for obs in list(ds.features.keys()) if not obs == 'raw_counts']\n",
        "obs_data = pd.DataFrame(ds.select_columns(observation_columns).data.to_pandas(),columns=observation_columns)\n",
        "lil = lil_matrix((len(ds),ds[0]['size']))\n",
        "lil.data = np.array(ds['raw_counts'],dtype=\"object\")\n",
        "lil.rows = np.array(ds['rows'],dtype=\"object\")\n",
        "adata = ad.AnnData(lil.tocsr(),obs=obs_data)\n",
        "adata.var_names = ds.features['raw_counts'].id.split(\",\")\n",
        "adata.var['gene_name'] = adata.var_names.str.upper()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XME3WOUl-BSf"
      },
      "source": [
        "Let's familiarise ourselves with the data!\n",
        "\n",
        "We are interested in the names of the cells we want to predict. They are saved in `adata.obs[\"LVL1\"]`.\n",
        "\n",
        "Additionally, we need to know how many distinct cell types/classes we have."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sjo5qHPFw0HY"
      },
      "outputs": [],
      "source": [
        "# get labels: the celltype\n",
        "num_types = adata.obs[\"LVL1\"].unique().shape[0]\n",
        "id2type = dict(enumerate(adata.obs[\"LVL1\"].astype(\"category\").cat.categories))\n",
        "\n",
        "celltypes_labels = np.array(adata.obs[\"LVL1\"].tolist())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VzPLNqRr-BSg"
      },
      "source": [
        "This is all summarized in this dictionary:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_j-x6weFm6UO",
        "outputId": "5d0c52f1-553e-460f-def7-7cf372b9bc89"
      },
      "outputs": [],
      "source": [
        "id2type"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HayHZKTC-BSi"
      },
      "source": [
        "Use the Helical package to get the embeddings of the gene expression profile.\n",
        "\n",
        "The only thing we need to specify is the column containing the names of the genes. (`gene_name` in this case)\n",
        "\n",
        "The resulting embeddings are the input features `x` for our smaller NN model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_cXrEgxHC-Wi"
      },
      "source": [
        "# scGPT"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "t0hNoTJdw0Ha",
        "outputId": "26c17bfa-13de-4ccb-beb4-a725a40dd97f"
      },
      "outputs": [],
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "scgpt_config = scGPTConfig(batch_size=50, device=device)\n",
        "scgpt = scGPT(configurer = scgpt_config)\n",
        "data = scgpt.process_data(adata, gene_names = \"gene_name\")\n",
        "x_scgpt = scgpt.get_embeddings(data)\n",
        "x_scgpt.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YxtVwacW-BSj"
      },
      "source": [
        "With the input features, we also need the corresponding labels `y`.\n",
        "\n",
        "They correspond to the cell type labels.\n",
        "\n",
        "As we have a categorical prediction, we one-hot-encode the 6 different cell type labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zHMZ7VX8m6UQ",
        "outputId": "d40c15bd-4ea5-4821-def3-7f253a3af469"
      },
      "outputs": [],
      "source": [
        "y = celltypes_labels\n",
        "num_classes = num_types\n",
        "\n",
        "encoder = LabelEncoder()\n",
        "y_encoded = encoder.fit_transform(y)\n",
        "y_encoded = to_categorical(y_encoded, num_classes=num_classes)\n",
        "y_encoded.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wQpfYP2z-BSk"
      },
      "source": [
        "## Define and train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iVtFwjSp4rgs",
        "outputId": "4b68dd99-6615-4912-97bb-8267135c40d5"
      },
      "outputs": [],
      "source": [
        "input_shape = (512,)\n",
        "\n",
        "# define the model\n",
        "head_model = Sequential()\n",
        "head_model.add(Dense(128, activation='relu', input_shape=input_shape))\n",
        "head_model.add(Dropout(0.4))\n",
        "head_model.add(Dense(32, activation='relu'))\n",
        "head_model.add(Dropout(0.4))\n",
        "head_model.add(Dense(num_classes, activation='softmax'))\n",
        "optimizer = Adam(learning_rate=0.001)\n",
        "head_model.compile(loss='categorical_crossentropy', optimizer=optimizer)\n",
        "head_model.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m_MksIuNw0Hb",
        "outputId": "1c316b1e-87b3-4ad5-d283-e6e810c44d19"
      },
      "outputs": [],
      "source": [
        "head_model_scgpt = head_model\n",
        "X_train, X_test, y_train, y_test = train_test_split(x_scgpt, y_encoded, test_size=0.1, random_state=42)\n",
        "\n",
        "# Setup callbacks\n",
        "log_dir = \"logs/fit/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "tensorboard_callback = TensorBoard(log_dir=log_dir, histogram_freq=1)\n",
        "\n",
        "history = head_model_scgpt.fit(X_train, y_train, epochs=50, batch_size=64, validation_data=(X_test, y_test), callbacks=[tensorboard_callback])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UX4bwC0xw0Hc",
        "outputId": "d726323c-ce91-4be6-a09c-608f1dda4bd7"
      },
      "outputs": [],
      "source": [
        "predictions_nn = head_model_scgpt.predict(X_test)\n",
        "y_pred = np.argmax(predictions_nn, axis=1)\n",
        "y_true = np.argmax(y_test, axis=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uy8_LASN-BSl"
      },
      "source": [
        "## Present the results\n",
        "- on the test set and,\n",
        "- a separate, unseen evaluation set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Lkv-mfDuKUP"
      },
      "outputs": [],
      "source": [
        "def get_evaluations(name_data_set, y_true, y_pred) -> dict:\n",
        "  accuracy = accuracy_score(y_true, y_pred)\n",
        "  precision = precision_score(y_true, y_pred, average='macro')\n",
        "  f1 = f1_score(y_true, y_pred, average='macro')\n",
        "  recall = recall_score(y_true, y_pred, average='macro')\n",
        "  print(f\"{name_data_set} accuracy: {(accuracy*100):.1f}%\")\n",
        "  print(f\"{name_data_set} precision: {(precision*100):.1f}%\")\n",
        "  print(f\"{name_data_set} f1: {(f1*100):.1f}%\")\n",
        "  print(f\"{name_data_set} recall: {(recall*100):.1f}%\")\n",
        "  return {\n",
        "      \"accuracy\": accuracy,\n",
        "      \"precision\": precision,\n",
        "      \"f1\": f1,\n",
        "      \"recall\": recall,\n",
        "  }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q3THhrm7uS99",
        "outputId": "6737b828-df9c-4507-c577-80082247918b"
      },
      "outputs": [],
      "source": [
        "get_evaluations(\"Test set\", y_true, y_pred)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PGwcZiNaLkBv"
      },
      "source": [
        "Load the unseen evaluation set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "ds = load_dataset(\"helical-ai/yolksac_human\",trust_remote_code=True, split=\"train[70%:]\",download_mode=\"reuse_cache_if_exists\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "observation_columns = [obs for obs in list(ds.features.keys()) if not obs == 'raw_counts']\n",
        "obs_data = pd.DataFrame(ds.select_columns(observation_columns).data.to_pandas(),columns=observation_columns)\n",
        "lil = lil_matrix((len(ds),ds[0]['size']))\n",
        "lil.data = np.array(ds['raw_counts'],dtype=\"object\")\n",
        "lil.rows = np.array(ds['rows'],dtype=\"object\")\n",
        "adata_unseen = ad.AnnData(lil.tocsr(),obs=obs_data)\n",
        "adata_unseen.var_names = ds.features['raw_counts'].id.split(\",\")\n",
        "adata_unseen.var['gene_name'] = adata_unseen.var_names.str.upper()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "18bomcpxlxtY",
        "outputId": "ab17404e-20c0-4590-b730-492de07e55bd"
      },
      "outputs": [],
      "source": [
        "data_unseen = scgpt.process_data(adata_unseen, gene_names=\"gene_name\")\n",
        "x_unseen = scgpt.get_embeddings(data_unseen)\n",
        "predictions_nn_unseen = head_model_scgpt.predict(x_unseen)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CTDOE-0uLFAq"
      },
      "source": [
        "We should double check that the cell types are mapped to the correct id numbers for both the training data and this new data set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h0Z6ZTBEJuuY",
        "outputId": "afab8809-b893-42fc-865d-4149f181a3ab"
      },
      "outputs": [],
      "source": [
        "num_types = adata_unseen.obs[\"LVL1\"].unique().shape[0]\n",
        "id2type_unseen = dict(enumerate(adata_unseen.obs[\"LVL1\"].astype(\"category\").cat.categories))\n",
        "id2type_unseen == id2type"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "loLEKJr3qZvL",
        "outputId": "32ec87e0-96d3-4d5e-a1ae-6d3244e80ddb"
      },
      "outputs": [],
      "source": [
        "y_true_unseen = np.array(adata_unseen.obs[\"LVL1\"].tolist())\n",
        "y_pred_unseen = [id2type[prediction] for prediction in np.argmax(predictions_nn_unseen, axis=1)]\n",
        "\n",
        "scgpt_results = get_evaluations(\"Evaluation set\", y_true_unseen, y_pred_unseen)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RgMbzX-z-BSm"
      },
      "source": [
        "Plot a confusion matrix to visualise the classification performance for each the cell type. This is done for the evalation set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "3f-apdJ4-BSm",
        "outputId": "3a195685-3539-4eb3-a90c-f029ecaf7949"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import confusion_matrix\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "set_predicted_cell_types = list(adata_unseen.obs[\"LVL1\"].unique())\n",
        "for i in set(y_pred_unseen):\n",
        "    if i not in set_predicted_cell_types:\n",
        "        set_predicted_cell_types.remove(i)\n",
        "\n",
        "cm = confusion_matrix(y_true_unseen, y_pred_unseen)\n",
        "cm = cm.astype(\"float\") / cm.sum(axis=1)[:, np.newaxis]\n",
        "cm = pd.DataFrame(cm, index=set_predicted_cell_types[:cm.shape[0]], columns=set_predicted_cell_types[:cm.shape[1]])\n",
        "plt.figure(figsize=(12, 12))\n",
        "sns.heatmap(cm, annot=True, fmt=\".1f\", cmap=\"Blues\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XdG0swPrC5sh"
      },
      "source": [
        "# Geneformer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's do the same with the Geneformer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 285,
          "referenced_widgets": [
            "9be7ad77cfbb4aebbf850bce0d703286",
            "96cf6bd80b23485a91309f97fceae386",
            "6793d6aa02fe4992b2a3079935804e68",
            "8b87146682714f9394465afb70fb6296",
            "8602ffeb4a5548f3b88e0011ff526b60",
            "dd28b99c3789438ba123044f9d8bd481",
            "11d0b1de79104633bd46099d3fc9f845",
            "b0387ec16b2346d6a690f8bd8ef7aae9",
            "0fcbdfa61d8c47499ea91c045e6bdb88",
            "475e8bc25e844eb7bda64024f6fd328b",
            "06b297ce68a54494925355922e6a3e60",
            "129fbaa5b93146839d3e095e066689b2",
            "75250bb34c7a4816bb50bab813367823",
            "18e7db16f29b4c5fb37b12b7ecabd994",
            "811ae6f2713942459f44a8ad0b78f1c8",
            "8333081f6e6844fb9872d9d14b9798f4",
            "1a9866fd593a4773a9336548d9fb4e7c",
            "6105972c9a224bb2b83c39d498d82694",
            "ac8fd8833f8b434c98ebf9de7fd7c956",
            "746baf43caea4ee2831817691b968573",
            "94072eee12234d52906f701dbaec4ea7",
            "0ba9a6eb6f0c4d63ae804a09f5d3b7a5"
          ]
        },
        "id": "5P_7jNYDqbYp",
        "outputId": "2fcd4258-77bb-47f1-94fd-6e9742f02e04"
      },
      "outputs": [],
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "if 'rows' in adata.obs:\n",
        "    adata.obs['rows'] = adata.obs['rows'].astype(str)\n",
        "geneformer_config = GeneformerConfig(batch_size=50, device=device)\n",
        "geneformer = Geneformer(configurer = geneformer_config)\n",
        "\n",
        "data_geneformer = geneformer.process_data(adata, gene_names = \"gene_name\")\n",
        "x_geneformer = geneformer.get_embeddings(data_geneformer)\n",
        "x_geneformer.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ITj-Gd-TDIcu",
        "outputId": "4946c091-3340-4332-e208-2ab2fc46f15e"
      },
      "outputs": [],
      "source": [
        "head_model_geneformer = head_model\n",
        "X_train, X_test, y_train, y_test = train_test_split(x_geneformer, y_encoded, test_size=0.1, random_state=42)\n",
        "\n",
        "# Setup callbacks\n",
        "log_dir = \"logs/fit/\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "tensorboard_callback = TensorBoard(log_dir=log_dir, histogram_freq=1)\n",
        "\n",
        "history = head_model_geneformer.fit(X_train, y_train, epochs=50, batch_size=64, validation_data=(X_test, y_test), callbacks=[tensorboard_callback])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 276
        },
        "id": "jL-tQg35L_cT",
        "outputId": "bd9c6f55-3a87-4ef7-e693-3c1cb47dd0d5"
      },
      "outputs": [],
      "source": [
        "# %tensorboard --logdir logs/fit # uncomment to use tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 156,
          "referenced_widgets": [
            "042588def40e42848c464da311dfa395",
            "9f06f76e74ba4d4493a2b3810fa78501",
            "afde76acb7cd4642967e74fc01cc5e60",
            "bed90884638d44a1a0e8f2002485a49c",
            "577aff58497a41afb68c49c4e4a3483c",
            "e6639753498c44abbb25f2dd39f224a6",
            "0ae5dd4c69154c36ac6409002dbfdbe3",
            "98077a70b5aa4738b2f9ce21fc8a76f8",
            "6c0eae3868aa4ea98c7c1106aba4cb84",
            "08d07e9439494aa895b90bda386cd2d9",
            "c284a26712fc46a49f40081765203b11",
            "9e58f857ee8a43ac94614686ba82cf82",
            "21078359f8b64f1f9e5b70fbb559ac5e",
            "0536305995ae41b989c00348dec37a50",
            "da40c55164824b2e81fef416cea0bdf3",
            "45e0ffdc378040c8b8fb62172c321be5",
            "031f15e503284b6e989e599834bc0cfc",
            "d6a09da0033a4fff88ce7ad041ea5151",
            "44164c48e4774fd28c8e7514e26dfec3",
            "7fba8d16f48b451ba2695e7d8d0c605b",
            "8f571f239ad54eb19d35da67e47ddea9",
            "fb07aa128c734c1396a8d2e8e46dfd7e"
          ]
        },
        "id": "ycWcn3Ra1YbG",
        "outputId": "8e38e24d-261d-4d77-b704-a8d3c6df9b73"
      },
      "outputs": [],
      "source": [
        "data_unseen_geneformer = geneformer.process_data(adata_unseen, gene_names = \"gene_name\")\n",
        "x_unseen_geneformer = geneformer.get_embeddings(data_unseen_geneformer)\n",
        "predictions_nn_unseen_geneformer = head_model_geneformer.predict(x_unseen_geneformer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eR8jNVcrEMRR",
        "outputId": "5d4ff27a-3bbb-409d-db36-61f787dc1b0e"
      },
      "outputs": [],
      "source": [
        "y_true_unseen = np.array(adata_unseen.obs[\"LVL1\"].tolist())\n",
        "y_pred_unseen = [id2type[prediction] for prediction in np.argmax(predictions_nn_unseen_geneformer, axis=1)]\n",
        "\n",
        "geneformer_results = get_evaluations(\"Evaluation set\", y_true_unseen, y_pred_unseen)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 990
        },
        "id": "sK4UqSqQFSaS",
        "outputId": "b1d94806-ed84-4061-d50b-d6b567df66f4"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "fig, axs = plt.subplots(2, 2, figsize=(10, 10))\n",
        "\n",
        "values_1 = [scgpt_results[\"accuracy\"], geneformer_results[\"accuracy\"]]\n",
        "x = [\"scGPT\", \"Geneformer\"]\n",
        "axs[0, 0].bar(x, values_1, width=0.4)\n",
        "axs[0, 0].set_title(\"Accuracy\")\n",
        "axs[0, 0].set_ylim([0, 1])\n",
        "\n",
        "values_2 = [scgpt_results[\"precision\"], geneformer_results[\"precision\"]]\n",
        "axs[0, 1].bar(x, values_2, width=0.4)\n",
        "axs[0, 1].set_title(\"Precision\")\n",
        "axs[0, 1].set_ylim([0, 1])\n",
        "\n",
        "values_3 = [scgpt_results[\"f1\"], geneformer_results[\"f1\"]]\n",
        "axs[1, 0].bar(x, values_3, width=0.4)\n",
        "axs[1, 0].set_title(\"F1\")\n",
        "axs[1, 0].set_ylim([0, 1])\n",
        "\n",
        "values_4 = [scgpt_results[\"recall\"], geneformer_results[\"recall\"]]\n",
        "axs[1, 1].bar(x, values_4, width=0.4)\n",
        "axs[1, 1].set_title(\"Recall\")\n",
        "axs[1, 1].set_ylim([0, 1])\n",
        "\n",
        "fig.suptitle(\"scGPT vs. Geneformer \\n Probing Comparison\")\n",
        "fig.tight_layout()\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W4yqDgBPMNhX"
      },
      "source": [
        "\n",
        "## scGPT\n",
        "- Accuracy: 99.2%\n",
        "- Precision: 84.8%\n",
        "- Recall: 85.6%\n",
        "- Macro F1: 87.3%\n",
        "\n",
        "## Geneformer\n",
        "- Accuracy: 99.0%\n",
        "- Precision: 79.2%\n",
        "- Recall: 78.1%\n",
        "- Macro F1: 77.8%"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        " (c) Helical 2024 - Developed by the Helical Team"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11.8"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "031f15e503284b6e989e599834bc0cfc": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "042588def40e42848c464da311dfa395": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_9f06f76e74ba4d4493a2b3810fa78501",
              "IPY_MODEL_afde76acb7cd4642967e74fc01cc5e60",
              "IPY_MODEL_bed90884638d44a1a0e8f2002485a49c"
            ],
            "layout": "IPY_MODEL_577aff58497a41afb68c49c4e4a3483c"
          }
        },
        "0536305995ae41b989c00348dec37a50": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_44164c48e4774fd28c8e7514e26dfec3",
            "max": 157,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_7fba8d16f48b451ba2695e7d8d0c605b",
            "value": 157
          }
        },
        "06b297ce68a54494925355922e6a3e60": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "08d07e9439494aa895b90bda386cd2d9": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0ae5dd4c69154c36ac6409002dbfdbe3": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0ba9a6eb6f0c4d63ae804a09f5d3b7a5": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0fcbdfa61d8c47499ea91c045e6bdb88": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "11d0b1de79104633bd46099d3fc9f845": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "129fbaa5b93146839d3e095e066689b2": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_75250bb34c7a4816bb50bab813367823",
              "IPY_MODEL_18e7db16f29b4c5fb37b12b7ecabd994",
              "IPY_MODEL_811ae6f2713942459f44a8ad0b78f1c8"
            ],
            "layout": "IPY_MODEL_8333081f6e6844fb9872d9d14b9798f4"
          }
        },
        "18e7db16f29b4c5fb37b12b7ecabd994": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ac8fd8833f8b434c98ebf9de7fd7c956",
            "max": 157,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_746baf43caea4ee2831817691b968573",
            "value": 157
          }
        },
        "1a9866fd593a4773a9336548d9fb4e7c": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "21078359f8b64f1f9e5b70fbb559ac5e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_031f15e503284b6e989e599834bc0cfc",
            "placeholder": "​",
            "style": "IPY_MODEL_d6a09da0033a4fff88ce7ad041ea5151",
            "value": "100%"
          }
        },
        "44164c48e4774fd28c8e7514e26dfec3": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "45e0ffdc378040c8b8fb62172c321be5": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "475e8bc25e844eb7bda64024f6fd328b": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "577aff58497a41afb68c49c4e4a3483c": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6105972c9a224bb2b83c39d498d82694": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "6793d6aa02fe4992b2a3079935804e68": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b0387ec16b2346d6a690f8bd8ef7aae9",
            "max": 7844,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_0fcbdfa61d8c47499ea91c045e6bdb88",
            "value": 7844
          }
        },
        "6c0eae3868aa4ea98c7c1106aba4cb84": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "746baf43caea4ee2831817691b968573": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "75250bb34c7a4816bb50bab813367823": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1a9866fd593a4773a9336548d9fb4e7c",
            "placeholder": "​",
            "style": "IPY_MODEL_6105972c9a224bb2b83c39d498d82694",
            "value": "100%"
          }
        },
        "7fba8d16f48b451ba2695e7d8d0c605b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "811ae6f2713942459f44a8ad0b78f1c8": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_94072eee12234d52906f701dbaec4ea7",
            "placeholder": "​",
            "style": "IPY_MODEL_0ba9a6eb6f0c4d63ae804a09f5d3b7a5",
            "value": " 157/157 [03:35&lt;00:00,  1.17s/it]"
          }
        },
        "8333081f6e6844fb9872d9d14b9798f4": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8602ffeb4a5548f3b88e0011ff526b60": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8b87146682714f9394465afb70fb6296": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_475e8bc25e844eb7bda64024f6fd328b",
            "placeholder": "​",
            "style": "IPY_MODEL_06b297ce68a54494925355922e6a3e60",
            "value": " 7844/7844 [00:02&lt;00:00, 3109.57 examples/s]"
          }
        },
        "8f571f239ad54eb19d35da67e47ddea9": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "94072eee12234d52906f701dbaec4ea7": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "96cf6bd80b23485a91309f97fceae386": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_dd28b99c3789438ba123044f9d8bd481",
            "placeholder": "​",
            "style": "IPY_MODEL_11d0b1de79104633bd46099d3fc9f845",
            "value": "Map: 100%"
          }
        },
        "98077a70b5aa4738b2f9ce21fc8a76f8": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9be7ad77cfbb4aebbf850bce0d703286": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_96cf6bd80b23485a91309f97fceae386",
              "IPY_MODEL_6793d6aa02fe4992b2a3079935804e68",
              "IPY_MODEL_8b87146682714f9394465afb70fb6296"
            ],
            "layout": "IPY_MODEL_8602ffeb4a5548f3b88e0011ff526b60"
          }
        },
        "9e58f857ee8a43ac94614686ba82cf82": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_21078359f8b64f1f9e5b70fbb559ac5e",
              "IPY_MODEL_0536305995ae41b989c00348dec37a50",
              "IPY_MODEL_da40c55164824b2e81fef416cea0bdf3"
            ],
            "layout": "IPY_MODEL_45e0ffdc378040c8b8fb62172c321be5"
          }
        },
        "9f06f76e74ba4d4493a2b3810fa78501": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e6639753498c44abbb25f2dd39f224a6",
            "placeholder": "​",
            "style": "IPY_MODEL_0ae5dd4c69154c36ac6409002dbfdbe3",
            "value": "Map: 100%"
          }
        },
        "ac8fd8833f8b434c98ebf9de7fd7c956": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "afde76acb7cd4642967e74fc01cc5e60": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_98077a70b5aa4738b2f9ce21fc8a76f8",
            "max": 7844,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6c0eae3868aa4ea98c7c1106aba4cb84",
            "value": 7844
          }
        },
        "b0387ec16b2346d6a690f8bd8ef7aae9": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "bed90884638d44a1a0e8f2002485a49c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_08d07e9439494aa895b90bda386cd2d9",
            "placeholder": "​",
            "style": "IPY_MODEL_c284a26712fc46a49f40081765203b11",
            "value": " 7844/7844 [00:02&lt;00:00, 4942.70 examples/s]"
          }
        },
        "c284a26712fc46a49f40081765203b11": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d6a09da0033a4fff88ce7ad041ea5151": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "da40c55164824b2e81fef416cea0bdf3": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8f571f239ad54eb19d35da67e47ddea9",
            "placeholder": "​",
            "style": "IPY_MODEL_fb07aa128c734c1396a8d2e8e46dfd7e",
            "value": " 157/157 [03:34&lt;00:00,  1.17s/it]"
          }
        },
        "dd28b99c3789438ba123044f9d8bd481": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e6639753498c44abbb25f2dd39f224a6": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fb07aa128c734c1396a8d2e8e46dfd7e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}