SeekRank.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5QUyhxIqdaXI"
   },
   "source": [
    "<img src=\"https://raw.githubusercontent.com/sokrypton/ColabFold/main/.github/ColabFold_Marv_Logo_Small.png\" height=\"200\" align=\"right\" style=\"height:240px\">\n",
    "\n",
    "##Homologous Protein search and ranking using ESM-1b and ColabFold database\n",
    "Find homologous proteins with similar/better experimental measure. You will need to provide a fasta file with numerical value associared with your sequences e.g. Kcat/Km, toxicity, ... The Colab will search homologous sequences to input proteins with [ColabFold](https://github.com/sokrypton/ColabFold), sorted by the numerical value predicted by a [ESM-1b](https://github.com/facebookresearch/esm) based-regressor.\n",
    "\n",
    "1.    Prepare a fasta file containing the sequence and numerical value representing the sequence (see the example fasta format below)\n",
    "```\n",
    ">seq_1 8.104\n",
    "PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK\n",
    ">seq_2 2.04\n",
    "MIAQIHILEGRSDEDKETLIRRVSEAISRSLDAPLTSVRVIIMEMAKGHFGGELASK\n",
    "```\n",
    "\n",
    "2.   `Runtime` -> `Run all` upload the file (option appears under \"File upload\")\n",
    "\n",
    "\n",
    "<b>Caution!</b><br>\n",
    "If your dataset has high ratio of same numerical values, correlation calculation might cause an error!\n",
    "<br><br>\n",
    "\n",
    "For more detail about what's happening in this colab, please look into the explanation on the bottom.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "I2pbvUlPC5wp"
   },
   "outputs": [],
   "source": [
    "#@title File upload.\n",
    "from google.colab import files\n",
    "import io\n",
    "import warnings\n",
    "#separator = \"\" #@param {type:\"string\"}\n",
    "#if separator == \"\":\n",
    "#  separator = \"\\t\"\n",
    "#@markdown - Run this cell and upload a fasta file\n",
    "#@markdown - Click third cell (meaning skip one below this), then hit `Runtime` -> `Run after`\n",
    "# Read in fasta file\n",
    "uploaded = files.upload()\n",
    "names = []\n",
    "values = []\n",
    "name_value = {}\n",
    "sequences = []\n",
    "seq = \"\"\n",
    "for k, v in uploaded.items():\n",
    "  for line in v.decode(\"utf-8\").split(\"\\n\"):\n",
    "    if len(line) > 0 and line[0] == \">\":\n",
    "      names.append(line.split()[0][1:])\n",
    "      values.append(float(line.split()[-1]))\n",
    "      name_value[line.split()[0][1:]] = float(line.split()[-1])\n",
    "      if seq != \"\":\n",
    "        sequences.append(seq.upper())\n",
    "        seq = \"\"\n",
    "    else:\n",
    "      if line != \"\":\n",
    "        seq += line\n",
    "  sequences.append(seq)\n",
    "assert len(names) == len(values)\n",
    "assert len(values) == len(sequences)\n",
    "if len(list(set(values))) == 1:\n",
    "  warnings.warn(\"Warning: only one value found. Changing procedure to get distance, not a regressor score\")\n",
    "  do_distance = True\n",
    "\n",
    "# Write a fasta file for esm1b (write proteins with <=1024 amino acid length)\n",
    "cnt = 0\n",
    "with open(\"/content/esm1b_input.fasta\", 'w') as new:\n",
    "  for i in range(len(names)):\n",
    "    if len(sequences[i]) < 1023:\n",
    "      new.write(\">%s\\n%s\\n\"%(names[i], sequences[i]))\n",
    "      cnt += 1\n",
    "\n",
    "# Make a warning if cnt is smaller than 40\n",
    "if cnt < 40:\n",
    "  warnings.warn(\"Warning: Training data size smaller than 40! We cannot ensure the result quality if the size of training data is too small\")\n",
    "if cnt < 10:\n",
    "  warnings.warn(\"Warning: Training data size smaller than 10, cannot conduct five fold validation, the last layer will be chosen from ESM-1b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "PIXo4meq_ccI"
   },
   "outputs": [],
   "source": [
    "#@title Options\n",
    "do_distance = False #@param {type:\"boolean\"}\n",
    "#@markdown  - Select `do_distance` when you only have positive samples with no numerical labels.<br>If you selected this option, please set all the numerical values to 1.\n",
    "Regressor = \"Random Forest Regressor\" #@param [\"Random Forest Regressor\", \"Linear Regression\", \"K Neighbors Regressor\", \"Logistic Regression\"]\n",
    "#@markdown - Use Logstic Regression only when you have labels of 0 and 1\n",
    "iteration = 5 #@param {type:\"integer\"}\n",
    "#@markdown - Decide how many times you will repeat five fold validation<br>\n",
    "n_neighbors = 5 #@param {type:\"integer\"}\n",
    "#@markdown - Set `n_neighbors` for K Neighbors Regressor, this has to be same or smaller than size of training data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "hPMmxFWm6DjK"
   },
   "outputs": [],
   "source": [
    "#@title Install Colabfold, ESM-1b and ProtT5 (Currently doesn't install ProtT5)\n",
    "!pip install git+https://github.com/facebookresearch/esm.git\n",
    "!pip install \"colabfold[alphafold] @ git+https://github.com/sokrypton/ColabFold\"\n",
    "#!pip install -q SentencePiece transformers\n",
    "\n",
    "from sklearn.linear_model import LogisticRegression, LinearRegression\n",
    "from sklearn.neighbors import KNeighborsRegressor\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "\n",
    "import requests\n",
    "import hashlib\n",
    "import tarfile\n",
    "import time\n",
    "import pickle\n",
    "import os\n",
    "import re\n",
    "\n",
    "import random\n",
    "import tqdm.notebook\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import matplotlib.patheffects\n",
    "from matplotlib import collections as mcoll\n",
    "\n",
    "import os\n",
    "import requests\n",
    "import tqdm.notebook\n",
    "import random\n",
    "import tarfile\n",
    "\n",
    "import torch\n",
    "\n",
    "from esm import Alphabet, FastaBatchedDataset, ProteinBertModel, pretrained, MSATransformer\n",
    "try:\n",
    "    model_esm\n",
    "except NameError:\n",
    "    model_esm, alphabet_esm = pretrained.load_model_and_alphabet(\"esm1b_t33_650M_UR50S\")\n",
    "\n",
    "import pickle\n",
    "\n",
    "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
    "\n",
    "def run_mmseqs2(x, prefix, use_env=True, use_filter=True,\n",
    "                use_templates=False, filter=None, host_url=\"https://api.colabfold.com\"):\n",
    "  \n",
    "  def submit(seqs, mode, N=101):    \n",
    "    n,query = N,\"\"\n",
    "    for seq in seqs: \n",
    "      query += f\">{n}\\n{seq}\\n\"\n",
    "      n += 1\n",
    "      \n",
    "    while True:\n",
    "      try:\n",
    "        # https://requests.readthedocs.io/en/latest/user/advanced/#advanced\n",
    "        # \"good practice to set connect timeouts to slightly larger than a multiple of 3\"\n",
    "        res = requests.post(f'{host_url}/ticket/msa', data={'q':query,'mode': mode}, timeout=6.02)\n",
    "      except requests.exceptions.Timeout:\n",
    "        continue\n",
    "      break\n",
    "\n",
    "    try: out = res.json()\n",
    "    except ValueError: out = {\"status\":\"UNKNOWN\"}\n",
    "    return out\n",
    "\n",
    "  def status(ID):\n",
    "    while True:\n",
    "      try:\n",
    "        res = requests.get(f'{host_url}/ticket/{ID}', timeout=6.02)\n",
    "      except requests.exceptions.Timeout:\n",
    "        continue\n",
    "      break\n",
    "\n",
    "    try: out = res.json()\n",
    "    except ValueError: out = {\"status\":\"UNKNOWN\"}\n",
    "    return out\n",
    "\n",
    "  def download(ID, path):\n",
    "    while True:\n",
    "      try:\n",
    "        res = requests.get(f'{host_url}/result/download/{ID}', timeout=6.02)\n",
    "      except requests.exceptions.Timeout:\n",
    "        continue\n",
    "      break\n",
    "\n",
    "    with open(path,\"wb\") as out: out.write(res.content)\n",
    "  \n",
    "  # process input x\n",
    "  seqs = [x] if isinstance(x, str) else x\n",
    "  \n",
    "  # compatibility to old option\n",
    "  if filter is not None:\n",
    "    use_filter = filter\n",
    "    \n",
    "  # setup mode\n",
    "  if use_filter:\n",
    "    mode = \"env\" if use_env else \"all\"\n",
    "  else:\n",
    "    mode = \"env-nofilter\" if use_env else \"nofilter\"\n",
    "\n",
    "  mode += \"-m8output\"\n",
    "  \n",
    "  # define path\n",
    "  path = f\"{prefix}_{mode}\"\n",
    "  if not os.path.isdir(path): os.mkdir(path)\n",
    "\n",
    "  # call mmseqs2 api\n",
    "  tar_gz_file = f'{path}/out.tar.gz'\n",
    "  N,REDO = 101,True\n",
    "  \n",
    "  # deduplicate and keep track of order\n",
    "  seqs_unique = sorted(list(set(seqs)))\n",
    "  Ms = [N+seqs_unique.index(seq) for seq in seqs]\n",
    "  \n",
    "  # lets do it!\n",
    "  if not os.path.isfile(tar_gz_file):\n",
    "    TIME_ESTIMATE = 150 * len(seqs_unique)\n",
    "    with tqdm.notebook.tqdm(total=TIME_ESTIMATE, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
    "      while REDO:\n",
    "        pbar.set_description(\"SUBMIT\")\n",
    "        \n",
    "        # Resubmit job until it goes through\n",
    "        out = submit(seqs_unique, mode, N)\n",
    "        while out[\"status\"] in [\"UNKNOWN\",\"RATELIMIT\"]:\n",
    "          # resubmit\n",
    "          time.sleep(5 + random.randint(0,5))\n",
    "          out = submit(seqs_unique, mode, N)\n",
    "\n",
    "        if out[\"status\"] == \"ERROR\":\n",
    "          raise Exception(f'MMseqs2 API is giving errors. Please confirm your input is a valid protein sequence. If error persists, please try again an hour later.')\n",
    "\n",
    "        if out[\"status\"] == \"MAINTENANCE\":\n",
    "          raise Exception(f'MMseqs2 API is undergoing maintenance. Please try again in a few minutes.')\n",
    "\n",
    "        # wait for job to finish\n",
    "        ID,TIME = out[\"id\"],0\n",
    "        pbar.set_description(out[\"status\"])\n",
    "        while out[\"status\"] in [\"UNKNOWN\",\"RUNNING\",\"PENDING\"]:\n",
    "          t = 5 + random.randint(0,5)\n",
    "          time.sleep(t)\n",
    "          out = status(ID)    \n",
    "          pbar.set_description(out[\"status\"])\n",
    "          if out[\"status\"] == \"RUNNING\":\n",
    "            TIME += t\n",
    "            pbar.update(n=t)\n",
    "          #if TIME > 900 and out[\"status\"] != \"COMPLETE\":\n",
    "          #  # something failed on the server side, need to resubmit\n",
    "          #  N += 1\n",
    "          #  break\n",
    "        \n",
    "        if out[\"status\"] == \"COMPLETE\":\n",
    "          if TIME < TIME_ESTIMATE:\n",
    "            pbar.update(n=(TIME_ESTIMATE-TIME))\n",
    "          REDO = False\n",
    "\n",
    "      # Download results\n",
    "      download(ID, tar_gz_file)\n",
    "\n",
    "  # prep list of a3m files\n",
    "  m8_files = [f\"{path}/uniref.m8\"]\n",
    "  if use_env: m8_files.append(f\"{path}/bfd.mgnify30.metaeuk30.smag30.m8\")\n",
    "  \n",
    "  # extract a3m files\n",
    "  if not os.path.isfile(m8_files[0]):\n",
    "    with tarfile.open(tar_gz_file) as tar_gz:\n",
    "      tar_gz.extractall(path)  \n",
    "\n",
    "  m8_lines = {}\n",
    "  for m8_file in m8_files:\n",
    "    for line in open(m8_file,\"r\"):\n",
    "      if len(line) > 0:\n",
    "        if \"\\x00\" in line:\n",
    "          line = line.replace(\"\\x00\",\"\")\n",
    "        M = int(line.split()[0])\n",
    "        if M not in m8_lines: m8_lines[M] = []\n",
    "        m8_lines[M].append(line)\n",
    "  \n",
    "  # return results\n",
    "  m8_lines = [\"\".join(m8_lines[n]) for n in Ms]\n",
    "  \n",
    "  if isinstance(x, str):\n",
    "    return m8_lines[0]\n",
    "  else:\n",
    "    return m8_lines\n",
    "\n",
    "# Code for embedding extraction (esm-1b)\n",
    "# Extract for every layer\n",
    "\n",
    "def esm_embedding(input_file, model, alphabet, nogpu=False, \n",
    "                  repr_layers=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33]):\n",
    "    #model, alphabet = pretrained.load_model_and_alphabet(\"esm1b_t33_650M_UR50S\")\n",
    "    model.eval()\n",
    "    if isinstance(model, MSATransformer):\n",
    "        raise ValueError(\n",
    "            \"This script currently does not handle models with MSA input (MSA Transformer).\"\n",
    "        )\n",
    "    if torch.cuda.is_available() and not nogpu:\n",
    "        model = model.cuda()\n",
    "        print(\"Transferred model to GPU\")\n",
    "\n",
    "    dataset = FastaBatchedDataset.from_file(input_file)\n",
    "    batches = dataset.get_batch_indices(4096, extra_toks_per_seq=1)\n",
    "    data_loader = torch.utils.data.DataLoader(\n",
    "        dataset, collate_fn=alphabet.get_batch_converter(), batch_sampler=batches\n",
    "    )\n",
    "    print(f\"Read input fasta file with {len(dataset)} sequences\")\n",
    "\n",
    "    assert all(-(model.num_layers + 1) <= i <= model.num_layers for i in repr_layers)\n",
    "    repr_layers = [(i + model.num_layers + 1) % (model.num_layers + 1) for i in repr_layers]\n",
    "\n",
    "    result_list = []\n",
    "    esm_names = []\n",
    "    for each in range(len(repr_layers)):\n",
    "      result_list.append([])\n",
    "    # This means inference with maximum performance layer\n",
    "    if len(repr_layers) == 1:\n",
    "      result_list = []\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (labels, strs, toks) in enumerate(data_loader):\n",
    "            print(\n",
    "                f\"Processing {batch_idx + 1} of {len(batches)} batches ({toks.size(0)} sequences)\"\n",
    "            )\n",
    "            if torch.cuda.is_available() and not nogpu:\n",
    "                toks = toks.to(device=\"cuda\", non_blocking=True)\n",
    "\n",
    "            out = model(toks, repr_layers=repr_layers, return_contacts=False)\n",
    "\n",
    "            logits = out[\"logits\"].to(device=\"cpu\")\n",
    "            representations = {\n",
    "                layer: t.to(device=\"cpu\") for layer, t in out[\"representations\"].items()\n",
    "            }\n",
    "\n",
    "            for i, label in enumerate(labels):\n",
    "                esm_names.append(label)\n",
    "                result = {\"label\": label}\n",
    "                # Call clone on tensors to ensure tensors are not views into a larger representation\n",
    "                # See https://github.com/pytorch/pytorch/issues/1995\n",
    "                # Get mean representations\n",
    "                result[\"mean_representations\"] = {\n",
    "                    layer: t[i, 1 : len(strs[i]) + 1].mean(0).clone()\n",
    "                    for layer, t in representations.items()\n",
    "                }\n",
    "                if len(repr_layers) > 1:\n",
    "                  for each in repr_layers:\n",
    "                    result_list[each].append(result[\"mean_representations\"][each].tolist())\n",
    "                else:\n",
    "                  result_list.append(result[\"mean_representations\"][repr_layers[0]].tolist())\n",
    "\n",
    "                \"\"\"\n",
    "                torch.save(\n",
    "                    result,\n",
    "                    args.output_file,\n",
    "                )\n",
    "                \"\"\"\n",
    "    # Output the result\n",
    "    return result_list, esm_names\n",
    "\n",
    "\n",
    "# Code for embedding extraction (prott5)\n",
    "# Extract for every layer\n",
    "# Requires too much resource, make it disable for now (Don't know the specific reason for now)\n",
    "\"\"\"\n",
    "from transformers import T5EncoderModel, T5Tokenizer\n",
    "import gc\n",
    "\n",
    "def prott5_embedding(sequence_list, model, tokenizer):\n",
    "  #tokenizer = T5Tokenizer.from_pretrained(\"Rostlab/prot_t5_xl_uniref50\", do_lower_case=False )\n",
    "  #model = T5EncoderModel.from_pretrained(\"Rostlab/prot_t5_xl_uniref50\")\n",
    "  gc.collect()\n",
    "  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "  model = model.to(device)\n",
    "  model = model.eval()\n",
    "  sequences_Example = [re.sub(r\"[UZOB]\", \"X\", sequence) for sequence in sequence_list]\n",
    "  ids = tokenizer.batch_encode_plus(sequences_Example, add_special_tokens=True, padding=True)\n",
    "  input_ids = torch.tensor(ids['input_ids']).to(device)\n",
    "  attention_mask = torch.tensor(ids['attention_mask']).to(device)\n",
    "  with torch.no_grad():\n",
    "    embedding = model(input_ids=input_ids,attention_mask=attention_mask)\n",
    "  embedding = embedding.last_hidden_state.cpu().numpy()\n",
    "  features = [] \n",
    "  for seq_num in range(len(embedding)):\n",
    "      seq_len = (attention_mask[seq_num] == 1).sum()\n",
    "      seq_emd = embedding[seq_num][:seq_len-1]\n",
    "      features.append(seq_emd)\n",
    "  return features\n",
    "\"\"\"\n",
    "\n",
    "def sort_criterion(elem):\n",
    "    return elem[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "VEGk8NKqDaUo"
   },
   "outputs": [],
   "source": [
    "#@title Compute Embeddings of input sequences from ESM-1b\n",
    "#@markdown - esm-1b only takes in sequences shorter than 1024 amino acids<br>Proteins longer than 1024 amino acids will be excluded\n",
    "#model_name = \"esm-1b, prott5\" #@param [\"esm-1b\", \"prott5\", \"esm-1b, prott5\"]\n",
    "#model_name = \"esm-1b\"\n",
    "esm1b_query_embedding = None\n",
    "prott5_query_embedding = None\n",
    "# Run esm-1b\n",
    "#if model_name == \"esm-1b, prott5\" or model_name == \"esm-1b\":\n",
    "#model_esm, alphabet_esm = pretrained.load_model_and_alphabet(\"esm1b_t33_650M_UR50S\")\n",
    "esm1b_query_embedding, esm1b_names = esm_embedding(\"/content/esm1b_input.fasta\", model_esm, alphabet_esm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "jrC6sBjlWUuO"
   },
   "outputs": [],
   "source": [
    "#@title Five Fold Validation with Machine Learning\n",
    "#@markdown - Compute the layer showing best performance\n",
    "n_neighbors = min(n_neighbors, cnt)\n",
    "regressor = None\n",
    "if Regressor == \"Random Forest Regressor\":\n",
    "  regressor = RandomForestRegressor()\n",
    "elif Regressor == \"Linear Regression\":\n",
    "  regressor = LinearRegression()\n",
    "elif Regressor == \"Logistic Regression\":\n",
    "  regressor = LogisticRegression()\n",
    "elif Regressor == \"K Neighbors Regressor\":\n",
    "  regressor = KNeighborsRegressor(n_neighbors=n_neighbors)\n",
    "\n",
    "# Let's do five fold validation!\n",
    "# Don't do five fold validation if the train data size is smaller than 10 (quit & make the max_layer = 33)\n",
    "max_layer = None\n",
    "if do_distance:\n",
    "  max_layer = len(esm1b_query_embedding) - 1\n",
    "  regressor = KNeighborsRegressor(n_neighbors=1)\n",
    "elif cnt < 10:\n",
    "  warnings.warn(\"Warning: Size of data smaller than 10, skipping five fold validation, last layer of ESM-1b is chosen.\")\n",
    "  max_layer = len(esm1b_query_embedding) - 1\n",
    "else:\n",
    "  # Training data size is big enough. Let's do the five fold validation for {iteration} times!\n",
    "  ff_average_esm = []\n",
    "  #ff_average_t5 = []\n",
    "  #for i in range(len(t5_embeddings_all_layers)):\n",
    "    #ff_average_t5.append([])\n",
    "  for i in range(len(esm1b_query_embedding)):\n",
    "    ff_average_esm.append([])\n",
    "  for ff in range(iteration):\n",
    "    correlations_esm=[]\n",
    "    mean_esm = []\n",
    "    std_esm = []\n",
    "    for layer in range(len(esm1b_query_embedding)):\n",
    "      esm_embeddings = esm1b_query_embedding[layer]\n",
    "      #pair the embedding, kcat, kcat_km and random shuffle, split into training and test set (esm)\n",
    "      all_data = []\n",
    "      for i in range(len(esm_embeddings)):\n",
    "        all_data.append([esm_embeddings[i], name_value[esm1b_names[i]]])\n",
    "      seed = np.random.randint(1,100000)\n",
    "      np.random.seed(seed)\n",
    "      np.random.shuffle(all_data)\n",
    "\n",
    "      #Actual Five fold validation\n",
    "      pearson_r = []\n",
    "      for iter in range(5):\n",
    "        training_set = []\n",
    "        test_set = []\n",
    "        for i in range(len(all_data)):\n",
    "          if i % 5 == iter:\n",
    "            test_set.append(all_data[i])\n",
    "          else:\n",
    "            training_set.append(all_data[i])\n",
    "\n",
    "        train_x = []\n",
    "        train_y = []\n",
    "        for i in range(len(training_set)):\n",
    "          train_x.append(training_set[i][0])\n",
    "          train_y.append(training_set[i][1])\n",
    "\n",
    "        test_x = []\n",
    "        test_y = []\n",
    "        for i in range(len(test_set)):\n",
    "          test_x.append(test_set[i][0])\n",
    "          test_y.append(test_set[i][1])\n",
    "        regressor.fit(train_x, train_y)\n",
    "        test_result = regressor.predict(test_x)\n",
    "        pearson_r.append(np.corrcoef(np.array(test_y), np.array(test_result))[0, 1])\n",
    "      pearson_r = np.array(pearson_r)\n",
    "      for cor in pearson_r:\n",
    "        correlations_esm.append(cor)\n",
    "      mean_esm.append(np.mean(pearson_r))\n",
    "      ff_average_esm[layer].append(np.mean(pearson_r))\n",
    "      std_esm.append(np.std(pearson_r))\n",
    "      #print(\"Layer: \", layer)\n",
    "      #print(\"All cor: \", pearson_r, \" mean: \", np.mean(pearson_r), \" std: \", np.std(pearson_r))\n",
    "    x_t5 = []\n",
    "    x_error_t5 = []\n",
    "    x_esm = []\n",
    "    x_error_esm = []\n",
    "    #num = 0\n",
    "    #for i in range(len(t5_embeddings_all_layers)):\n",
    "    #  x_error_t5.append(num)\n",
    "    #  for j in range(5):\n",
    "    #    x_t5.append(num)\n",
    "    #  num += 1\n",
    "    num = 0\n",
    "    for i in range(len(esm1b_query_embedding)):\n",
    "      x_error_esm.append(num)\n",
    "      for j in range(5):\n",
    "        x_esm.append(num)\n",
    "      num += 1\n",
    "    \"\"\"\n",
    "    fig1 = plt.figure(2 * ff)\n",
    "    plt.scatter(x_t5, correlations_t5)\n",
    "    plt.errorbar(x_error_t5, mean_t5, std_t5, linestyle='None', marker='^', ecolor='red', mfc='red', mec='red', ms=10)\n",
    "    plt.title(\"Prottrans t5\")\n",
    "    plt.ylim(0,1)\n",
    "    plt.xlabel(\"Layer\")\n",
    "    plt.ylabel(\"Correlation\")\n",
    "    \"\"\"\n",
    "    fig2 = plt.figure(2 * ff + 1)\n",
    "    plt.scatter(x_esm, correlations_esm)\n",
    "    plt.errorbar(x_error_esm, mean_esm, std_esm, linestyle='None', marker='^', ecolor='red', mfc='red', mec='red', ms=10)\n",
    "    plt.title(\"esm1b\")\n",
    "    plt.ylim(0,1)\n",
    "    plt.xlabel(\"Layer\")\n",
    "    plt.ylabel(\"Correlation\")\n",
    "    plt.show()\n",
    "    print(\"%dth iteration\"%(ff + 1))\n",
    "\n",
    "  #get overall average of five trials, and report the best layer\n",
    "  #ff_average_t5 = np.array(ff_average_t5)\n",
    "  ff_average_esm = np.array(ff_average_esm)\n",
    "  #ff_average_t5_final = np.mean(ff_average_t5, axis=1)\n",
    "  ff_average_esm_final = np.mean(ff_average_esm, axis=1)\n",
    "  #print(\"ff_average_t5 max: \", np.max(ff_average_t5_final), \", argmax: \", np.argmax(ff_average_t5_final))\n",
    "  print(\"ff_average_esm max: \", np.max(ff_average_esm_final), \", argmax: \", np.argmax(ff_average_esm_final))\n",
    "  max_layer = np.argmax(ff_average_esm_final)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "vsyiUAZMWI2p"
   },
   "outputs": [],
   "source": [
    "#@title Search query sequences against Colabfold database and report the top ranking targets\n",
    "amino_acid = [\"G\", \"A\", \"L\", \"M\", \"K\", \"F\", \"W\", \"Q\", \"E\", \"S\", \"P\", \"V\", \"I\", \"C\", \"Y\", \"H\", \"R\", \"N\", \"D\", \"T\"]\n",
    "#target_number = 20 #@param = {type:\"integer\"}\n",
    "m8_output = run_mmseqs2(sequences, \"tmp\")\n",
    "with open(\"/content/esm1b_target.m8\", 'w') as new:\n",
    "  for each in m8_output:\n",
    "    for line in each.split(\"\\n\"):\n",
    "      if line != \"\":\n",
    "        seq = line.split()[-1].upper()\n",
    "        new_seq = \"\"\n",
    "        for char in seq:\n",
    "          if char in amino_acid:\n",
    "            new_seq += char\n",
    "        new.write(\"%s %s\\n\"%(line.split()[1], new_seq))\n",
    "!sort -u -k2,2 esm1b_target.m8 | sort -k1,1 | awk 'BEGIN {cnt=2;name=\"\";} {if (name == $1) {print $1\"_colab_\"cnt, $2;cnt += 1} else {print;name=$1;cnt=2}}' > esm1b_nonredundant.m8\n",
    "\n",
    "# Gather the search result and make input file for esm1b\n",
    "#searched_sequences = []\n",
    "#searched_names = []\n",
    "with open(\"/content/esm1b_nonredundant.m8\") as f:\n",
    "  with open(\"/content/esm1b_target_input.fasta\", 'w') as new:\n",
    "    line = f.readline()\n",
    "    while line:\n",
    "      if len(line[:-1].split()[1]) < 1023:\n",
    "        new.write(\">%s\\n%s\\n\"%(line.split()[0], line[:-1].split()[1].upper()))\n",
    "      line = f.readline()\n",
    "#assert len(searched_names) == len(searched_sequences)\n",
    "\n",
    "# Get the embeddings of the search result\n",
    "target_embeddings, target_names = esm_embedding(\"/content/esm1b_target_input.fasta\", model_esm, alphabet_esm, repr_layers=[max_layer])\n",
    "\n",
    "# Retrain the regressor with all input data\n",
    "esm_embeddings = esm1b_query_embedding[max_layer]\n",
    "all_x = []\n",
    "all_y = []\n",
    "for i in range(len(esm_embeddings)):\n",
    "  all_x.append(esm_embeddings[i])\n",
    "  all_y.append(name_value[esm1b_names[i]])\n",
    "regressor.fit(all_x, all_y)\n",
    "\n",
    "# Get the prediction with regressor and sort\n",
    "target_result = []\n",
    "if do_distance:\n",
    "  target_results = regressor.kneighbors(target_embeddings, n_neighbors=cnt)\n",
    "  target_dist = np.sum(target_results[0], axis = 1)\n",
    "  for i in range(len(target_dist)):\n",
    "    target_result.append([target_names[i], target_dist[i]])\n",
    "  target_result.sort(key = sort_criterion)\n",
    "else:\n",
    "  target_results = regressor.predict(target_embeddings)\n",
    "  for i in range(len(target_results)):\n",
    "      target_result.append([target_names[i], target_results[i]])\n",
    "  target_result.sort(key=sort_criterion)\n",
    "  target_result.reverse()\n",
    "\n",
    "\n",
    "# Save the result\n",
    "with open(\"/content/top_ranked_target_tmp.tsv\",'w') as new:\n",
    "  for i in range(len(target_result)):\n",
    "    new.write(\"%s\\t%f\\n\"%(target_result[i][0], target_result[i][1]))\n",
    "# Get the sequence also\n",
    "!awk 'BEGIN {OFS = \"\\t\";} NR==FNR {if (substr($1, 1, 1) == \">\") {name = substr($1, 2, length($1) - 1);getline;a[name]=$0;next;}} {if ($1 in a) {print $1, $2, a[$1]} else {print error}}' esm1b_target_input.fasta top_ranked_target_tmp.tsv > top_ranked_target_list.tsv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "nw3wT9472yk_"
   },
   "outputs": [],
   "source": [
    "#@title Download the result file\n",
    "#@markdown there is only one result file:<br>\n",
    "#@markdown - <b>top_ranked_target_list.tsv</b>, containing names, predicted values, and sequences of the targets on each column respectively, sorted by predicted values.<br>\n",
    "#@markdown - Result file might have names containing \"\\_colab_%d\" at the end of the name. This happens when there are more than one sequences with same name.\n",
    "\n",
    "#!zip -FSr enzyme_colabfold_search.result.zip \"top_ranked_target.list\"\n",
    "files.download(\"top_ranked_target_list.tsv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ajfvEIZCBSsS"
   },
   "source": [
    "##What's happening in here\n",
    "1.   Compute the embeddings of every layer of [ESM-1b](https://github.com/facebookresearch/esm) for each input sequence\n",
    "2.   Conduct five fold validations n times with a regressor set by user (five and random forest regressor for default respectively), and choose the layer that shows the best performance\n",
    "3.   Search through [colabfold database](https://colabfold.mmseqs.com/), using the input sequences as query\n",
    "4.   Compute the embeddings of selected layer of ESM-1b for each search result\n",
    "5.   Compute the predicted numerical values for each search result using a regressor trained by all input sequences\n",
    "6.   Rank the search results by the predicted numerical value and download the result file (see the very bottom of the colab for detail about result file)\n",
    "\n",
    "##To see the github, please visit:\n",
    "https://github.com/steineggerlab/SeekRank\n",
    "\n",
    "##Pulibcation\n",
    "Eom, Hyunuk, et al. \"Discovery of highly active kynureninases for cancer immunotherapy through protein language model.\" Nucleic Acids Research 53.1 (2025), doi: https://doi.org/10.1093/nar/gkae1245"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}