notebooks/extract_localization_from_ProtBert_using_light_attention.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  },
  "colab": {
   "name": "extract_localization_from_ProtBert_using_light_attention.ipynb",
   "provenance": [],
   "collapsed_sections": [],
   "toc_visible": true
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GyFqUpIC877d"
   },
   "source": [
    "# Colab initialization\n",
    "- install the pipeline in the colab runtime\n",
    "- download files neccessary for this example"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q6w5lYWY877f",
    "outputId": "217e25d9-dbd7-48d7-dd24-327c346d0f50"
   },
   "source": [
    "!pip3 install -U pip > /dev/null\n",
    "!pip3 install -U \"bio-embeddings[all] @ git+https://gitlab.lrz.de/sacdallago/bio_embeddings.git@light_attention\" > /dev/null"
   ],
   "execution_count": 1,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The system cannot find the path specified.\n",
      "The system cannot find the path specified.\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "QhY5StD_877g",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "9954627f-c6da-41bf-961c-b86db32a418e"
   },
   "source": [
    "!wget http://data.bioembeddings.com/public/embeddings/feature_models/light_attention/la_bert_subcellular_localization.pt --output-document la_bert_subcellular_localization.pt\n",
    "!wget http://data.bioembeddings.com/public/embeddings/feature_models/light_attention/la_bert_solubility.pt --output-document la_bert_solubility.pt"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "text": [
      "--2021-04-04 07:51:13--  http://data.bioembeddings.com/public/embeddings/feature_models/light_attention/la_bert_subcellular_localization.pt\n",
      "Resolving data.bioembeddings.com (data.bioembeddings.com)... 143.95.108.236\n",
      "Connecting to data.bioembeddings.com (data.bioembeddings.com)|143.95.108.236|:80... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 227320161 (217M) [text/plain]\n",
      "Saving to: ‘localization_unprocessed.pt’\n",
      "\n",
      "localization_unproc 100%[===================>] 216.79M  76.3MB/s    in 2.8s    \n",
      "\n",
      "2021-04-04 07:51:16 (76.3 MB/s) - ‘localization_unprocessed.pt’ saved [227320161/227320161]\n",
      "\n",
      "--2021-04-04 07:51:16--  http://data.bioembeddings.com/public/embeddings/feature_models/light_attention/la_bert_solubility.pt\n",
      "Resolving data.bioembeddings.com (data.bioembeddings.com)... 143.95.108.236\n",
      "Connecting to data.bioembeddings.com (data.bioembeddings.com)|143.95.108.236|:80... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 227317089 (217M) [text/plain]\n",
      "Saving to: ‘solubility_unprocessed.pt’\n",
      "\n",
      "solubility_unproces 100%[===================>] 216.79M  77.0MB/s    in 2.8s    \n",
      "\n",
      "2021-04-04 07:51:19 (77.0 MB/s) - ‘solubility_unprocessed.pt’ saved [227317089/227317089]\n",
      "\n"
     ],
     "name": "stdout"
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "is_executing": false
    },
    "id": "lvBzbnyn877h"
   },
   "source": [
    "# Extract solubility and subcellular localization predictions from ProtTrans T5 using Light Attention\n",
    "In this notebook we will extract annotations from ProtTrans T5 embeddings via trained Light Attention models that can predict solubility and subcellular localization"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "U-Ol083p877h"
   },
   "source": [
    "from bio_embeddings.embed import ProtTransBertBFDEmbedder"
   ],
   "execution_count": 14,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v1De4kYq877h"
   },
   "source": [
    "We initialize the ProtTrans T5 embedder."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "P_euXvNz877i",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "077a57d1-94ce-45f5-ba46-d423b993a097"
   },
   "source": [
    "embedder = ProtTransBertBFDEmbedder()"
   ],
   "execution_count": 15,
   "outputs": [
    {
     "output_type": "stream",
     "text": [
      "prottrans_bert_bfd.zip: 1.56GB [00:18, 84.5MB/s]                            \n"
     ],
     "name": "stderr"
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ya8qR2Bh877i"
   },
   "source": [
    "We select an AA sequence. In this case, the sequence is that of [Aspartate aminotransferase, mitochondrial](https://www.uniprot.org/uniprot/P12345)\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "FnS1qCt2877i"
   },
   "source": [
    "target_sequence = \"MALLHSARVLSGVASAFHPGLAAAASARASSWWAHVEMGPPDPILGVTEAYKRDTNSKKMNLGVGAYRDDNGKPYVLPSVRKAEAQIAAKGLDKEYLPIGGLAEFCRASAELALGENSEVVKSGRFVTVQTISGTGALRIGASFLQRFFKFSRDVFLPKPSWGNHTPIFRDAGMQLQSYRYYDPKTCGFDFTGALEDISKIPEQSVLLLHACAHNPTGVDPRPEQWKEIATVVKKRNLFAFFDMAYQGFASGDGDKDAWAVRHFIEQGINVCLCQSYAKNMGLYGERVGAFTVICKDADEAKRVESQLKILIRPMYSNPPIHGARIASTILTSPDLRKQWLQEVKGMADRIIGMRTQLVSNLKKEGSTHSWQHITDQIGMFCFTGLKPEQVERLTKEFSIYMTKDGRISVAGVTSGNVGYLAHAIHQVTK\""
   ],
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3fvVxHwJ877j"
   },
   "source": [
    "We produce the embeddings of the above sequence. Since we only have one sequence, we use the simple `embed` function, rather than the `embed_many` or `embed_batch`, which we would instead use if we had multiple sequences to embed."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "F1PmgQHa877j"
   },
   "source": [
    "embedding = embedder.embed(target_sequence)"
   ],
   "execution_count": 17,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "m8t4wCiB877j"
   },
   "source": [
    "The `bio_embeddings` pipeline includes some models trained on embeddings for the prediction of Secondary Structure and Subcellular Localization. In the following we make use of these models.\n",
    "\n",
    "To speed up processing, we have downloaded the model weights of the supervised subcellular localization and secondary structure prediction models from [here](http://maintenance.dallago.us/public/embeddings/feature_models/seqvec/)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "brBbIS-N877j"
   },
   "source": [
    "from bio_embeddings.extract.light_attention import LightAttentionAnnotationExtractor"
   ],
   "execution_count": 19,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ldxMslDV877k"
   },
   "source": [
    "annotations_extractor = LightAttentionAnnotationExtractor(subcellular_location_checkpoint_file='la_bert_subcellular_localization.pt', membrane_checkpoint_file='la_bert_solubility.pt')"
   ],
   "execution_count": 20,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "39Kf9_XE877k"
   },
   "source": [
    "annotations = annotations_extractor.get_subcellular_location(embedding)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UDiasSNm877k"
   },
   "source": [
    "Let's see what annotations are available from Light Attention"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "DoVZlsQ0877k",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "ec477faf-b1d5-4762-c20c-817a524c0b78"
   },
   "source": [
    "annotations._fields"
   ],
   "execution_count": 33,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "('localization', 'membrane')"
      ]
     },
     "metadata": {
      "tags": []
     },
     "execution_count": 33
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nw0-AdWG877l"
   },
   "source": [
    "Let's print the subcellular localization predicted via the Light Attention on the ProtBert embeddings"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "dcCa4za_877l",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "49fba49c-3769-4def-ebcf-ea6c5b5e372d"
   },
   "source": [
    "print(f\"The subcellular localization predicted from the embedding is: {annotations.localization.value}\")\n",
    "print(f\"The solubility/membrane-boundness predicted from the embedding is: {annotations.membrane.value}\")"
   ],
   "execution_count": 36,
   "outputs": [
    {
     "output_type": "stream",
     "text": [
      "The subcellular localization predicted from the embedding is: Mitochondrion\n",
      "The embedding is soluble: Soluble\n"
     ],
     "name": "stdout"
    }
   ]
  }
 ]
}