ported over updates for 3.1.6 (#219)

Co-authored-by: chetan <[email protected]>

Author: cmishra

cookbook/tutorials/3_gfp_design.ipynb

@@ -18,7 +18,7 @@
     "   generate many designs and filter them with a set of computational filters and\n",
     "   ranking mechanisms.\n",
     "\n",
-    "And we validated a small number of the generated designs in a wet lab, which of course you can also do... but this notebook isn't very helpful with that!"
+    "And we validated a small number of the generated designs in a wet lab, which of course you can also do... but this notebook isn't very helpful with that!\n"
    ]
   },
   {
@@ -29,12 +29,12 @@
    "source": [
     "## Set up the notebook and model (via the Forge API).\n",
     "\n",
-    "We begin by installing the [esm package](https://github.com/evolutionaryscale/esm) and py3Dmol, which will allow us to visualize our generations, and then importing necessary packages."
+    "We begin by installing the [esm package](https://github.com/evolutionaryscale/esm) and py3Dmol, which will allow us to visualize our generations, and then importing necessary packages.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {
     "id": "VgTZdaIMQ44H"
    },
@@ -50,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 2,
    "metadata": {
     "id": "poK5NTaXRGcX"
    },
@@ -75,17 +75,16 @@
     "id": "vmVYm2uQ7m-5"
    },
    "source": [
-    "\n",
-    "ESM3 is a frontier generative model for biology. It is  scalable due to its ability to tokenize sequence, structure, and function and use a (nearly standard) transformer architecture while still being able to reason across all modalities simulateously.\n",
+    "ESM3 is a frontier generative model for biology. It is scalable due to its ability to tokenize sequence, structure, and function and use a (nearly standard) transformer architecture while still being able to reason across all modalities simulateously.\n",
     "\n",
     "The largest ESM3 (98 billion parameters) was trained with 1.07e24 FLOPs on 2.78 billion proteins and 771 billion unique tokens. To create esmGFP we used the 7 billion parameter variant of ESM3. We'll use this model via the [EvolutionaryScale Forge](https://forge.evolutionaryscale.ai) API.\n",
     "\n",
-    "Grab a token from [the Forge console](https://forge.evolutionaryscale.ai/console) and add it below. Note that your token is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new token and delete old, unused ones. It is also recommended to paste the token directly into an environment variable or use a utility like `getpass` as shown below so tokens are not accidentally shared or checked into code repositories."
+    "Grab a token from [the Forge console](https://forge.evolutionaryscale.ai/console) and add it below. Note that your token is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new token and delete old, unused ones. It is also recommended to paste the token directly into an environment variable or use a utility like `getpass` as shown below so tokens are not accidentally shared or checked into code repositories.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 3,
    "metadata": {
     "id": "zNrU9Q2SYonX"
    },
@@ -100,12 +99,12 @@
     "id": "9jIc4OZyh2oE"
    },
    "source": [
-    "We then create a model stub that behaves somewhat like a PyTorch model, but under the hood it sends the inputs to the Forge server, runs the inputs through the neural network weights on that remote server, and then returns the output tensors here in this notebook. This stub can also be used in the EvolutionaryScale SDK to simplify a lot of the operations around generation, folding, and generally using the sampling. This is important because iterative sampling is key to getting the best performance from ESM3, and the SDK manages a lot of the complexity around implementing these standard routines."
+    "We then create a model stub that behaves somewhat like a PyTorch model, but under the hood it sends the inputs to the Forge server, runs the inputs through the neural network weights on that remote server, and then returns the output tensors here in this notebook. This stub can also be used in the EvolutionaryScale SDK to simplify a lot of the operations around generation, folding, and generally using the sampling. This is important because iterative sampling is key to getting the best performance from ESM3, and the SDK manages a lot of the complexity around implementing these standard routines.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 4,
    "metadata": {
     "id": "Tna_mjGOjdXA"
    },
@@ -138,7 +137,7 @@
     "\n",
     "Prompt engineering is a bit of an art and a bit of a science, so one typically needs to experiment to get a prompt that produces a desired result. Also because we use sampling to generate from the model the results of different generations from the same prompt will vary. Some prompts tend to have higher success rates requiring only a few generations to get a candidate protein design. Other more difficult prompts may require thousands of generations! The models are more controllable with alignment.\n",
     "\n",
-    "The model we will be using is the raw pretrained (unaligned) model, but we've worked a lot on this prompt so one can typically get an interesting design with only a few generations."
+    "The model we will be using is the raw pretrained (unaligned) model, but we've worked a lot on this prompt so one can typically get an interesting design with only a few generations.\n"
    ]
   },
   {
@@ -147,7 +146,7 @@
     "id": "qtwnyA1BngWy"
    },
    "source": [
-    "We'll construct our prompt from fragments of the [1qy3](https://www.rcsb.org/structure/1qy3) sequence and structure from the PDB. The following code fetches data from the PDB and then uses ESM3's tokenizers to convert the sequence and structure to tokens that can be passed into the model. Once can see that both the amino acid type and the coordinates are converted into one discrete token per sequence position."
+    "We'll construct our prompt from fragments of the [1qy3](https://www.rcsb.org/structure/1qy3) sequence and structure from the PDB. The following code fetches data from the PDB and then uses ESM3's tokenizers to convert the sequence and structure to tokens that can be passed into the model. Once can see that both the amino acid type and the coordinates are converted into one discrete token per sequence position.\n"
    ]
   },
   {
@@ -161,7 +160,6 @@
     "template_gfp = ESMProtein.from_protein_chain(\n",
     "    ProteinChain.from_rcsb(\"1qy3\", chain_id=\"A\")\n",
     ")\n",
-    "\n",
     "template_gfp_tokens = model.encode(template_gfp)\n",
     "\n",
     "print(\"Sequence tokens:\")\n",
@@ -183,7 +181,7 @@
    "source": [
     "We'll now build a prompt. Specifically we'll specify 4 amino acid identities at positions near where we want the chromophore to form, and 2 amino acid identities on the beta barrel that are known to support chromophore formation.\n",
     "\n",
-    "Furthermore we'll specify the structure should be similar to the 1qy3 structure at all these positions by adding tokens from the encoded 1qy3 structure to the structure track of our prompt. We'll also specify a few more positions (along the alpha helix kink)."
+    "Furthermore we'll specify the structure should be similar to the 1qy3 structure at all these positions by adding tokens from the encoded 1qy3 structure to the structure track of our prompt. We'll also specify a few more positions (along the alpha helix kink).\n"
    ]
   },
   {
@@ -229,7 +227,7 @@
    "source": [
     "The output shows the original 1qy3 sequence and the our prompt sequence track amino acid identities and the positions that have a token on the structure track. ESM3 will then be tasked with filling in the structure and sequence at the remaining masked (underscore) positions.\n",
     "\n",
-    "One small note, we introduced the mutation A93R in our prompt. This isn't a mistake. Using Alanine at this position causes the chromophore to mature extremely slowly (which is how we are able to measure the precyclized structure of GFP!). However we don't want to wait around for our GFPs to glow so we go with Arginine at this position."
+    "One small note, we introduced the mutation A93R in our prompt. This isn't a mistake. Using Alanine at this position causes the chromophore to mature extremely slowly (which is how we are able to measure the precyclized structure of GFP!). However we don't want to wait around for our GFPs to glow so we go with Arginine at this position.\n"
    ]
   },
   {
@@ -256,7 +254,8 @@
    "source": [
     "%%time\n",
     "\n",
-    "num_tokens_to_decode = (prompt.structure == 4096).sum().item()\n",
+    "num_tokens_to_decode = min((prompt.structure == 4096).sum().item(), 20)\n",
+    "\n",
     "\n",
     "structure_generation = model.generate(\n",
     "    prompt,\n",
@@ -276,9 +275,7 @@
     ")\n",
     "\n",
     "# Decodes structure tokens to backbone coordinates.\n",
-    "structure_generation_protein = model.decode(structure_generation)\n",
-    "\n",
-    "print(\"\")"
+    "structure_generation_protein = model.decode(structure_generation)"
    ]
   },
   {
@@ -287,7 +284,7 @@
     "id": "0HARel94tJfI"
    },
    "source": [
-    "Now let's visualize our generated structure. This will probably look like the familiar GFP beta barrel around an alpha helix."
+    "Now let's visualize our generated structure. This will probably look like the familiar GFP beta barrel around an alpha helix.\n"
    ]
   },
   {
@@ -316,7 +313,7 @@
    "source": [
     "At this point we only want to continue the generation if this design is a close match to a wildtype GFP at the active site, has some structural difference across the full protein (otherwise it would end up being very sequence-similar to wildtype GFP), and overall still looks like the classic GFP alpha helix in a beta barrel structure.\n",
     "\n",
-    "Of course when generating many designs we cannot look at each one manually, so we adopt some automated rejection sampling criteria based on the overall structure RMSD and the constrained site RMSD for our generated structure being faithful to the prompt. If these checks pass then we'll try to design a sequence for this structure. If not, one should go back up a few cells and design another structure until it passes these computational screens. (Or not... this is your GFP design!)"
+    "Of course when generating many designs we cannot look at each one manually, so we adopt some automated rejection sampling criteria based on the overall structure RMSD and the constrained site RMSD for our generated structure being faithful to the prompt. If these checks pass then we'll try to design a sequence for this structure. If not, one should go back up a few cells and design another structure until it passes these computational screens. (Or not... this is your GFP design!)\n"
    ]
   },
   {
@@ -354,7 +351,7 @@
     "\n",
     "Now we have a backbone with some structural variation but that also matches the GFP constrained site, and we want to design a sequence that folds to this structure. We can use the prior generation, which is exactly our original prompt plus the new structure tokens representing the backbone, to prompt ESM3 again.\n",
     "\n",
-    "One we have designed a sequence we'll want to confirm that sequence is a match for our structure, so we'll remove all other conditioning from the prompt and fold the sequence. Conveniently with ESM3, folding a sequence is simply generating a set of structure tokens conditioned on the amino acid sequence. In this case we want the model's highest confidence generation (with no diversity) so we sample with a temperature of zero."
+    "One we have designed a sequence we'll want to confirm that sequence is a match for our structure, so we'll remove all other conditioning from the prompt and fold the sequence. Conveniently with ESM3, folding a sequence is simply generating a set of structure tokens conditioned on the amino acid sequence. In this case we want the model's highest confidence generation (with no diversity) so we sample with a temperature of zero.\n"
    ]
   },
   {
@@ -394,7 +391,7 @@
     "id": "v_zK7TDCzEX3"
    },
    "source": [
-    "We now have a candidate GFP sequence!"
+    "We now have a candidate GFP sequence!\n"
    ]
   },
   {
@@ -414,7 +411,7 @@
     "id": "LBvQYpR_zQAK"
    },
    "source": [
-    "We can align this sequence against the original template to see how similar it is to avGFP. One might also want to search against all known fluorescent proteins to assess the novelty of this potential GFP."
+    "We can align this sequence against the original template to see how similar it is to avGFP. One might also want to search against all known fluorescent proteins to assess the novelty of this potential GFP.\n"
    ]
   },
   {
@@ -455,7 +452,7 @@
    "source": [
     "We now recheck our computational metrics for the constrained site. If we see the constrained site is not a match then we'd want to try designing the sequence again. If many attempts to design a sequence that matches the structure fail, then it's likely the structure is not easily designable and we may want to reject this structure generation as well!\n",
     "\n",
-    "At this point the backbone RMSD doesn't matter very much to us, so long as the sequence is adequately distant to satisfy our scientific curiosity!"
+    "At this point the backbone RMSD doesn't matter very much to us, so long as the sequence is adequately distant to satisfy our scientific curiosity!\n"
    ]
   },
   {
@@ -487,7 +484,7 @@
     "id": "0cIeC4Lg1Bz9"
    },
    "source": [
-    "An now we can visualize the final structure prediction of our candidate GFP design."
+    "An now we can visualize the final structure prediction of our candidate GFP design.\n"
    ]
   },
   {
@@ -511,7 +508,7 @@
     "id": "VrNuZHeHRWuP"
    },
    "source": [
-    "Before considering this sequence for wet lab validation, we run a joint optimization of the sequence and structure. The outputs of that process are then passed through stringent computational filters and then many designs from many starting points are ranked by a number of computational scores to select the final designs sent for testing. We'll walk through that process in a different notebook."
+    "Before considering this sequence for wet lab validation, we run a joint optimization of the sequence and structure. The outputs of that process are then passed through stringent computational filters and then many designs from many starting points are ranked by a number of computational scores to select the final designs sent for testing. We'll walk through that process in a different notebook.\n"
    ]
   },
   {
@@ -520,7 +517,7 @@
     "id": "c3jSQrJa1Tfi"
    },
    "source": [
-    "If you've made it this far it's worth noting that this isn't the only method to prompt ESM3 to design a GFP, it's just the one we used to report the successful generation of esmGFP in our paper. We hope you'll try different techniques to generate from ESM3. We're interested to hear what works for you!"
+    "If you've made it this far it's worth noting that this isn't the only method to prompt ESM3 to design a GFP, it's just the one we used to report the successful generation of esmGFP in our paper. We hope you'll try different techniques to generate from ESM3. We're interested to hear what works for you!\n"
    ]
   }
  ],
@@ -543,7 +540,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,

esm/sdk/__init__.py

@@ -30,7 +30,6 @@ def batch_executor(max_attempts: int = 10):
 
     Usage:
         with batch_executor() as executor:
-            for inputs in input_batches:
-                executor.submit(fn, inputs)
+            executor.execute_batch(fn, **kwargs)
     """
     return ForgeBatchExecutor(max_attempts)