Revert "sync 3.2.2.post1 (#270)" (#272)

Author: thomas-a-neil

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

@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn.functional as F
+
 from esm.pretrained import (
     ESM3_function_decoder_v0,
     ESM3_sm_open_v0,
@@ -12,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/snippets/fold_invfold.py

@@ -2,6 +2,7 @@
 from typing import cast
 
 import numpy as np
+
 from esm.sdk.api import (
     ESM3InferenceClient,
     ESMProtein,

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/2_embed.ipynb

@@ -49,18 +49,18 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Grab a token from [Forge](https://forge.evolutionaryscale.ai/) 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."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
     "from getpass import getpass\n",
     "\n",
-    "token = getpass(\"Token from Forge: \")"
+    "token = getpass(\"Token from Forge console: \")"
    ]
   },
   {

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"
@@ -79,18 +80,18 @@
     "\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 [Forge](https://forge.evolutionaryscale.ai/) 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"
+    "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": null,
+   "execution_count": 3,
    "metadata": {
     "id": "zNrU9Q2SYonX"
    },
    "outputs": [],
    "source": [
-    "token = getpass(\"Token from Forge: \")"
+    "token = getpass(\"Token from Forge console: \")"
    ]
   },
   {

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"
@@ -52,7 +53,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Grab a token from [Forge](https://forge.evolutionaryscale.ai/) 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."
    ]
   },
   {
@@ -63,7 +64,7 @@
    "source": [
     "from getpass import getpass\n",
     "\n",
-    "token = getpass(\"Token from Forge: \")\n",
+    "token = getpass(\"Token from Forge console: \")\n",
     "model = client(model=\"esm3-open\", url=\"https://forge.evolutionaryscale.ai\", token=token)"
    ]
   },

cookbook/tutorials/5_guided_generation.ipynb

@@ -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"
    ]
@@ -119,7 +120,7 @@
     "\n",
     "from esm.sdk import client\n",
     "\n",
-    "token = getpass(\"Token from Forge: \")\n",
+    "token = getpass(\"Token from Forge console: \")\n",
     "model = client(\n",
     "    model=\"esm3-medium-2024-08\", url=\"https://forge.evolutionaryscale.ai\", token=token\n",
     ")"

esm/__init__.py

@@ -1,2 +1 @@
-__version__ = "3.2.2.post1"
-
+__version__ = "3.2.2"

esm/layers/attention.py

@@ -5,10 +5,7 @@
 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