oss sync (#347)

Author: imathur1

.pre-commit-config.yaml

@@ -22,17 +22,17 @@ repos:
     args: [ --fix ]
   - id: ruff-format # formatter
     types_or: [python, jupyter]
-- repo: https://github.com/RobertCraigie/pyright-python
-  rev: v1.1.399
+- repo: local
   hooks:
-    - id: pyright
-      name: pyright
-      entry: pyright
-      language: system
-      types: [python]
-      pass_filenames: true # For speed, we only check the files that are changed
-      # Modal-app tutorial: deps (modal, abnumber) and dynamic decorators aren't resolvable in the lint env.
-      exclude: ^cookbook/tutorials/binder_design\.py$
+    - id: ty
+      name: ty
+      entry: ty check
+      language: system # ty is a pixi dev dep (pyproject.toml [tool.pixi.feature.dev]); pre-commit runs in the pixi env
+      # pass_filenames: false — ty checks the whole project (it needs the full module graph and
+      # is fast enough), and per-file invocation would bypass [tool.ty.src].exclude (astral-sh/ty#269).
+      pass_filenames: false
+      always_run: true
+      require_serial: true
 - repo: https://github.com/gitleaks/gitleaks
   rev: v8.24.2
   hooks:

README.md

@@ -4,7 +4,7 @@
 # A world model of protein biology: ESMC, ESMFold2, & ESM Atlas
 
 
-[ESMC & ESMFold2 Preprint](https://biohub.ai/papers/esm_protein.pdf) &sdot;  [Atlas](https://biohub.ai/esm/protein/atlas) &sdot; [Tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials) &sdot; [Slack](https://bit.ly/esm-slack)<br>
+[ESMC & ESMFold2 Preprint](https://www.biorxiv.org/content/10.64898/2026.06.03.729735) &sdot;  [Atlas](https://biohub.ai/esm/protein/atlas) &sdot; [Tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials) &sdot; [Slack](https://bit.ly/esm-slack)<br>
 </div>
 
 We are releasing a world model for protein biology: a scientific engine for prediction, design, and discovery. Built on the latest generation of Evolutionary Scale Modeling (ESM), this system learns from the protein sequences produced by evolution and uses that knowledge to represent, map, predict, and design proteins across scales — from atomic interactions to evolutionary relationships spanning billions of years. The system includes three artifacts: ESMC, ESMFold2, and ESM Atlas.
@@ -25,7 +25,7 @@ We are releasing a world model for protein biology: a scientific engine for pred
 </div>
 
 
-ESMFold2 is validated in the lab across five therapeutic targets. Inversion of ESMFold2 enables generation of de novo minibinders and antibody-derived scFvs with high hit rates, nanomolar affinities, target specificity, and functional activity. We've released the full protocol from target sequence to ranked binder design in this [notebook](https://github.com/Biohub/esm/blob/main/cookbook/tutorials/binder_design.ipynb). For additional details, please refer to the [preprint](https://biohub.ai/papers/esm_protein.pdf).
+ESMFold2 is validated in the lab across five therapeutic targets. Inversion of ESMFold2 enables generation of de novo minibinders and antibody-derived scFvs with high hit rates, nanomolar affinities, target specificity, and functional activity. We've released the full protocol from target sequence to ranked binder design in this [notebook](https://github.com/Biohub/esm/blob/main/cookbook/tutorials/binder_design.ipynb). For additional details, please refer to the [preprint](https://www.biorxiv.org/content/10.64898/2026.06.03.729735).
 
 <div align="center">
   <img src="_assets/esmfold2_binder.png" width="60%"/>
@@ -320,7 +320,7 @@ If you use ESM in your work, please cite one of the following:
             and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S.
             and Sercu, Tom and Rives, Alexander},
   year   = {2026},
-  url    = {https://biohub.ai/papers/esm_protein.pdf},
+  url    = {https://www.biorxiv.org/content/10.64898/2026.06.03.729735},
   note   = {Preprint}
 }
 ```

cookbook/tutorials/binder_design.ipynb

@@ -7,7 +7,7 @@
    "source": [
     "## [Tutorial](https://github.com/biohub/esm/tree/main/cookbook/tutorials): How to run minibinder + scFv design fully end-to-end.\n",
     "\n",
-    "In this notebook we will use [Modal](https://modal.com/) to parallelize binder design and synthesize a selection, using the protocol described in the ESMC and ESMFold2 paper titled [\"Language Modeling Materializes a World Model of Protein Biology\"](https://biohub.ai/papers/esm_protein.pdf).\n",
+    "In this notebook we will use [Modal](https://modal.com/) to parallelize binder design and synthesize a selection, using the protocol described in the ESMC and ESMFold2 paper titled [\"Language Modeling Materializes a World Model of Protein Biology\"](https://www.biorxiv.org/content/10.64898/2026.06.03.729735).\n",
     "\n",
     "Biohub used this approach to design minibinders and scFvs against five therapeutically relevant targets — PDGFRB, EGFR, PD-L1, CD45, and CTLA4 — spanning receptor tyrosine kinases, immune checkpoints, and cell-surface phosphatases. Binders exhibit nanomolar affinity, target specificity, and functional activity in laboratory assays.\n",
     "\n",

cookbook/tutorials/binder_design.py

@@ -9,7 +9,7 @@
 """
 Code for binder design with ESMFold2 and ESMC.
 
-As described in [Language Modeling Materializes a World Model of Protein Biology](https://biohub.ai/papers/esm_protein.pdf).
+As described in [Language Modeling Materializes a World Model of Protein Biology](https://www.biorxiv.org/content/10.64898/2026.06.03.729735).
 """
 
 import logging
@@ -1051,7 +1051,7 @@ def _apply_torch_compile(model: torch.nn.Module) -> None:
     def _maybe_compile_module(module: torch.nn.Module) -> None:
         if not isinstance(module, compile_targets):
             return
-        module.forward = torch.compile(module.forward)  # pyright: ignore
+        module.forward = torch.compile(module.forward)  # ty:ignore[invalid-assignment]
 
     model.apply(_maybe_compile_module)
 
@@ -1206,7 +1206,9 @@ def main(
         app.load(use_scaling_critics)
         run_fn = app.design
     else:
-        app = ESMFold2DesignModal(use_scaling_critics=use_scaling_critics)
+        app = ESMFold2DesignModal(
+            use_scaling_critics=use_scaling_critics  # ty:ignore[unknown-argument]
+        )
         run_fn = app.design.remote
 
     seq, trajectory, results = run_fn(

cookbook/tutorials/embed.ipynb

@@ -25,7 +25,8 @@
     "# If you are working in colab, uncomment these lines to install dependencies\n",
     "#! pip install esm@git+https://github.com/Biohub/esm.git@main\n",
     "#! pip install matplotlib\n",
-    "#! pip install seaborn"
+    "#! pip install seaborn\n",
+    "#! pip install remotezip"
    ]
   },
   {
@@ -137,47 +138,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--2026-05-28 16:44:15--  https://docs.google.com/uc?export=download&id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg\n",
-      "Resolving docs.google.com (docs.google.com)... 142.251.210.78\n",
-      "Connecting to docs.google.com (docs.google.com)|142.251.210.78|:443... connected.\n",
-      "HTTP request sent, awaiting response... 303 See Other\n",
-      "Location: https://drive.usercontent.google.com/download?id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg&export=download [following]\n",
-      "--2026-05-28 16:44:15--  https://drive.usercontent.google.com/download?id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg&export=download\n",
-      "Resolving drive.usercontent.google.com (drive.usercontent.google.com)... 192.178.50.65\n",
-      "Connecting to drive.usercontent.google.com (drive.usercontent.google.com)|192.178.50.65|:443... connected.\n",
-      "HTTP request sent, awaiting response... 200 OK\n",
-      "Length: 43132 (42K) [application/octet-stream]\n",
-      "Saving to: 'adk.csv'\n",
-      "\n",
-      "adk.csv             100%[===================>]  42.12K  --.-KB/s    in 0.02s   \n",
-      "\n",
-      "2026-05-28 16:44:16 (2.17 MB/s) - 'adk.csv' saved [43132/43132]\n",
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "!wget --no-check-certificate \"https://docs.google.com/uc?export=download&id=1SpOkL11MJxIgy99dqufvUNJuCiuhxuyg\" -O adk.csv"
+    "# Stream just the one CSV we need out of the 200MB archive, no full download.\n",
+    "from remotezip import RemoteZip\n",
+    "\n",
+    "DATA_URL = \"https://zenodo.org/records/15022271/files/data.zip\"\n",
+    "MEMBER = \"data/adk_ml_dataset.csv\"\n",
+    "\n",
+    "with RemoteZip(DATA_URL) as zf:\n",
+    "    with zf.open(MEMBER) as src, open(\"adk_ml_dataset.csv\", \"wb\") as dst:\n",
+    "        dst.write(src.read())"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import pandas as pd\n",
     "import seaborn as sns\n",
     "\n",
-    "adk_path = \"adk.csv\"\n",
+    "adk_path = \"adk_ml_dataset.csv\"\n",
     "df = pd.read_csv(adk_path)\n",
     "df = df[[\"org_name\", \"sequence\", \"lid_type\", \"temperature\"]]\n",
     "df = df[df[\"lid_type\"] != \"other\"]  # drop one structural class for simplicity"
@@ -231,7 +217,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -244,7 +230,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -322,9 +308,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Pixi (ESM)",
+   "display_name": "default",
    "language": "python",
-   "name": "pixi-esm"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -336,7 +322,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.12"
+   "version": "3.12.9"
   }
  },
  "nbformat": 4,

cookbook/tutorials/esmc_sae_feature_interpretation.ipynb

@@ -493,7 +493,7 @@
     "        for pos in top_positions:\n",
     "            if activations[pos] > 0:\n",
     "                ax.annotate(\n",
-    "                    f\"{sequence[pos]}{pos+1}\",\n",
+    "                    f\"{sequence[pos]}{pos + 1}\",\n",
     "                    (pos + 1, activations[pos]),\n",
     "                    textcoords=\"offset points\",\n",
     "                    xytext=(0, 6),\n",
@@ -508,7 +508,7 @@
     "    top_positions = np.argsort(activations)[::-1][:5]\n",
     "    for pos in top_positions:\n",
     "        if activations[pos] > 0:\n",
-    "            print(f\"  Position {pos+1} ({sequence[pos]}): {activations[pos]:.3f}\")\n",
+    "            print(f\"  Position {pos + 1} ({sequence[pos]}): {activations[pos]:.3f}\")\n",
     "    print()"
    ]
   },

cookbook/tutorials/gfp_design.ipynb

@@ -432,7 +432,7 @@
     "alignment = alignments[0]\n",
     "\n",
     "identity = align.get_sequence_identity(alignment)\n",
-    "print(f\"Sequence identity: {100*identity:.2f}%\")\n",
+    "print(f\"Sequence identity: {100 * identity:.2f}%\")\n",
     "\n",
     "print(\"\\nSequence alignment:\")\n",
     "fig = pl.figure(figsize=(8.0, 4.0))\n",

esm/layers/attention.py

@@ -8,9 +8,9 @@
 from esm.layers.rotary import RotaryEmbedding, TritonRotaryEmbedding
 
 try:
-    from flash_attn import flash_attn_varlen_qkvpacked_func  # type: ignore
+    from flash_attn import flash_attn_varlen_qkvpacked_func
 except (ImportError, RuntimeError):
-    flash_attn_varlen_qkvpacked_func = None
+    flash_attn_varlen_qkvpacked_func = None  # ty:ignore[invalid-assignment]
 
 
 class MultiHeadAttention(nn.Module):
@@ -135,7 +135,8 @@ def forward(
         )
         qkv_N3HD = self.rotary(qkv_N3HD, cu_seqlens, max_seqlen)
 
-        context_NHD = flash_attn_varlen_qkvpacked_func(  # type: ignore
+        assert flash_attn_varlen_qkvpacked_func is not None
+        context_NHD = flash_attn_varlen_qkvpacked_func(
             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/codebook.py

@@ -6,6 +6,10 @@
 
 
 class EMACodebook(nn.Module):
+    embeddings: torch.Tensor
+    N: torch.Tensor
+    z_avg: torch.Tensor
+
     def __init__(
         self,
         n_codes,
@@ -17,7 +21,7 @@ def __init__(
         super().__init__()
         self.register_buffer("embeddings", torch.randn(n_codes, embedding_dim))
         self.register_buffer("N", torch.zeros(n_codes))
-        self.register_buffer("z_avg", self.embeddings.data.clone())  # pyright: ignore[reportCallIssue]
+        self.register_buffer("z_avg", self.embeddings.data.clone())
 
         self.n_codes = n_codes
         self.embedding_dim = embedding_dim
@@ -50,9 +54,9 @@ def _init_embeddings(self, z):
         _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
         if dist.is_initialized():
             dist.broadcast(_k_rand, 0)
-        self.embeddings.data.copy_(_k_rand)  # pyright: ignore[reportCallIssue]
-        self.z_avg.data.copy_(_k_rand)  # pyright: ignore[reportCallIssue]
-        self.N.data.copy_(torch.ones(self.n_codes))  # pyright: ignore[reportCallIssue]
+        self.embeddings.data.copy_(_k_rand)
+        self.z_avg.data.copy_(_k_rand)
+        self.N.data.copy_(torch.ones(self.n_codes))
 
     def forward(self, z):
         # z: [b, t, c]
@@ -62,17 +66,14 @@ def forward(self, z):
         flat_inputs = z.view(-1, self.embedding_dim)
         distances = (
             (flat_inputs**2).sum(dim=1, keepdim=True)
-            - 2 * flat_inputs @ self.embeddings.t()  # pyright: ignore[reportCallIssue]
-            + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)  # pyright: ignore[reportCallIssue]
+            - 2 * flat_inputs @ self.embeddings.t()
+            + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)
         )  # [bt, c]
 
         encoding_indices = torch.argmin(distances, dim=1)
         encoding_indices = encoding_indices.view(*z.shape[:2])  # [b, t, ncode]
 
-        embeddings = F.embedding(
-            encoding_indices,
-            self.embeddings,  # pyright: ignore[reportArgumentType]
-        )  # [b, t, c] # pyright: ignore[reportArgumentType]
+        embeddings = F.embedding(encoding_indices, self.embeddings)  # [b, t, c]
 
         commitment_loss = 0.25 * F.mse_loss(z, embeddings.detach())
 
@@ -84,8 +85,8 @@ def forward(self, z):
         return embeddings_st, encoding_indices, commitment_loss
 
     def dictionary_lookup(self, encodings):
-        embeddings = F.embedding(encodings, self.embeddings)  # pyright: ignore[reportArgumentType]
+        embeddings = F.embedding(encodings, self.embeddings)
         return embeddings
 
     def soft_codebook_lookup(self, weights: torch.Tensor) -> torch.Tensor:
-        return weights @ self.embeddings  # pyright: ignore[reportOperatorIssue]
+        return weights @ self.embeddings

esm/layers/rotary.py

@@ -26,11 +26,9 @@
 from einops import rearrange, repeat
 
 try:
-    from flash_attn.ops.triton.rotary import (  # type:ignore
-        apply_rotary as apply_triton_rotary,
-    )
+    from flash_attn.ops.triton.rotary import apply_rotary as apply_triton_rotary
 except ImportError:
-    apply_triton_rotary = None
+    apply_triton_rotary = None  # ty:ignore[invalid-assignment]
 
 
 def rotate_half(x, interleaved=False):
@@ -169,23 +167,23 @@ def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
                 else:
                     inv_freq = self.inv_freq
             else:
-                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)  # pyright: ignore[reportArgumentType, reportCallIssue]
+                t = torch.arange(
+                    seqlen, device=device, dtype=self.inv_freq.dtype
+                )  # ty:ignore[no-matching-overload]
                 t /= self.scaling_factor
                 inv_freq = self.inv_freq
             # Don't do einsum, it converts fp32 to fp16 under AMP
             # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
-            freqs = torch.outer(t, inv_freq)  # pyright: ignore[reportArgumentType]
+            freqs = torch.outer(t, inv_freq)  # ty:ignore[invalid-argument-type]
 
             if self.scale is None:
                 self._cos_cached = torch.cos(freqs).to(dtype)
                 self._sin_cached = torch.sin(freqs).to(dtype)
             else:
                 power = (
-                    torch.arange(  # pyright: ignore[reportCallIssue]
-                        seqlen,
-                        dtype=self.scale.dtype,  # pyright: ignore[reportArgumentType]
-                        device=self.scale.device,  # pyright: ignore[reportArgumentType]
-                    )
+                    torch.arange(
+                        seqlen, dtype=self.scale.dtype, device=self.scale.device
+                    )  # ty:ignore[no-matching-overload]
                     - seqlen // 2
                 ) / self.scale_base
                 scale = self.scale.to(device=power.device) ** power.unsqueeze(-1)
@@ -225,7 +223,7 @@ def forward(
                     self.interleaved,
                     True,  # inplace=True
                 ),
-            )  # type: ignore
+            )
         else:
             assert False