IgLM.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "5AEsv1z5KXXA"
      },
      "source": [
        "# **IgLM**: Generative language modeling for antibody design\n",
        "\n",
        "Official notebook for [IgLM](https://www.biorxiv.org/content/10.1101/2021.12.13.472419v2), a generative language model for antibody sequence generation and infilling.  The code, data, and weights for this work are made available for non-commercial use. For commercial inquiries, please contact `dmalon11[at]jhu.edu`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0PsLNGK57LDq"
      },
      "outputs": [],
      "source": [
        "#@title Configure generation settings then press `Runtime` -> `Run all`\n",
        "\n",
        "import os\n",
        "import sys\n",
        "import uuid\n",
        "\n",
        "python_version = f\"{sys.version_info.major}.{sys.version_info.minor}\"\n",
        "\n",
        "job_name = uuid.uuid4().hex[:6]\n",
        "pred_dir = job_name\n",
        "os.makedirs(pred_dir, exist_ok=True)\n",
        "\n",
        "species = 'human' #@param [\"human\", \"mouse\", \"camel\", \"rat\", \"rabbit\", \"rhesus\"]\n",
        "chain = 'heavy' #@param [\"heavy\", \"light\"]\n",
        "\n",
        "#@markdown IgLM can be used to generate full-length antibody sequences or infill regions of existing sequences.\n",
        "mode = 'generate' #@param [\"generate\", \"infill\"]\n",
        "num_sequences = 20 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Predict structures with IgFold.\n",
        "predict_structures = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown **Full-length generation**\n",
        "\n",
        "#@markdown Initial residues may be provided to prompt sequence generation. This is useful to ensure full-length sequences are generated.\n",
        "use_prompt = False #@param {type:\"boolean\"}\n",
        "prompt = \"EVQ\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown **Sequence infilling**\n",
        "\n",
        "#@markdown For sequence infilling, such as CDR loop generation, a parent sequence must be provided with specific residue indices to infill.\n",
        "\n",
        "parent_sequence = \"EVQLVESGGGLVQPGGSLRLSCAASGFNIKEYYMHWVRQAPGKGLEWVGLIDPEQGNTIYDPKFQDRATISADNSKNTAYLQMNSLRAEDTAVYYCARDTAAYFDYWGQGTLVTVS\" #@param {type:\"string\"}\n",
        "infill_start = 98 #@param {type:\"integer\"}\n",
        "infill_end = 106 #@param {type:\"integer\"}\n",
        "\n",
        "print(f\"Running job {job_name}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "a2a3BsiE9AXI"
      },
      "outputs": [],
      "source": [
        "#@title Install dependencies\n",
        "\n",
        "PYTHON_VERSION = python_version\n",
        "\n",
        "torch_string = \"torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html\"\n",
        "\n",
        "if not os.path.isfile(\"IGLM_READY\"):\n",
        "  print(\"installing iglm...\")\n",
        "  os.system(f\"pip3 install {torch_string}\")\n",
        "  os.system(f\"pip install 'iglm' {torch_string}\")\n",
        "  os.system(\"pip install -q --no-warn-conflicts 'py3Dmol>=2.0.1' matplotlib seaborn\")\n",
        "  os.system(\"touch IGLM_READY\")\n",
        "\n",
        "if predict_structures and not os.path.isfile(\"IGFOLD_READY\"):\n",
        "  print(\"installing igfold...\")\n",
        "  os.system(f\"pip install 'igfold>=0.3.0' {torch_string}\")\n",
        "  os.system(\"touch IGFOLD_READY\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xFOTYxsP9Cz1"
      },
      "outputs": [],
      "source": [
        "#@title Generate sequences\n",
        "\n",
        "from iglm import IgLM\n",
        "from tqdm import tqdm\n",
        "\n",
        "species_tokens = {\n",
        "    \"human\": \"[HUMAN]\",\n",
        "    \"mouse\": \"[MOUSE]\",\n",
        "    \"camel\": \"[CAMEL]\",\n",
        "    \"rat\": \"[RAT]\",\n",
        "    \"rabbit\": \"[RABBIT]\",\n",
        "    \"rhesus\": \"[RHESUS]\",\n",
        "}\n",
        "chain_tokens = {\n",
        "    \"heavy\": \"[HEAVY]\",\n",
        "    \"light\": \"[LIGHT]\",\n",
        "}\n",
        "\n",
        "species_token = species_tokens[species]\n",
        "chain_token = chain_tokens[chain]\n",
        "\n",
        "iglm = IgLM()\n",
        "if mode == \"generate\":\n",
        "  if use_prompt and len(prompt) > 0:\n",
        "    prompt = prompt.upper()\n",
        "  else:\n",
        "    prompt = None\n",
        "  generated_seqs = iglm.generate(\n",
        "      chain_token,\n",
        "      species_token,\n",
        "      prompt_sequence=prompt,\n",
        "      num_to_generate=num_sequences,\n",
        "  )\n",
        "elif mode == \"infill\":\n",
        "  generated_seqs = iglm.infill(\n",
        "      parent_sequence,\n",
        "      chain_token,\n",
        "      species_token,\n",
        "      infill_range=(infill_start, infill_end),\n",
        "      num_to_generate=num_sequences,\n",
        "  )\n",
        "\n",
        "print(\"\\nSequences\")\n",
        "for i, sequence in enumerate(generated_seqs):\n",
        "  print(f\"{str(i).zfill(num_sequences // 10)}\\t{sequence}\")\n",
        "\n",
        "fasta_file = os.path.join(pred_dir, \"all_sequences.fasta\")\n",
        "chain_id = \"H\" if chain == \"heavy\" else \"L\"\n",
        "with open(fasta_file, \"w\") as f:\n",
        "  for i, sequence in enumerate(generated_seqs):\n",
        "    f.write(f\">{job_name}_{str(i).zfill(num_sequences // 10)}\\n{sequence}\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ajyElWbZ9EFF"
      },
      "outputs": [],
      "source": [
        "#@title Predict structures\n",
        "\n",
        "if predict_structures:\n",
        "  import contextlib\n",
        "  with open(os.devnull, 'w') as devnull:\n",
        "    with contextlib.redirect_stdout(devnull):\n",
        "      from igfold import IgFoldRunner\n",
        "\n",
        "      igfold = IgFoldRunner(num_models=1)\n",
        "\n",
        "      for i, sequence in tqdm(enumerate(generated_seqs), total=num_sequences):\n",
        "        pdb_file = os.path.join(pred_dir, f\"{job_name}_{str(i).zfill(num_sequences // 10)}.pdb\")\n",
        "        fasta_file = pdb_file.replace(\"pdb\", \"fasta\")\n",
        "        igfold.fold(\n",
        "            pdb_file,\n",
        "            sequences={chain_id: sequence},\n",
        "            do_refine=False,\n",
        "            do_renum=False,\n",
        "        )\n",
        "        os.system(f\"rm {fasta_file}\")\n",
        "else:\n",
        "  print(\"Not predicting structures\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gZBzjpMdJ77q"
      },
      "outputs": [],
      "source": [
        "#@title Download results\n",
        "\n",
        "#@markdown Download zip file containing generated sequences (and structures). If download fails, results are also accessible from file explorer on the left panel of the notebook.\n",
        "\n",
        "from google.colab import files\n",
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
        "\n",
        "!zip -FSr $job_name\".result.zip\" $pred_dir/ &> /dev/null\n",
        "files.download(f\"{job_name}.result.zip\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "IgLM.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.9.12 ('iglm')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.12"
    },
    "vscode": {
      "interpreter": {
        "hash": "84181e5f1827f203c248bfcd3a60e7e3a4ffc08f0a7dd8a443bd855d4ab14b5d"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}