synced oss and fixed pyright issues

Author: imathur1

cookbook/local/open_generate.ipynb

@@ -38,6 +38,7 @@
     "\n",
     "!pip install py3Dmol\n",
     "import py3Dmol\n",
+    "\n",
     "from esm.models.esm3 import ESM3\n",
     "from esm.sdk.api import ESMProtein, GenerationConfig\n",
     "from esm.utils.structure.protein_chain import ProteinChain"

cookbook/local/raw_forwards.py

@@ -13,9 +13,7 @@
 from esm.tokenization.function_tokenizer import (
     InterProQuantizedTokenizer as EsmFunctionTokenizer,
 )
-from esm.tokenization.sequence_tokenizer import (
-    EsmSequenceTokenizer,
-)
+from esm.tokenization.sequence_tokenizer import EsmSequenceTokenizer
 from esm.utils.structure.protein_chain import ProteinChain
 from esm.utils.types import FunctionAnnotation
 

cookbook/tutorials/1_esmprotein.ipynb

@@ -72,6 +72,7 @@
    "outputs": [],
    "source": [
     "from biotite.database import rcsb\n",
+    "\n",
     "from esm.sdk.api import ESMProtein\n",
     "from esm.utils.structure.protein_chain import ProteinChain\n",
     "from esm.utils.types import FunctionAnnotation\n",
@@ -496,9 +497,10 @@
     "# Functions for visualizing InterPro function annotations\n",
     "\n",
     "from dna_features_viewer import GraphicFeature, GraphicRecord\n",
-    "from esm.utils.function.interpro import InterPro, InterProEntryType\n",
     "from matplotlib import colormaps\n",
     "\n",
+    "from esm.utils.function.interpro import InterPro, InterProEntryType\n",
+    "\n",
     "\n",
     "def visualize_function_annotations(\n",
     "    annotations: list[FunctionAnnotation],\n",

cookbook/tutorials/3_gfp_design.ipynb

@@ -64,6 +64,7 @@
     "import matplotlib.pyplot as pl\n",
     "import py3Dmol\n",
     "import torch\n",
+    "\n",
     "from esm.sdk import client\n",
     "from esm.sdk.api import ESMProtein, GenerationConfig\n",
     "from esm.utils.structure.protein_chain import ProteinChain"

cookbook/tutorials/4_forge_generate.ipynb

@@ -36,6 +36,7 @@
     "\n",
     "!pip install py3Dmol\n",
     "import py3Dmol\n",
+    "\n",
     "from esm.sdk import client\n",
     "from esm.sdk.api import ESMProtein, GenerationConfig\n",
     "from esm.utils.structure.protein_chain import ProteinChain"

cookbook/tutorials/5_guided_generation.ipynb

@@ -14,13 +14,13 @@
     "3. Minimize a biophysical energy function\n",
     "4. Use experimental screening data to guide designs with a regression model\n",
     "\n",
-    "As long as your scoring function takes a protein as input and outputs a single score, you can use it to guide designs. To accomplish this, we use an implementation of derivative-free guidance inspired by Soft Value-Based Decoding described in [Li, et al 2024](https://arxiv.org/abs/2408.08252).\n",
+    "As long as your scoring function takes a protein as input and outputs a single score, you can use it to guide designs. To accomplish this, we use an implementation of derivative-free guidance inspired by Soft Value-Based Decoding described in [Li, et al 2024](https://arxiv.org/abs/2408.08252) and constrained optimization using the Modified Differential Method of Multipliers from [Platt & Barr 1987](https://proceedings.neurips.cc/paper_files/paper/1987/file/a1126573153ad7e9f44ba80e99316482-Paper.pdf)\n",
     "\n",
     "In this notebook we will walk through a few examples to illustrate how to use guided generation. \n",
     "\n",
     "1. Guide towards high pTM for improved generation quality\n",
     "2. Generate a protein with no cysteine (C) residues\n",
-    "3. Maximize protein globularity by minimizing the radius of gyration\n",
+    "3. Maximize protein globularity by minimizing the radius of gyration, while keeping pTM high\n",
     "\n"
    ]
   },
@@ -49,6 +49,7 @@
    "source": [
     "import biotite.structure as bs\n",
     "import py3Dmol\n",
+    "\n",
     "from esm.sdk.api import ESMProtein, GenerationConfig\n",
     "from esm.sdk.experimental import ESM3GuidedDecoding, GuidedDecodingScoringFunction"
    ]
@@ -269,6 +270,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Start from a fully masked protein\n",
+    "PROTEIN_LENGTH = 256\n",
+    "starting_protein = ESMProtein(sequence=\"_\" * PROTEIN_LENGTH)\n",
+    "\n",
+    "# Call guided_generate\n",
     "no_cysteine_protein = no_cysteine_guided_decoding.guided_generate(\n",
     "    protein=starting_protein,\n",
     "    num_decoding_steps=len(starting_protein) // 8,\n",
@@ -302,7 +308,20 @@
    "source": [
     "## Maximize Globularity\n",
     "\n",
-    "We use the radius of gyration as a proxy to maximize globularity, we also encourage generations to have high pTM"
+    "We use the radius of gyration as a proxy to maximize globularity, and we will also encourage generations to have high pTM by using constraints"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from esm.sdk.experimental import (\n",
+    "    ConstraintType,\n",
+    "    ESM3GuidedDecodingWithConstraints,\n",
+    "    GenerationConstraint,\n",
+    ")"
    ]
   },
   {
@@ -313,12 +332,11 @@
    "source": [
     "class RadiousOfGyrationScoringFunction(GuidedDecodingScoringFunction):\n",
     "    def __call__(self, protein: ESMProtein) -> float:\n",
+    "        # Use the negative radius of gyration as the score to maximize\n",
     "        score = -1 * self.radius_of_gyration(protein)\n",
     "\n",
-    "        assert protein.ptm is not None, \"Protein must have pTM scores to be scored\"\n",
-    "        if protein.ptm < 0.5:\n",
-    "            # Penalize proteins with low pTM scores\n",
-    "            score = score * 2\n",
+    "        # Re-scale the score to be in a similar magnitude as pTM\n",
+    "        score = score / 100.0\n",
     "\n",
     "        return score\n",
     "\n",
@@ -335,8 +353,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "radius_guided_decoding = ESM3GuidedDecoding(\n",
-    "    client=model, scoring_function=RadiousOfGyrationScoringFunction()\n",
+    "# Constrain generation to have pTM > 0.75\n",
+    "ptm_constraint = GenerationConstraint(\n",
+    "    scoring_function=PTMScoringFunction(),\n",
+    "    constraint_type=ConstraintType.GREATER_EQUAL,\n",
+    "    value=0.75,\n",
+    ")\n",
+    "\n",
+    "radius_guided_decoding = ESM3GuidedDecodingWithConstraints(\n",
+    "    client=model,\n",
+    "    scoring_function=RadiousOfGyrationScoringFunction(),\n",
+    "    constraints=[ptm_constraint],  # Add list of constraints\n",
+    "    damping=1.0,  # Damping factor for the MMDM algorithm\n",
+    "    learning_rate=10.0,  # Learning rate for the MMDM algorithm\n",
     ")"
    ]
   },
@@ -346,6 +375,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Start from a fully masked protein\n",
+    "PROTEIN_LENGTH = 256\n",
+    "starting_protein = ESMProtein(sequence=\"_\" * PROTEIN_LENGTH)\n",
+    "\n",
+    "# Call guided_generate\n",
     "radius_guided_protein = radius_guided_decoding.guided_generate(\n",
     "    protein=starting_protein,\n",
     "    num_decoding_steps=len(starting_protein) // 8,\n",
@@ -359,11 +393,32 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Visualize the trajectory of the constrained generation\n",
+    "radius_guided_decoding.visualize_latest_trajectory()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Visualize the generated protein\n",
     "view = py3Dmol.view(width=800, height=400)\n",
     "view.addModel(radius_guided_protein.to_pdb_string(), \"pdb\")\n",
     "view.setStyle({\"cartoon\": {\"color\": \"spectrum\"}})\n",
     "view.zoomTo()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Check pTM\n",
+    "radius_guided_protein.ptm"
+   ]
   }
  ],
  "metadata": {

esm/__init__.py

@@ -1,2 +1 @@
 __version__ = "3.2.1"
-

esm/layers/attention.py

@@ -5,15 +5,12 @@
 import torch.nn.functional as F
 from torch import nn
 
-from esm.layers.rotary import (
-    RotaryEmbedding,
-    TritonRotaryEmbedding,
-)
+from esm.layers.rotary import RotaryEmbedding, TritonRotaryEmbedding
 
 try:
-    from flash_attn import flash_attn_varlen_qkvpacked_func  # type:ignore
-except ImportError:
-    flash_attn_varlen_func = None
+    from flash_attn import flash_attn_varlen_qkvpacked_func
+except (ImportError, RuntimeError):
+    flash_attn_varlen_qkvpacked_func = None
 
 
 class MultiHeadAttention(nn.Module):
@@ -117,7 +114,7 @@ def forward(self, x, seq_id):
         )
         qkv_N3HD = self.rotary(qkv_N3HD, cu_seqlens, max_seqlen)
 
-        context_NHD = flash_attn_varlen_qkvpacked_func(
+        context_NHD = flash_attn_varlen_qkvpacked_func(  # type: ignore
             qkv_N3HD, cu_seqlens, max_seqlen, softmax_scale=self.d_head**-0.5
         )
         context_ND = einops.rearrange(context_NHD, "n h d -> n (h d)")

esm/layers/blocks.py

@@ -2,13 +2,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from esm.layers.attention import (
-    FlashMultiHeadAttention,
-    MultiHeadAttention,
-)
-from esm.layers.geom_attention import (
-    GeometricReasoningOriginalImpl,
-)
+from esm.layers.attention import FlashMultiHeadAttention, MultiHeadAttention
+from esm.layers.geom_attention import GeometricReasoningOriginalImpl
 from esm.utils.structure.affine3d import Affine3D
 
 

esm/layers/structure_proj.py

@@ -2,10 +2,7 @@
 import torch.nn as nn
 
 from esm.utils.constants.physics import BB_COORDINATES
-from esm.utils.structure.affine3d import (
-    Affine3D,
-    RotationMatrix,
-)
+from esm.utils.structure.affine3d import Affine3D, RotationMatrix
 
 
 class Dim6RotStructureHead(nn.Module):