load_pLM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b35553c4",
   "metadata": {},
   "source": [
    "# Loading pLMs from Wilke Lab Shared Folder  \n",
    "This repository is intended to facilitate the use of pLM models without using up all the space in your home directory.   \n",
    "\n",
    "All models are located at:   \n",
    "/stor/work/Wilke/wilkelab/pLMs_checkpoints/   \n",
    "\n",
    "\n",
    "NOTE: All these models are not in eval() mode by default. So if you are loading them for inference, DO NOT forget to do the following:   \n",
    "``` \n",
    "model.eval()\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc236e88",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34m/stor/work/Wilke/wilkelab/pLMs_checkpoints/\u001b[00m\n",
      "├── \u001b[01;34mAMPLIFY\u001b[00m\n",
      "│   ├── \u001b[01;34mAMPLIFY_120M\u001b[00m\n",
      "│   │   ├── amplify.py\n",
      "│   │   ├── config.json\n",
      "│   │   ├── config.yaml\n",
      "│   │   ├── model.safetensors\n",
      "│   │   ├── pytorch_model.pt\n",
      "│   │   ├── rmsnorm.py\n",
      "│   │   ├── rotary.py\n",
      "│   │   ├── special_tokens_map.json\n",
      "│   │   ├── tokenizer_config.json\n",
      "│   │   └── tokenizer.json\n",
      "│   └── \u001b[01;34mAMPLIFY_350M\u001b[00m\n",
      "│       ├── amplify.py\n",
      "│       ├── config.json\n",
      "│       ├── model.safetensors\n",
      "│       ├── README.md\n",
      "│       ├── rmsnorm.py\n",
      "│       ├── rotary.py\n",
      "│       ├── special_tokens_map.json\n",
      "│       ├── tokenizer_config.json\n",
      "│       └── tokenizer.json\n",
      "├── \u001b[01;34mESM1\u001b[00m\n",
      "│   ├── esm1b_t33_650M_UR50S-contact-regression.pt\n",
      "│   └── esm1b_t33_650M_UR50S.pt\n",
      "├── \u001b[01;34mESM2\u001b[00m\n",
      "│   ├── esm2_t12_35M_UR50D-contact-regression.pt\n",
      "│   ├── esm2_t12_35M_UR50D.pt\n",
      "│   ├── esm2_t30_150M_UR50D-contact-regression.pt\n",
      "│   ├── esm2_t30_150M_UR50D.pt\n",
      "│   ├── esm2_t33_650M_UR50D-contact-regression.pt\n",
      "│   ├── esm2_t33_650M_UR50D.pt\n",
      "│   ├── esm2_t36_3B_UR50D-contact-regression.pt\n",
      "│   ├── esm2_t36_3B_UR50D.pt\n",
      "│   ├── esm2_t48_15B_UR50D-contact-regression.pt\n",
      "│   ├── esm2_t48_15B_UR50D.pt\n",
      "│   ├── esm2_t6_8M_UR50D-contact-regression.pt\n",
      "│   └── esm2_t6_8M_UR50D.pt\n",
      "└── \u001b[01;34mESMC\u001b[00m\n",
      "    ├── esmc_300m_2024_12_v0.pth\n",
      "    └── esmc_600m_2024_12_v0.pth\n",
      "\n",
      "6 directories, 35 files\n"
     ]
    }
   ],
   "source": [
    "!tree /stor/work/Wilke/wilkelab/pLMs_checkpoints/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f7b183b",
   "metadata": {},
   "source": [
    "# Emptying massive space used by ESM2 3B and 15B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad77a2f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Models downloaded from Torch Hub (via ESM2 scripts) or Hugging Face will be stored in the cache.\n",
    "# Uncomment the following lines to clear the cache if needed:\n",
    "\n",
    "#rm -rf .cache/torch/hub/checkpoints/\n",
    "\n",
    "#rm -rf .cache/huggingface/hub/\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58b6f33b",
   "metadata": {},
   "source": [
    "# Loading ESM2 models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "72ad4f69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ESM2(\n",
       "  (embed_tokens): Embedding(33, 480, padding_idx=1)\n",
       "  (layers): ModuleList(\n",
       "    (0-11): 12 x TransformerLayer(\n",
       "      (self_attn): MultiheadAttention(\n",
       "        (k_proj): Linear(in_features=480, out_features=480, bias=True)\n",
       "        (v_proj): Linear(in_features=480, out_features=480, bias=True)\n",
       "        (q_proj): Linear(in_features=480, out_features=480, bias=True)\n",
       "        (out_proj): Linear(in_features=480, out_features=480, bias=True)\n",
       "        (rot_emb): RotaryEmbedding()\n",
       "      )\n",
       "      (self_attn_layer_norm): LayerNorm((480,), eps=1e-05, elementwise_affine=True)\n",
       "      (fc1): Linear(in_features=480, out_features=1920, bias=True)\n",
       "      (fc2): Linear(in_features=1920, out_features=480, bias=True)\n",
       "      (final_layer_norm): LayerNorm((480,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (contact_head): ContactPredictionHead(\n",
       "    (regression): Linear(in_features=240, out_features=1, bias=True)\n",
       "    (activation): Sigmoid()\n",
       "  )\n",
       "  (emb_layer_norm_after): LayerNorm((480,), eps=1e-05, elementwise_affine=True)\n",
       "  (lm_head): RobertaLMHead(\n",
       "    (dense): Linear(in_features=480, out_features=480, bias=True)\n",
       "    (layer_norm): LayerNorm((480,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import esm\n",
    "import torch\n",
    "\n",
    "model_path = '/stor/work/Wilke/wilkelab/pLMs_checkpoints/'\n",
    "\n",
    "supported_models = [\n",
    "            'esm2_t6_8M_UR50D', 'esm2_t12_35M_UR50D', 'esm2_t30_150M_UR50D', \n",
    "            'esm2_t33_650M_UR50D', 'esm2_t36_3B_UR50D', 'esm2_t48_15B_UR50D'\n",
    "            ]\n",
    "\n",
    "# seting a device\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Choose from the list of models above\n",
    "model_name = 'esm2_t12_35M_UR50D'  \n",
    "\n",
    "checkpoint_path = os.path.join(model_path, 'ESM2', model_name + '.pt')\n",
    "\n",
    "model, alphabet = esm.pretrained.load_model_and_alphabet(checkpoint_path)\n",
    "\n",
    "\n",
    "# Load the model to the device\n",
    "model = model.to(device)\n",
    "\n",
    "#if loading the model for inference, such as extracting embeddings, you can set the model to evaluation mode:\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33cd192a",
   "metadata": {},
   "source": [
    "## Example of use os ESM2 models:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "498b525f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extracting embeddings using the ESM2 script available in the ESM2 repository\n",
    "\n",
    "model_full_path = os.path.join(model_path, 'ESM2', 'esm2_t6_8M_UR50D.pt')\n",
    "\n",
    "#python scripts/extract.py model_full_path data/prot_seqs.fasta embeddings/esm2_650M/prot_seqs/ --repr_layers 6 --include mean"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5075496",
   "metadata": {},
   "source": [
    "# Loading ESMC models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef23c0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from esm.models.esmc import ESMC\n",
    "from esm.tokenization import get_esmc_model_tokenizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dd8b5985",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the models locally\n",
    "def ESMC_300M_202412(model_path: str, device: torch.device | str = \"cpu\"):\n",
    "    with torch.device(device):\n",
    "        model = ESMC(\n",
    "            d_model=960, n_heads=15, n_layers=30, tokenizer=get_esmc_model_tokenizers()\n",
    "        )\n",
    "    state_dict = torch.load(model_path, map_location=device, weights_only=True)\n",
    "    model.load_state_dict(state_dict)\n",
    "    # Convert model parameters to torch.bfloat16 or torch.float32\n",
    "    model = model.to(torch.float32)\n",
    "    return model\n",
    "\n",
    "\n",
    "def ESMC_600M_202412(model_path: str, device: torch.device | str = \"cpu\"):\n",
    "    with torch.device(device):\n",
    "        model = ESMC(\n",
    "            d_model=1152, n_heads=18, n_layers=36, tokenizer=get_esmc_model_tokenizers()\n",
    "        )\n",
    "    state_dict = torch.load(model_path, map_location=device, weights_only=True)\n",
    "    model.load_state_dict(state_dict)\n",
    "    # Convert model parameters to float32\n",
    "    model = model.to(torch.float32)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "feb88d56",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ESMC(\n",
       "  (embed): Embedding(64, 960)\n",
       "  (transformer): TransformerStack(\n",
       "    (blocks): ModuleList(\n",
       "      (0-29): 30 x UnifiedTransformerBlock(\n",
       "        (attn): MultiHeadAttention(\n",
       "          (layernorm_qkv): Sequential(\n",
       "            (0): LayerNorm((960,), eps=1e-05, elementwise_affine=True)\n",
       "            (1): Linear(in_features=960, out_features=2880, bias=False)\n",
       "          )\n",
       "          (out_proj): Linear(in_features=960, out_features=960, bias=False)\n",
       "          (q_ln): LayerNorm((960,), eps=1e-05, elementwise_affine=True)\n",
       "          (k_ln): LayerNorm((960,), eps=1e-05, elementwise_affine=True)\n",
       "          (rotary): RotaryEmbedding()\n",
       "        )\n",
       "        (ffn): Sequential(\n",
       "          (0): LayerNorm((960,), eps=1e-05, elementwise_affine=True)\n",
       "          (1): Linear(in_features=960, out_features=5120, bias=False)\n",
       "          (2): SwiGLU()\n",
       "          (3): Linear(in_features=2560, out_features=960, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (norm): LayerNorm((960,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (sequence_head): Sequential(\n",
       "    (0): Linear(in_features=960, out_features=960, bias=True)\n",
       "    (1): GELU(approximate='none')\n",
       "    (2): LayerNorm((960,), eps=1e-05, elementwise_affine=True)\n",
       "    (3): Linear(in_features=960, out_features=64, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_name = 'esmc_300m_2024_12_v0.pth'  \n",
    "checkpoint_path = os.path.join(path, 'ESMC', model_name)\n",
    "\n",
    "model = ESMC_300M_202412(checkpoint_path)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c25ce6da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 5, 2],\n",
       "        [0, 5, 2],\n",
       "        [0, 5, 2],\n",
       "        [0, 5, 2],\n",
       "        [0, 5, 2],\n",
       "        [0, 5, 2],\n",
       "        [0, 5, 2]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use the model tokenizer\n",
    "model._tokenize(\"AAAAAAA\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "827e2850",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ESMC(\n",
       "  (embed): Embedding(64, 1152)\n",
       "  (transformer): TransformerStack(\n",
       "    (blocks): ModuleList(\n",
       "      (0-35): 36 x UnifiedTransformerBlock(\n",
       "        (attn): MultiHeadAttention(\n",
       "          (layernorm_qkv): Sequential(\n",
       "            (0): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "            (1): Linear(in_features=1152, out_features=3456, bias=False)\n",
       "          )\n",
       "          (out_proj): Linear(in_features=1152, out_features=1152, bias=False)\n",
       "          (q_ln): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "          (k_ln): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "          (rotary): RotaryEmbedding()\n",
       "        )\n",
       "        (ffn): Sequential(\n",
       "          (0): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "          (1): Linear(in_features=1152, out_features=6144, bias=False)\n",
       "          (2): SwiGLU()\n",
       "          (3): Linear(in_features=3072, out_features=1152, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (norm): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (sequence_head): Sequential(\n",
       "    (0): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "    (1): GELU(approximate='none')\n",
       "    (2): LayerNorm((1152,), eps=1e-05, elementwise_affine=True)\n",
       "    (3): Linear(in_features=1152, out_features=64, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_name = 'esmc_600m_2024_12_v0.pth'  \n",
    "checkpoint_path = os.path.join(path, 'ESMC', model_name)\n",
    "\n",
    "model = ESMC_600M_202412(checkpoint_path)\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53731bf9",
   "metadata": {},
   "source": [
    "# Loading AMPLIFY Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07d59736",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModel, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "c99e04d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def AMPLIFY(model_checkpoint):\n",
    "    model = AutoModel.from_pretrained(model_checkpoint, trust_remote_code=True)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, trust_remote_code=True)\n",
    "    return model, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3578a8b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AMPLIFY(\n",
       "  (encoder): Embedding(27, 640, padding_idx=0)\n",
       "  (transformer_encoder): ModuleList(\n",
       "    (0-23): 24 x EncoderBlock(\n",
       "      (q): Linear(in_features=640, out_features=640, bias=False)\n",
       "      (k): Linear(in_features=640, out_features=640, bias=False)\n",
       "      (v): Linear(in_features=640, out_features=640, bias=False)\n",
       "      (wo): Linear(in_features=640, out_features=640, bias=False)\n",
       "      (resid_dropout): Dropout(p=0, inplace=False)\n",
       "      (ffn): SwiGLU(\n",
       "        (w12): Linear(in_features=640, out_features=3424, bias=False)\n",
       "        (w3): Linear(in_features=1712, out_features=640, bias=False)\n",
       "      )\n",
       "      (attention_norm): RMSNorm()\n",
       "      (ffn_norm): RMSNorm()\n",
       "      (ffn_dropout): Dropout(p=0, inplace=False)\n",
       "    )\n",
       "  )\n",
       "  (layer_norm_2): RMSNorm()\n",
       "  (decoder): Linear(in_features=640, out_features=27, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_name = 'AMPLIFY_120M/'\n",
    "checkpoint_path = os.path.join(path, 'AMPLIFY', model_name)\n",
    "model, tokenizer = AMPLIFY(model_checkpoint)\n",
    "model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}