Copy disabled (too large)
Showing preview only (18,319K chars total). The displayed content is truncated. Use the JSON API for full output.
Repository: evolutionaryscale/esm
Branch: main
Commit: 827ec128e4cd
Files: 202
Total size: 17.4 MB
Directory structure:
gitextract_5y54ehgk/
├── .github/
│ ├── scripts/
│ │ └── airtable_issue_sync.py
│ └── workflows/
│ ├── airtable-issue-sync.yaml
│ └── ci.yml
├── .gitignore
├── .pre-commit-config.yaml
├── CONTRIBUTIONS.md
├── LICENSE.md
├── README.md
├── THIRD_PARTY_NOTICE.md
├── _assets/
│ └── ESM3_README.md
├── cookbook/
│ ├── local/
│ │ ├── README.md
│ │ ├── open_generate.ipynb
│ │ └── raw_forwards.py
│ ├── snippets/
│ │ ├── README.md
│ │ ├── esm3.py
│ │ ├── esmc.py
│ │ ├── fold_invfold.py
│ │ ├── sae.py
│ │ ├── sae_example.py
│ │ └── sparse_utils.py
│ └── tutorials/
│ ├── README.md
│ ├── binder_design.ipynb
│ ├── binder_design.py
│ ├── embed.ipynb
│ ├── esm3_generate.ipynb
│ ├── esm3_guided_generation.ipynb
│ ├── esmc_finetune.ipynb
│ ├── esmc_layer_sweep.ipynb
│ ├── esmc_mutation_scoring.ipynb
│ ├── esmc_sae_feature_interpretation.ipynb
│ ├── esmfold2.ipynb
│ ├── esmfold2_local_applesilicon.ipynb
│ ├── esmfold2_local_gpu.ipynb
│ ├── esmprotein.ipynb
│ ├── g3l5_chainA.a3m
│ ├── g3l5_chainB.a3m
│ └── gfp_design.ipynb
├── esm/
│ ├── __init__.py
│ ├── data/
│ │ ├── ParentChildTreeFile.txt
│ │ ├── entry_list_safety_29026.list
│ │ ├── interpro_29026_to_keywords_58641.csv
│ │ ├── keyword_idf_safety_filtered_58641.npy
│ │ └── keyword_vocabulary_safety_filtered_58641.txt
│ ├── layers/
│ │ ├── attention.py
│ │ ├── blocks.py
│ │ ├── codebook.py
│ │ ├── ffn.py
│ │ ├── geom_attention.py
│ │ ├── regression_head.py
│ │ ├── rotary.py
│ │ ├── structure_proj.py
│ │ └── transformer_stack.py
│ ├── models/
│ │ ├── esm3.py
│ │ ├── esmc/
│ │ │ ├── __init__.py
│ │ │ ├── checkpoint_layout.py
│ │ │ ├── compatibility.py
│ │ │ ├── config.py
│ │ │ ├── kernels.py
│ │ │ ├── layers.py
│ │ │ ├── model.py
│ │ │ ├── sae.py
│ │ │ └── tokenizer.py
│ │ ├── esmfold2/
│ │ │ ├── __init__.py
│ │ │ ├── config.py
│ │ │ ├── conformers.py
│ │ │ ├── constants.py
│ │ │ ├── experimental.py
│ │ │ ├── hf_adapter.py
│ │ │ ├── hf_checkpoint.py
│ │ │ ├── kernels/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── fused_attention_pair_bias.py
│ │ │ │ ├── fused_dropout_residual.py
│ │ │ │ ├── fused_dual_gemm.py
│ │ │ │ ├── fused_ln_residual.py
│ │ │ │ ├── fused_lnlin_swiglu.py
│ │ │ │ ├── trimul_einsum_triton.py
│ │ │ │ └── trimul_with_residual.py
│ │ │ ├── layers.py
│ │ │ ├── model.py
│ │ │ ├── output.py
│ │ │ ├── paired_msa.py
│ │ │ ├── prepare_input.py
│ │ │ ├── processor.py
│ │ │ ├── protein_utils.py
│ │ │ └── types.py
│ │ ├── function_decoder.py
│ │ ├── hub.py
│ │ └── vqvae.py
│ ├── pretrained.py
│ ├── sdk/
│ │ ├── __init__.py
│ │ ├── api.py
│ │ ├── base_forge_client.py
│ │ ├── experimental/
│ │ │ ├── __init__.py
│ │ │ ├── constrained_generation.py
│ │ │ └── guided_generation.py
│ │ ├── forge.py
│ │ ├── retry.py
│ │ ├── sagemaker.py
│ │ └── validation.py
│ ├── tokenization/
│ │ ├── __init__.py
│ │ ├── function_tokenizer.py
│ │ ├── residue_tokenizer.py
│ │ ├── sasa_tokenizer.py
│ │ ├── sequence_tokenizer.py
│ │ ├── ss_tokenizer.py
│ │ ├── structure_tokenizer.py
│ │ └── tokenizer_base.py
│ ├── utils/
│ │ ├── constants/
│ │ │ ├── api.py
│ │ │ ├── esm3.py
│ │ │ ├── models.py
│ │ │ └── physics.py
│ │ ├── decoding.py
│ │ ├── encoding.py
│ │ ├── forge_context_manager.py
│ │ ├── function/
│ │ │ ├── encode_decode.py
│ │ │ ├── interpro.py
│ │ │ ├── lsh.py
│ │ │ └── tfidf.py
│ │ ├── generation.py
│ │ ├── misc.py
│ │ ├── msa/
│ │ │ ├── __init__.py
│ │ │ ├── filter_sequences.py
│ │ │ └── msa.py
│ │ ├── noise_schedules.py
│ │ ├── parsing.py
│ │ ├── residue_constants.py
│ │ ├── sampling.py
│ │ ├── sequential_dataclass.py
│ │ ├── structure/
│ │ │ ├── affine3d.py
│ │ │ ├── aligner.py
│ │ │ ├── atom_indexer.py
│ │ │ ├── input_builder.py
│ │ │ ├── metrics.py
│ │ │ ├── mmcif_parsing.py
│ │ │ ├── molecular_complex.py
│ │ │ ├── normalize_coordinates.py
│ │ │ ├── predicted_aligned_error.py
│ │ │ ├── protein_chain.py
│ │ │ ├── protein_complex.py
│ │ │ └── protein_structure.py
│ │ ├── system.py
│ │ └── types.py
│ └── widgets/
│ ├── components/
│ │ ├── function_annotator.py
│ │ ├── results_visualizer.py
│ │ ├── sasa_prompt_selector.py
│ │ ├── secondary_structure_prompt_selector.py
│ │ ├── sequence_prompt_selector.py
│ │ └── structure_prompt_selector.py
│ ├── utils/
│ │ ├── clients.py
│ │ ├── drawing/
│ │ │ ├── colors.py
│ │ │ ├── draw_category_array.py
│ │ │ ├── draw_function_annotations.py
│ │ │ └── draw_protein_structure.py
│ │ ├── indexing.py
│ │ ├── parsing.py
│ │ ├── printing.py
│ │ ├── prompting.py
│ │ ├── protein_import.py
│ │ ├── serialization.py
│ │ └── types.py
│ └── views/
│ ├── esm3_generation_launcher.py
│ ├── esm3_prompt_preview.py
│ ├── esm3_prompt_selector.py
│ ├── generation.py
│ ├── inverse_folding.py
│ ├── login.py
│ └── prediction.py
├── pyproject.toml
├── tests/
│ ├── Makefile
│ ├── __init__.py
│ ├── compatibility/
│ │ ├── __init__.py
│ │ ├── compatibility_test.py
│ │ ├── esmc_legacy_contract_test.py
│ │ ├── esmfold2_hf_adapter_test.py
│ │ └── esmfold2_hf_checkpoint_test.py
│ ├── conftest.py
│ ├── models/
│ │ ├── __init__.py
│ │ ├── esmc_test.py
│ │ ├── esmfold2_api_test.py
│ │ ├── esmfold2_builds_test.py
│ │ ├── esmfold2_cpu_only_test.py
│ │ ├── esmfold2_execution_test.py
│ │ ├── esmfold2_inputs_test.py
│ │ ├── esmfold2_msa_test.py
│ │ ├── esmfold2_sampler_test.py
│ │ ├── esmfold2_test.py
│ │ └── prepare_input_test.py
│ ├── oss_pytests/
│ │ ├── Dockerfile
│ │ ├── requirements.txt
│ │ ├── test_oss_client.py
│ │ ├── test_output_attentions.py
│ │ └── test_placeholder.py
│ ├── regenerate_reference.py
│ ├── sdk/
│ │ ├── __init__.py
│ │ └── forge_context_manager_test.py
│ └── utils/
│ ├── __init__.py
│ ├── input_builder_test.py
│ ├── misc_test.py
│ ├── molecular_complex_test.py
│ ├── msa_test.py
│ └── sampling_test.py
└── tools/
└── README.md
================================================
FILE CONTENTS
================================================
================================================
FILE: .github/scripts/airtable_issue_sync.py
================================================
#!/usr/bin/env python3
"""Sync a GitHub issue into an Airtable base.
Env: AIRTABLE_TOKEN, AIRTABLE_BASE, AIRTABLE_TABLE, and (from GitHub Actions)
GITHUB_EVENT_PATH. Optional: LABEL_TYPE_MAP, PRODUCT_KEYWORD_MAP,
DEFAULT_PRODUCTS, DRY_RUN.
Everything written comes out of the webhook payload, apart from the "Automated"
tag on "Issue Source" that marks the row as machine-written; the script makes no
calls to GitHub.
An issue is matched to an existing row by URL and updated in place, otherwise a
row is created. Fields that already hold a value are left as they are.
The Airtable token needs schema.bases:read, data.records:read and
data.records:write.
"""
import json
import os
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
API = "https://api.airtable.com/v0"
F_TYPE = "Type of Request"
F_DATE = "Submission Date"
F_PRODUCT = "Tool or Product"
F_SOURCE = "Issue Source"
F_DESCRIPTION = "Issue Description"
F_GITHUB_ID = "Github Username"
F_ORIGINAL_Q = "Original Q Location"
F_COMPLETE = "Issue Complete?"
F_DETAILS = "Additional Details"
# The sync key, "<owner>/<repo>#<number>". Held alongside the issue URL so a row
# somebody entered by hand is still matched on the URL alone.
F_REFERENCE = "Issue tracker link"
SOURCE_GITHUB = "Github"
# Stamped on every row this script touches, so rows that arrived by automation
# can be told apart from the ones people enter by hand.
SOURCE_AUTOMATED = "Automated"
STATE_OPEN = "Incomplete"
STATE_CLOSED = "Complete"
COMPLETE_STATES = {"complete", "out of scope"}
# Airtable long text holds 100k characters; stay well under it.
MAX_TEXT_CHARS = 50_000
REQUEST_TIMEOUT_SECONDS = 30
# GitHub label (lowercased) -> "Type of Request" option. Extend with
# LABEL_TYPE_MAP rather than editing this.
DEFAULT_LABEL_TYPE_MAP = {
"bug": "Bug",
"defect": "Bug",
"error": "Error",
"crash": "Error",
"enhancement": "Feature Request",
"feature": "Feature Request",
"feature request": "Feature Request",
"question": "Support",
"support": "Support",
"help wanted": "Support",
"documentation": "Support",
"docs": "Support",
"feedback": "General Feedback",
"access": "Access issue",
"permissions": "Access issue",
}
# Substring (lowercased) -> "Tool or Product" option, scanned over title, body
# and labels.
DEFAULT_PRODUCT_KEYWORD_MAP = {
"esmfold2": "ESMFold2",
"esmfold 2": "ESMFold2",
"esm atlas": "ESM Atlas",
"metagenomic atlas": "ESM Atlas",
"biohub platform": "Biohub Platform",
"esmc": "ESMC",
"esm3": "ESM3",
"esm 3": "ESM3",
"binder": "Binder",
"sae": "SAE",
}
def require_env(name):
v = os.environ.get(name)
if not v:
sys.exit(f"Missing required env var: {name}")
return v
def _lower_keys(raw):
return {k.lower(): v for k, v in json.loads(raw).items()}
AIRTABLE_TOKEN = require_env("AIRTABLE_TOKEN")
AIRTABLE_BASE = require_env("AIRTABLE_BASE")
AIRTABLE_TABLE = require_env("AIRTABLE_TABLE")
LABEL_TYPE_MAP = {
**DEFAULT_LABEL_TYPE_MAP,
**_lower_keys(os.environ.get("LABEL_TYPE_MAP", "{}")),
}
PRODUCT_KEYWORD_MAP = {
**DEFAULT_PRODUCT_KEYWORD_MAP,
**_lower_keys(os.environ.get("PRODUCT_KEYWORD_MAP", "{}")),
}
DEFAULT_PRODUCTS = json.loads(os.environ.get("DEFAULT_PRODUCTS", "[]"))
DRY_RUN = os.environ.get("DRY_RUN", "").lower() in ("1", "true", "yes")
_warned = set()
def warn(msg):
if msg not in _warned:
print(f"WARN: {msg}")
_warned.add(msg)
def _request(req, attempts=4):
for attempt in range(attempts):
try:
with urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT_SECONDS) as resp:
return json.loads(resp.read().decode("utf-8"))
except urllib.error.HTTPError as err:
if err.code in (429, 500, 502, 503) and attempt < attempts - 1:
time.sleep(2**attempt)
continue
detail = err.read().decode("utf-8", "replace")
raise RuntimeError(
f"{req.get_method()} {req.full_url} -> {err.code}: {detail}"
)
raise RuntimeError(f"{req.get_method()} {req.full_url}: retries exhausted")
def airtable(path, method="GET", body=None):
data = json.dumps(body).encode("utf-8") if body is not None else None
req = urllib.request.Request(API + path, data=data, method=method)
req.add_header("Authorization", f"Bearer {AIRTABLE_TOKEN}")
req.add_header("Content-Type", "application/json")
return _request(req)
def table_path():
return f"/{AIRTABLE_BASE}/{urllib.parse.quote(AIRTABLE_TABLE, safe='')}"
def get_schema():
"""Field name -> field schema for the target table."""
tables = airtable(f"/meta/bases/{AIRTABLE_BASE}/tables")["tables"]
for t in tables:
if AIRTABLE_TABLE in (t["id"], t["name"]):
return {f["name"]: f for f in t["fields"]}
names = ", ".join(f'"{t["name"]}"' for t in tables)
sys.exit(
f'Table "{AIRTABLE_TABLE}" not found in base {AIRTABLE_BASE}. Have: {names}'
)
def truncate(text, url):
if len(text) <= MAX_TEXT_CHARS:
return text
return text[:MAX_TEXT_CHARS] + f"\n\n... (truncated; read the full issue at {url})"
def match_choices(fschema, values):
"""Keep the values that already exist as options, in the schema's casing."""
options = [c["name"] for c in fschema.get("options", {}).get("choices", [])]
by_lower = {o.lower(): o for o in options}
matched = []
for v in values:
canonical = by_lower.get(str(v).strip().lower())
if canonical is None:
warn(f'"{fschema["name"]}" has no option matching "{v}"; dropping it.')
elif canonical not in matched:
matched.append(canonical)
return matched
def coerce(fschema, value):
"""Value in the shape Airtable wants, or None if it should not be written."""
t = fschema["type"]
if t in ("multipleSelects", "singleSelect"):
values = value if isinstance(value, list) else [value]
matched = match_choices(fschema, [v for v in values if v])
if not matched:
return None
return matched if t == "multipleSelects" else matched[0]
if t in ("multilineText", "singleLineText", "richText"):
return str(value) if value else None
if t == "date":
return str(value)[:10] if value else None
if t == "checkbox":
return bool(value)
warn(f'field "{fschema["name"]}" (type {t}) is not written by this sync; skipping.')
return None
def is_empty(current):
return current is None or current == "" or current == [] or current is False
class Row:
"""Fields to write, validated against the live schema."""
def __init__(self, schema, existing):
self.schema = schema
self.existing = (existing or {}).get("fields", {})
self.fields = {}
def _prepare(self, name, value):
fschema = self.schema.get(name)
if not fschema:
warn(f'field "{name}" not found in the table; skipping.')
return None
return coerce(fschema, value)
def own(self, name, value):
"""Keep this field equal to the GitHub value."""
prepared = self._prepare(name, value)
if prepared is not None and self.existing.get(name) != prepared:
self.fields[name] = prepared
def fill(self, name, value):
"""Write this field only while it is still empty."""
if not is_empty(self.existing.get(name)):
return
prepared = self._prepare(name, value)
if prepared is not None:
self.fields[name] = prepared
def add(self, name, values):
"""Make sure these values are present without dropping the others.
``own`` replaces the cell, which on a multi-select would discard a
source somebody selected by hand. This sync only needs its own tags to
be there.
"""
current = self.existing.get(name) or []
if not isinstance(current, list):
current = [current]
prepared = self._prepare(name, list(current) + list(values))
if prepared is None:
return
wanted = prepared if isinstance(prepared, list) else [prepared]
if sorted(map(str, wanted)) != sorted(map(str, current)):
self.fields[name] = prepared
def label_names(issue):
return [
str(lb["name"] if isinstance(lb, dict) else lb)
for lb in issue.get("labels") or []
]
def request_types(issue):
types = []
for label in label_names(issue):
mapped = LABEL_TYPE_MAP.get(label.strip().lower())
if mapped and mapped not in types:
types.append(mapped)
return types
def products(issue):
haystack = "\n".join(
[
issue.get("title") or "",
issue.get("body") or "",
" ".join(label_names(issue)),
]
).lower()
found = []
# Longest keyword first so "esmfold2" is not shadowed by a shorter match.
for keyword in sorted(PRODUCT_KEYWORD_MAP, key=len, reverse=True):
product = PRODUCT_KEYWORD_MAP[keyword]
if product in found:
continue
if re.search(rf"(?<![a-z0-9]){re.escape(keyword)}(?![a-z0-9])", haystack):
found.append(product)
return found or list(DEFAULT_PRODUCTS)
def details_block(issue, repo_full_name):
labels = label_names(issue)
return "\n".join(
[
f"Repository: {repo_full_name}",
f"Issue: #{issue['number']}",
f"Title: {issue.get('title') or ''}",
f"Author: @{(issue.get('user') or {}).get('login', 'unknown')}",
f"Labels: {', '.join(labels) if labels else '(none)'}",
]
)
def escape_formula(value):
return str(value).replace("\\", "\\\\").replace('"', '\\"')
def find_record(sync_key, url):
formula = (
f'OR({{{F_REFERENCE}}}="{escape_formula(sync_key)}",'
f' {{{F_ORIGINAL_Q}}}="{escape_formula(url)}")'
)
query = urllib.parse.urlencode({"filterByFormula": formula, "maxRecords": "10"})
records = airtable(f"{table_path()}?{query}").get("records", [])
if len(records) > 1:
ids = ", ".join(r["id"] for r in records)
warn(f"{len(records)} rows match {sync_key} ({ids}); updating the first.")
return records[0] if records else None
def desired_state(action, current):
"""Follow the issue's open/closed state without discarding another value."""
values = current if isinstance(current, list) else [current] if current else []
is_complete = any(str(v).strip().lower() in COMPLETE_STATES for v in values)
if action == "closed":
return None if is_complete else STATE_CLOSED
if action == "reopened":
return STATE_OPEN if (not values or is_complete) else None
return None
def build_row(schema, existing, issue, repo_full_name, action):
url = issue["html_url"]
title = issue.get("title") or f"Issue #{issue['number']}"
body = (issue.get("body") or "").strip() or "(No description was provided.)"
login = (issue.get("user") or {}).get("login")
row = Row(schema, existing)
row.own(F_ORIGINAL_Q, url)
row.own(F_REFERENCE, f"{repo_full_name}#{issue['number']}")
row.add(F_SOURCE, [SOURCE_GITHUB, SOURCE_AUTOMATED])
row.own(F_GITHUB_ID, login)
row.fill(F_DATE, issue.get("created_at"))
row.fill(F_DESCRIPTION, truncate(f"{title}\n\n{body}", url))
row.fill(F_DETAILS, details_block(issue, repo_full_name))
row.fill(F_TYPE, request_types(issue))
row.fill(F_PRODUCT, products(issue))
state = (
desired_state(action, row.existing.get(F_COMPLETE)) if existing else STATE_OPEN
)
if state:
row.own(F_COMPLETE, [state])
return row.fields
def main():
with open(require_env("GITHUB_EVENT_PATH"), encoding="utf-8") as fh:
event = json.load(fh)
action = event.get("action")
issue = event.get("issue")
if not issue:
print(f'No issue in payload (action="{action}"); nothing to do.')
return
if (issue.get("user") or {}).get("type") == "Bot":
print(f"Issue #{issue['number']} was opened by a bot; skipping.")
return
repo = event.get("repository") or {}
repo_full_name = repo.get("full_name") or os.environ.get("GITHUB_REPOSITORY", "")
sync_key = f"{repo_full_name}#{issue['number']}"
schema = get_schema()
existing = find_record(sync_key, issue["html_url"])
fields = build_row(schema, existing, issue, repo_full_name, action)
if not fields:
print(f'{sync_key}: nothing to change (action="{action}").')
return
if DRY_RUN:
target = existing["id"] if existing else "(new record)"
print(f"DRY_RUN {sync_key} -> {target}\n{json.dumps(fields, indent=2)}")
return
if existing:
airtable(f"{table_path()}/{existing['id']}", "PATCH", {"fields": fields})
print(
f'Updated {sync_key} as {existing["id"]} (action="{action}"): '
f"{', '.join(sorted(fields))}."
)
else:
resp = airtable(table_path(), "POST", {"fields": fields})
print(f'Created {sync_key} as {resp["id"]} (action="{action}").')
if __name__ == "__main__":
try:
main()
except Exception as err: # noqa: BLE001
sys.exit(f"Sync failed: {err}")
================================================
FILE: .github/workflows/airtable-issue-sync.yaml
================================================
name: Sync issues to Airtable
on:
issues:
types: [opened, edited, reopened, closed, labeled, unlabeled]
permissions:
contents: read
concurrency:
group: airtable-sync-${{ github.event.issue.number }}
cancel-in-progress: false
jobs:
sync:
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.11"
- name: Sync issue to Airtable
run: python .github/scripts/airtable_issue_sync.py
env:
AIRTABLE_TOKEN: ${{ secrets.AIRTABLE_TOKEN }}
# "All Feedback Post Launch". Pinned here rather than in a repository
# variable so the destination is reviewable with the code that writes
# to it; neither id grants access on its own.
AIRTABLE_BASE: app9zoO6eKABI4dog
AIRTABLE_TABLE: tbllUJRwhQg35gzbT
================================================
FILE: .github/workflows/ci.yml
================================================
name: ESM Tests
on:
pull_request:
branches:
- "**"
workflow_dispatch:
merge_group:
types: [checks_requested]
permissions:
contents: read
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
test-precommit:
runs-on: ubuntu-24.04
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Setup Environment
uses: prefix-dev/setup-pixi@v0.9.0
with:
pixi-version: v0.70.2
cache: false
cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
- name: Check formatting and typing
run: pixi run --environment dev lint-all
test-esm:
runs-on: ubuntu-24.04
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Setup Environment
uses: prefix-dev/setup-pixi@v0.9.0
with:
pixi-version: v0.70.2
cache: false
cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
- name: Run tests
run: pixi run --environment dev cov-test
- name: Run Docker tests
env:
DOCKER_TAG: ${{ github.sha }}
BIOHUB_URL: https://biohub.ai/
ESM_API_KEY: ${{ secrets.ESM_API_KEY }}
run: |
set -e
cd tests
make build-oss-ci
make start-docker-oss URL=${{ env.BIOHUB_URL }} DOCKER_TAG=${{ env.DOCKER_TAG }} ESM_API_KEY=${{ env.ESM_API_KEY }}
shell: pixi run --environment dev bash -e {0}
- name: cleanup docker containers if they're hanging
if: ${{ always() }}
run: |
docker rm --force --volumes $(docker ps -q) || true
# Next two steps adds a comment to any PR that reports tests + code coverage.
- name: Pytest Coverage Comment
if: ${{ github.event_name == 'pull_request' }}
uses: MishaKav/pytest-coverage-comment@v1.1.47
id: coverageComment
with:
title: Coverage Report
pytest-coverage-path: pytest-coverage.txt
junitxml-path: pytest.xml
report-only-changed-files: true
- name: Pytest coverage GitHub summary
if: ${{ github.event_name == 'pull_request' }}
run: |
echo '${{ steps.coverageComment.outputs.coverageHtml }}' >> $GITHUB_STEP_SUMMARY
# `requires-python` no longer caps at 3.13, so the versions above it need
# exercising. Tests only: the Docker suite and the coverage comment stay on the
# 3.12 job, so they run once, against one environment and one set of secrets.
test-esm-python-versions:
runs-on: ubuntu-24.04
strategy:
fail-fast: false
matrix:
environment: [dev-py313, dev-py314]
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Setup Environment
uses: prefix-dev/setup-pixi@v0.9.0
with:
pixi-version: v0.70.2
environments: ${{ matrix.environment }}
cache: false
- name: Run tests
run: pixi run --environment ${{ matrix.environment }} cov-test
================================================
FILE: .gitignore
================================================
esm.egg-info
# pixi environments
.pixi
*.egg-info
*.pyc
# pytest --cov artifacts
.coverage
.coverage.*
coverage.xml
htmlcov/
pytest.xml
.ruff_cache/
================================================
FILE: .pre-commit-config.yaml
================================================
# See https://pre-commit.com for more information
# See https://pre-commit.com/hooks.html for more hooks
exclude: (fasta|pdb|cif|mds|json)$
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v3.2.0
hooks:
- id: trailing-whitespace
- id: end-of-file-fixer
- id: check-yaml
- id: check-added-large-files
exclude: pixi.lock
- id: check-merge-conflict
- repo: https://github.com/seddonym/import-linter
rev: v1.12.1
hooks:
- id: import-linter
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.16 # keep in sync with the monorepo .pre-commit-config.yaml
hooks:
- id: ruff-check
args: [ --fix ]
- id: ruff-format # formatter
types_or: [python, jupyter]
- repo: local
hooks:
- 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:
- id: gitleaks
================================================
FILE: CONTRIBUTIONS.md
================================================
We welcome community contributions to help make this package better!
## Contributing
_to be written_
## Testing
Improving test coverage and automating the process is high priority but not done yet. For now, to test:
0. Ensure you are testing in a clean environment - we use micromamba for this
```bash
micromamba create -n esm
micromamba activate esm
micromamba install -c conda-forge python=3.10
# in root level of repo
pip install -e .
pip install examples/requirements.txt
python -c 'from huggingface_hub import login; login()'
```
```bash
# Ensure package can correctly interact with forge.
# This will require an API key from forge.evolutionaryscale.ai
ESM_API_KEY=$ESM_API_KEY PYTHONPATH='.' python cookbook/snippets/esm3.py
```
1. Ensure the following scripts run without errors. Most have a pip install command installing the published `esm` package - comment this out so your release candidate version is tested and not the already published version.
```bash
ESM_API_KEY=$ESM_API_KEY PYTHONPATH='.' python cookbook/snippets/esm3.py
pip install treon
treon cookbook/tutorials/1_esmprotein.ipynb
treon cookbook/tutorials/2_embed.ipynb
treon cookbook/tutorials/3_gfp_design.ipynb
treon cookbook/tutorials/4_forge_generate.ipynb
# requires a GPU
python cookbook/snippets/esm3.py
python cookbook/snippets/esmc.py
python cookbook/local/raw_forwards.py
```
================================================
FILE: LICENSE.md
================================================
**License (MIT)**
Copyright 2026 Chan Zuckerberg Biohub, Inc.
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
================================================
FILE: README.md
================================================
<div align="center">
<img src="_assets/header.png" style="width: 60%; height: auto;" />
# A world model of protein biology: ESMC, ESMFold2, & ESM Atlas
[ESMC & ESMFold2 Preprint](https://www.biorxiv.org/content/10.64898/2026.06.03.729735) ⋅ [Atlas](https://biohub.ai/esm/protein/atlas) ⋅ [Tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials) ⋅ [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.
**[ESMC](https://biohub.ai/esm/protein)** is a state-of-the-art protein language model that has learned the rules of protein biology from training on billions of protein sequences. ESMC defines a new scaling frontier relative to ESM2, achieving stronger performance in emergent long-range structural understanding as model scale increases.
<div align="center">
<img src="_assets/esmc_graphic.png" width="40%"/>
</div>
**[ESMFold2](https://huggingface.co/biohub/ESMFold2)**, built on the ESMC 6B model, is a state-of-the-art structure prediction model that has been validated for the design of protein-protein interactions. ESMFold2 surpasses other models in DockQ pass-rate on Foldbench protein-protein and antibody-antigen complexes, and can be used in single-sequence mode for an order of magnitude speedup in folding.
<div align="center">
<img src="_assets/esmfold2_folding.png" width="60%"/>
</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://www.biorxiv.org/content/10.64898/2026.06.03.729735).
<div align="center">
<img src="_assets/esmfold2_binder.png" width="60%"/>
</div>
The **[ESM Atlas](https://biohub.ai/esm/protein/atlas)** is a map of 6.8 billion proteins covering the full breadth of life’s biodiversity. ESMFold2’s folding throughput enabled the prediction of more than one billion predicted structures. The Atlas is organized according to the internal world model of ESMC. We make this world model interpretable by training sparse autoencoders (SAEs). SAEs are unsupervised neural networks trained to decompose ESMC internal representations into a sparse set of ~16,000 interpretable features that reveal the functional relationships between proteins that ESMC has learned. Each feature is summarized in natural language with an agentic pipeline that maps features onto known biology from protein databases. We release a collection of SAEs trained on different model scales, layers, and at different levels of granularity. Learn more about how to use the ESM Atlas on the [Biohub Platform](https://biohub.ai/).
For information on using ESM3, see the [ESM3 README](https://github.com/Biohub/esm/blob/main/_assets/ESM3_README.md).
## Table of Contents
- [ESMC](#esmc)
- [ESMC Sparse Autoencoders](#esmc-sparse-autoencoders)
- [ESMFold2](#esmfold2)
- [Batch Inference](#batch-inference)
- [Frontier-Safety](#frontier-safety)
- [Licenses](#licenses)
- [Citations](#citations)
## ESMC
<a name="esmc"></a>
[ESMC](https://biohub.ai/esm/protein) is a state-of-the-art protein language model that has learned representations of protein biology from training on billions of protein sequences.
Codebase, model weights, and model variants for ESMC are available through [Hugging Face](https://huggingface.co/collections/biohub/esmc-model-family).
There are two primary ways of running the ESM models: through the [**Biohub Platform**](https://biohub.ai/) or locally with Hugging Face. The Biohub Platform enables users to easily run inference with ESM models with minimal setup. Users interested in customizing or fine-tuning ESM models can use the models from Hugging Face.
### Running ESMC Through Hugging Face
<a name="running-esmc-through-hugging-face"></a>
First, install `esm` from PyPI:
```
pip install esm
```
Then use the following code to run ESMC using weights from Hugging Face:
```python
import torch
from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer
# example GFP sequence
sequences = ["MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"]
model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-6B", device="cuda").eval()
tokenizer = EsmcTokenizer()
inputs = tokenizer(sequences, return_tensors="pt", padding=True)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.inference_mode():
output = model(**inputs)
```
By default, the model returns only the final layer representations. To return hidden states from **all transformer layers**, set:
```python
output = model(**inputs, output_hidden_states=True)
```
### Running ESMC Through the Biohub Platform
<a name="running-esmc-through-biohub-platform"></a>
The code below shows how to access ESMC using the Biohub Platform. API tokens can be created in the [developer console](https://biohub.ai/developer-console/api-keys).
Note that our API migrated from forge.evolutionaryscale.ai to [biohub.ai](https://biohub.ai), so some code classes reference “Forge”.
To get started with ESM, install the python library using `pip`:
```
pip install esm
```
Then import the necessary libraries and instantiate your desired model.
```py
from esm.sdk import esmc_client
from esm.sdk.api import ESMProtein, LogitsConfig
# Human carbonic anhydrase II (PDB 2CBA)
protein = ESMProtein(
sequence=(
"MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDD"
"SQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVL"
"GIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPPLLECVTWIVLKEP"
"ISVSSEQVLKFRKLNFNGEGEPEELMVDNWRPAQPLKNRQIKASFK"
)
)
model = esmc_client(
model="esmc-600m-2024-12", url="https://biohub.ai", token="<your API token>"
)
protein_tensor = model.encode(protein)
logits_output = model.logits(
protein_tensor, LogitsConfig(sequence=True, return_embeddings=True)
)
print(logits_output.logits, logits_output.embeddings)
```
For tutorials on how to use ESMC, see our [tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials).
## ESMC Sparse Autoencoders (SAE)
<a name="esmc-sparse-autoencoders"></a>
Sparse autoencoders (SAE) are an unsupervised method for decomposing representations of large transformer language models into interpretable units. We released SAEs trained on ESMC to reveal the interpretable units of functional organization that ESMC's world model has learned.
The sparse autoencoder used in the Atlas and analyzed in the paper, `ESMC-6B-sae-layer60-k64-codebook16384`, is built on the ESMC 6B model. We also provide human-interpretable, agent-generated feature descriptions for this SAE's codebook.
Codebase, model weights, and model variants for ESMC SAEs are available through [Hugging Face](https://huggingface.co/collections/biohub/esmc-saes-for-hidden-states-all-layers).
### Running SAEs Through Hugging Face
First, install `esm` from PyPI:
```
pip install esm
```
Then use the following code to set up an ESMC SAE using weights from Hugging Face:
```python
import torch
from esm.models.esmc import EsmcForMaskedLM, EsmcSaeModel, EsmcTokenizer
sequence = "MGSNKSKPKDASQRRRSLEPAENVHGAGGGAFPASQTPSKPASADGHRGPSAAFAPAAAEPKLFGGFNSSDTVTSPQRAGPLAGGVTTFVALYDYESRTETDLSFKKGERLQIVNNTEGDWWLAHSLSTGQTGYIPSNYVAPSDSIQAEEWYFGKITRRESERLLLNAENPRGTFLVRESETTKGAYCLSVSDFDNAKGLNVKHYKIRKLDSGGFYITSRTQFNSLQQLVAYYSKHADGLCHRLTTVCPTSKPQTQGLAKDAWEIPRESLRLEVKLGQGCFGEVWMGTWNGTTRVAIKTLKPGTMSPEAFLQEAQVMKKLRHEKLVQLYAVVSEEPIYIVTEYMSKGSLLDFLKGETGKYLRLPQLVDMAAQIASGMAYVERMNYVHRDLRAANILVGENLVCKVADFGLARLIEDNEYTARQGAKFPIKWTAPEAALYGRFTIKSDVWSFGILLTELTTKGRVPYPGMVNREVLDQVERGYRMPCPPECPESLHDLMCQCWRKEPEERPTFEYLQAFLEDYFTSTEPQYQPGENL"
model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-6B", device="cuda").eval()
tokenizer = EsmcTokenizer()
sae = EsmcSaeModel.from_pretrained(
"biohub/ESMC-6B-sae-k64-codebook16384",
allow_patterns=["config.json", "layer_30.safetensors", "layer_60.safetensors"],
device=model.device,
)
sae.initialize_layers([30, 60])
model.add_sae_models([sae.layers["30"], sae.layers["60"]])
inputs = tokenizer(sequence, return_tensors="pt", padding=True)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.inference_mode():
output = model(**inputs)
output.sae_outputs["layer60"] # sparse.coo tensor
print(output.sae_outputs["layer60"].shape)
```
### Running SAEs Through The Biohub Platform
For a tutorial on using SAEs using the Biohub Platform, see [here](https://github.com/Biohub/esm/blob/main/cookbook/tutorials/esmc_sae_feature_interpretation.ipynb).
## ESMFold2
<a name="esmfold2"></a>
[ESMFold2](https://huggingface.co/biohub/ESMFold2) is a state-of-the-art protein structure prediction model that combines ESMC (6B parameter) language model embeddings with a diffusion-based structure prediction architecture.
The model predicts high-resolution, all-atom 3D protein structures directly from amino acid sequences, with optional multiple sequence alignment (MSA) input for enhanced accuracy on challenging targets. ESMFold2 achieves state-of-the-art performance matching or exceeding AlphaFold3 across diverse evaluation datasets, while offering improved computational efficiency through optimized diffusion sampling and architectural innovations.
Codebase, model weights, and model variants for ESMFold2 are available through [Hugging Face](https://huggingface.co/biohub/ESMFold2)
### Running ESMFold2 Through Hugging Face
<a name="running-esmfold2-through-hugging-face"></a>
First, install `esm` from PyPI:
```
pip install esm
```
Then use the following code to run ESMFold2 locally using weights from Hugging Face:
```python
from esm.models.esmfold2 import (
DNAInput,
ESMFold2InputBuilder,
EsmFold2Model,
LigandInput,
Modification,
ProteinInput,
StructurePredictionInput,
)
HHAI_SEQ = (
"MIEIKDKQLTGLRFIDLFAGLGGFRLALESCGAECVYSNEWDKYAQEVYEMNFGEKPEGDITQVNEKTIPDH"
"DILCAGFPCQAFSISGKQKGFEDSRGTLFFDIARIVREKKPKVVFMENVKNFASHDNGNTLEVVKNTMNELD"
"YSFHAKVLNALDYGIPQKRERIYMICFRNDLNIQNFQFPKPFELNTFVKDLLLPDSEVEHLVIDRKDLVMTN"
"QEIEQTTPKTVRLGIVGKGGQGERIYSTRGIAITLSAYGGGIFAKTGGYLVNGKTRKLHPRECARVMGYPDS"
"YKVHPSTSQAYKQFGNSVVINVLQYIAYNIGSSLNFKPY"
)
model = EsmFold2Model.from_pretrained("biohub/ESMFold2", device="cuda").eval()
spi = StructurePredictionInput(
sequences=[
ProteinInput(id="A", sequence=HHAI_SEQ),
DNAInput(
id="B",
sequence="GATAGCGCTATC",
modifications=[Modification(position=5, ccd="C36")],
),
DNAInput(
id="C",
sequence="TGATAGCGCTATC",
modifications=[Modification(position=6, ccd="C36")],
),
LigandInput(id="L", ccd=["SAH"]),
]
)
result = ESMFold2InputBuilder().fold(
model, spi, num_loops=20, num_sampling_steps=100, num_diffusion_samples=1, seed=0
)
print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
with open("1mht_pred.cif", "w") as f:
f.write(result.complex.to_mmcif())
```
> **AMD ROCm users:** use ROCm 6.4 with PyTorch 2.9 or newer.
### Running ESMFold2 Through the Biohub Platform
Install the `esm` Python package
```
pip install esm
```
Import the necessary libraries.
```py
from esm.sdk.forge import SequenceStructureForgeInferenceClient
from esm.sdk.api import FoldingConfig
from esm.utils.structure.input_builder import ProteinInput, StructurePredictionInput
```
Call the inference client with the selected model of choice and replace <your API token> with your token name.
```py
client = SequenceStructureForgeInferenceClient(model="esmfold2-fast-2026-05", url="https://biohub.ai", token="<your API token>")
# Human carbonic anhydrase II (PDB 2CBA)
ca2_sequence = (
"MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDD"
"SQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVL"
"GIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPPLLECVTWIVLKEP"
"ISVSSEQVLKFRKLNFNGEGEPEELMVDNWRPAQPLKNRQIKASFK"
)
ca2_input = StructurePredictionInput(
sequences=[ProteinInput(id="A", sequence=ca2_sequence)]
)
config = FoldingConfig(
num_loops=20,
num_sampling_steps=100
)
result = client.fold_all_atom(ca2_input, config=config)
with open("result.cif", "w") as f:
f.write(result.complex.to_mmcif())
```
For tutorials on how to use ESMFold2, see our [tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials).
## Parallel Job Execution
For jobs that require processing multiple inputs, the Parallel Executor sends parallel requests, executing them concurrently and efficiently while respecting rate limits and adapting to request latency. The following example provides an example of using the parallel_executor context manager when embedding sequences.
```py
from esm.sdk.forge import ESMCForgeInferenceClient
from esm.sdk.api import ESMProtein, LogitsConfig, LogitsOutput, ESMProteinError
from esm.sdk import parallel_executor
def embed_sequence(client: ESMCForgeInferenceClient, sequence: str) -> LogitsOutput:
protein = ESMProtein(sequence=sequence)
protein_tensor = client.encode(protein)
if isinstance(protein_tensor, ESMProteinError):
raise protein_tensor
output = client.logits(protein_tensor, LogitsConfig(sequence=True, return_embeddings=True))
return output
sequences = ["A", "AA", "AAA"]
client = ESMCForgeInferenceClient(model="esmc-6b-2024-12", url="https://biohub.ai", token="<your API token>")
# Usage Example:
# To execute a batch job, wrap your function inside the batch executor context manager.
# Syntax:
# with parallel_executor() as executor:
# outputs = executor.execute_batch(user_func=<your_function>, **kwargs)
with parallel_executor() as executor:
outputs = executor.execute_batch(user_func=embed_sequence, client=client, sequence=sequences)
```
## Frontier Safety
<a name="frontier-safety"></a>
Biohub has established a safety team to assess the benefits and potential risks of our models and tools prior to release, and develop mitigations where necessary. To do this, we follow a structured approach that includes assessing both biosafety and biosecurity risks as well as existing, comparable open-source models and tools. We actively engage with the scientific community, stakeholders and domain experts to advance innovation as well as best practices for responsible development. Risk assessment was conducted for each of the components of this release, including our ESMC Cambrian models, ESMFold2, ESMC SAEs, ESM Atlas, and binder design system.
Informed by our risk assessments, we are releasing the source code and model weights for ESMC 6B, ESMFold2, and ESMC SAEs. We are also releasing our ESM Atlas dataset and binder design system openly. Biohub values open science, and we share our research with the scientific community so that others can evaluate, reproduce, and build upon our work.
Evaluations: Prior to release, we conducted evaluations to inform our understanding of capability uplift for specific misuse-relevant functional tasks. The full details of these evaluations are available in our corresponding paper appendix.
The Biohub Platform: We implement guardrails that detect and restrict the use of keywords and sequences corresponding to controlled pathogens and toxins on our freely accessible platform. For further details regarding these guardrails, please refer to our Biohub Platform Resources page. We recognize there are many legitimate reasons to use AI models to understand and model these sequences and proteins. If you are a researcher whose work is impacted by these guardrails, you can request elevated access to our platform via [biohub.ai](https://biohub.ai).
Please follow our [Acceptable Use Policy](https://biohub.org/acceptable-use-policy/) when using the model.
## Licenses
<a name="licenses"></a>
These models are available under the [MIT license](https://github.com/Biohub/esm/blob/main/LICENSE.md).
## Citations
<a name="citations"></a>
If you use ESM in your work, please cite one of the following:
#### ESMC, SAEs, and ESMFold2
```
@misc{candido2026language,
title = {Language Modeling Materializes a World Model of Protein Biology},
author = {Candido, Salvatore and Hayes, Thomas and Derry, Alexander and Rao, Roshan
and Lin, Zeming and Verkuil, Robert and Wu, Bryan and Lee, Jin Sub
and Bruguera, Elise S. and Keval, Jehan A. and Kopylov, Mykhailo
and Pak, John E. and Wu, Wesley and Thomas, Neil and Mataraso, Samson
and Hsu, Alvin and Trotman-Grant, Ashton C. and Fatras, Kilian
and dos Santos Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil
and Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth
and Sitwala, Hrishita and Yu, Yue and Wiggert, Marius
and Carlin, Dylan Alexander and Goering, Anthony W. and Blazejewski, Tomasz
and Sandora, McCullen and Hla, Michael and Jia, Tina Z.
and Kloker, Leon H. and Sofroniew, Nicholas J. and Uehara, Masatoshi
and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S.
and Sercu, Tom and Rives, Alexander},
year = {2026},
url = {https://www.biorxiv.org/content/10.64898/2026.06.03.729735},
note = {Preprint}
}
```
#### ESM3
```
@article {hayes2024simulating,
author = {Hayes, Thomas and Rao, Roshan and Akin, Halil and Sofroniew, Nicholas J. and Oktay, Deniz and Lin, Zeming and Verkuil, Robert and Tran, Vincent Q. and Deaton, Jonathan and Wiggert, Marius and Badkundri, Rohil and Shafkat, Irhum and Gong, Jun and Derry, Alexander and Molina, Raul S. and Thomas, Neil and Khan, Yousuf A. and Mishra, Chetan and Kim, Carolyn and Bartie, Liam J. and Nemeth, Matthew and Hsu, Patrick D. and Sercu, Tom and Candido, Salvatore and Rives, Alexander},
title = {Simulating 500 million years of evolution with a language model},
year = {2025},
doi = {10.1126/science.ads0018},
URL = {http://dx.doi.org/10.1126/science.ads0018},
journal = {Science}
}
```
#### ESM Github (Code / Weights)
```
@software{evolutionaryscale_2024,
author = {{EvolutionaryScale Team}},
title = {evolutionaryscale/esm},
year = {2024},
publisher = {Zenodo},
doi = {10.5281/zenodo.14219303},
URL = {https://doi.org/10.5281/zenodo.14219303}
}
```
================================================
FILE: THIRD_PARTY_NOTICE.md
================================================
The code in this repository depends on the following third-party libraries:
| Library | License | Link |
|----------|----------|----------|
| flash-attn | BSD | https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE |
| PyTorch | BSD | https://github.com/pytorch/pytorch/blob/main/LICENSE |
| xformers | BSD | https://github.com/facebookresearch/xformers/blob/main/LICENSE |
| jaxtyping | MIT | https://github.com/patrick-kidger/jaxtyping/blob/main/LICENSE |
| einops | MIT | https://github.com/arogozhnikov/einops/blob/main/LICENSE |
| omegaconf | BSD | https://github.com/omry/omegaconf/blob/master/LICENSE |
| attrs | MIT | https://github.com/python-attrs/attrs/blob/main/LICENSE |
| scipy | BSD-3-Clause | https://github.com/scipy/scipy/blob/main/LICENSE.txt<br>https://github.com/scipy/scipy/blob/main/LICENSES_bundled.txt |
| lightning / torchmetrics | Apache 2.0 | https://github.com/Lightning-AI/torchmetrics/blob/master/LICENSE |
================================================
FILE: _assets/ESM3_README.md
================================================
# ESM3 README
[ESM3](https://www.science.org/doi/10.1126/science.ads0018) is a frontier generative model for biology, able to jointly reason across three fundamental biological properties of proteins: sequence, structure, and function. These three data modalities are represented as tracks of discrete tokens at the input and output of ESM3. You can present the model with a combination of partial inputs across the tracks, and ESM3 will provide output predictions for all the tracks.
ESM3 is a *generative* masked language model. You can prompt it with partial sequence, structure, and function keywords, and iteratively sample masked positions until all positions are unmasked. This iterative sampling is what the `.generate()` function does.
<!---->
<img src="./esm3_diagram.png" alt="ESM3 Diagram" width="400" />
The ESM3 architecture is highly scalable due to its transformer backbone and all-to-all reasoning over discrete token sequences. At its largest scale, ESM3 was trained with 1.07e24 FLOPs on 2.78 billion proteins and 771 billion unique tokens, and has 98 billion parameters.
Learn more by reading the paper [(Hayes et al., 2024)](https://www.science.org/doi/10.1126/science.ads0018).
## [ESM3 Family](https://huggingface.co/collections/biohub/esm3-model-family)
The code for ESM3 is available from Github and weights for esm3-sm-open-v1 is available on [Hugging Face](https://huggingface.co/collections/biohub/esm3-model-family). Other weights are available when the model is accessed through Biohub.
| Model | Model Size | Release Date | Note |
| :---- | :---- | :---- | :---- |
| **Flagship Models** | | | Most users will be interested in using one of these models. |
| esm3-large-2024-03 | 98B | 2024-03 | |
| esm3-medium-2024-08 | 7B | 2024-08 | |
| esm3-small-2024-08 | 1.4B | 2024-08 | |
| **Published Models** | | | These models were used to generate all of the results in the ESM3 paper and are provided to facilitate reproducibility. |
| esm3-large-2024-03 | 98B | 2024-03 | |
| esm3-medium-2024-03 | 7B | 2024-03 | |
| esm3-small-2024-03 | 1.4B | 2024-03 | |
## Quickstart for ESM3
<a name="quickstart-esm3"></a>
### Running ESM3 Through Biohub
First install the python library using `pip`:
```
pip install esm
```
Then import the necessary libraries and instantiate your model. Use your token from the [Biohub platform](https://biohub.ai")
```py
from esm.sdk.forge import ESM3ForgeInferenceClient
from esm.sdk import client
from esm.sdk.api import ESMProtein, ESMProteinError, LogitsConfig, LogitsOutput
model: ESM3InferenceClient = esm.sdk.client("esm3-medium-2024-08", token="<your API token>")
```
### Running ESM3 Locally
The following code demonstrates how to run ESM3 locally and generate a simple sequence prompt. The weights are stored on [Hugging Face](https://huggingface.co/biohub/esm3-sm-open-v1).
First install the python library using `pip`:
```
pip install esm
```
Then import the necessary libraries for your model.
```py
from huggingface_hub import login
from esm.models.esm3 import ESM3
from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig
# Will instruct you how to get an API key from huggingface hub, make one with "Read" permission.
login()
# This will download the model weights and instantiate the model on your machine.
model: ESM3InferenceClient = ESM3.from_pretrained("esm3-sm-open-v1").to("cuda") # or "cpu"
# Generate a completion for a partial Carbonic Anhydrase (2vvb)
prompt = "___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________"
protein = ESMProtein(sequence=prompt)
# Generate the sequence, then the structure. This will iteratively unmask the sequence track.
protein = model.generate(protein, GenerationConfig(track="sequence", num_steps=8, temperature=0.7))
# We can show the predicted structure for the generated sequence.
protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
protein.to_pdb("./generation.pdb")
# Then we can do a round trip design by inverse folding the sequence and recomputing the structure
protein.sequence = None
protein = model.generate(protein, GenerationConfig(track="sequence", num_steps=8))
protein.coordinates = None
protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
protein.to_pdb("./round_tripped.pdb")
```
## Tutorials for ESM3
<a name="tutorials-esm3"></a>
For tutorials on how to use ESM3, see our Tutorials [here](https://github.com/Biohub/esm/tree/main/cookbook/tutorials).
## Responsible Development
<a name="responsible-development"></a>
Biohub has established a safety team to assess the benefits and potential risks of our models and tools prior to release, and develop mitigations where necessary. To do this, we follow a structured approach that includes assessing both biosafety and biosecurity risks as well as existing, comparable open-source models and tools. We actively engage with the scientific community, stakeholders and domain experts to advance innovation as well as best practices for responsible development. Risk assessment was conducted for ESM3.
Please follow our [Acceptable Use Policy](https://biohub.org/acceptable-use-policy/) when using the model.
## Licenses
<a name="licenses"></a>
These models are available under the [MIT license](https://github.com/Biohub/esm/blob/main/LICENSE.md).
## Citations
<a name="citations"></a>
If you use ESM in your work, please cite one of the following:
#### ESM3
```
@article {hayes2024simulating,
author = {Hayes, Thomas and Rao, Roshan and Akin, Halil and Sofroniew, Nicholas J. and Oktay, Deniz and Lin, Zeming and Verkuil, Robert and Tran, Vincent Q. and Deaton, Jonathan and Wiggert, Marius and Badkundri, Rohil and Shafkat, Irhum and Gong, Jun and Derry, Alexander and Molina, Raul S. and Thomas, Neil and Khan, Yousuf A. and Mishra, Chetan and Kim, Carolyn and Bartie, Liam J. and Nemeth, Matthew and Hsu, Patrick D. and Sercu, Tom and Candido, Salvatore and Rives, Alexander},
title = {Simulating 500 million years of evolution with a language model},
year = {2025},
doi = {10.1126/science.ads0018},
URL = {http://dx.doi.org/10.1126/science.ads0018},
journal = {Science}
}
```
================================================
FILE: cookbook/local/README.md
================================================
Examples utilizing the open model run locally.
================================================
FILE: cookbook/local/open_generate.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# ESM3\n",
"\n",
"ESM3 is a frontier generative model for biology, able to jointly reason across three fundamental biological properties of proteins: sequence, structure, and function. These three data modalities are represented as tracks of discrete tokens at the input and output of ESM3. You can present the model with a combination of partial inputs across the tracks, and ESM3 will provide output predictions for all the tracks.\n",
"\n",
"ESM3 is a generative masked language model. You can prompt it with partial sequence, structure, and function keywords, and iteratively sample masked positions until all positions are unmasked. This iterative sampling is what the `.generate()` function does.\n",
"\n",
"\n",
"\n",
"The ESM3 architecture is highly scalable due to its transformer backbone and all-to-all reasoning over discrete token sequences. At its largest scale, ESM3 was trained with 1.07e24 FLOPs on 2.78 billion proteins and 771 billion unique tokens, and has 98 billion parameters.\n",
"Here we present `esm3-open-small`. With 1.4B parameters it is the smallest and fastest model in the family, trained specifically to be open sourced. ESM3-open is available under a non-commercial license.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Imports\n",
"\n",
"If you're running in Colab, you probably want to get a GPU runtime first (Runtime > Change runtime type > T4 GPU).\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%set_env TOKENIZERS_PARALLELISM=false\n",
"!pip install esm\n",
"import numpy as np\n",
"import torch\n",
"\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"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Load `esm-open-small` on GPU\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from esm.utils.misc import huggingfacehub_login\n",
"\n",
"huggingfacehub_login() # will prompt you to get an API key and accept the ESM3 license.\n",
"model = ESM3.from_pretrained(\"esm3_sm_open_v1\", device=torch.device(\"cuda\"))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Alternatively, you could use the Forge API running the model remotely, and use the local `client` to call the API just like you're used to with the model running locally on your GPU:\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# from getpass import getpass\n",
"# token = getpass(\"Token from Forge console: \")\n",
"# model = client(\n",
"# model=\"esm3-large-2024-03\",\n",
"# url=\"https://biohub.ai\",\n",
"# token=token,\n",
"# )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Let's construct a prompt for ESM3, focusing on the task of scaffolding a motif from a natural protein\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"First, we can use the `ProteinChain` class from the `esm` sdk to grab a protein structure from the PDB.\n",
"We'll work with a human renal (kidney) dipeptidase (a protein that breaks up two amino acids bound together). Renal dipeptidases are of particular interest because they metabolize certain antibiotics.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pdb_id = \"1ITU\" # PDB ID corresponding to Renal Dipeptidase\n",
"chain_id = \"A\" # Chain ID corresponding to Renal Dipeptidase in the PDB structure\n",
"renal_dipep_chain = ProteinChain.from_rcsb(pdb_id, chain_id)\n",
"# Alternatively, we could have used ProteinChain.from_pdb() to load a protein structure from a local PDB file"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `ProteinChain` class is a object that makes it easy to work with protein structures. It contains a `sequence` attribute that contains the amino acid sequence of the protein\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(renal_dipep_chain.sequence)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"`ProteinChain` also contains an `atom37_positions` numpy array that contains the atomic coordinates of each of the residues in the protein.\n",
"\n",
"The shape of the array is `(n_residues, 37, 3)` where `n_residues` is the number of residues in the protein and 37 is the number of possible distinct atoms that may be present across all amino acids (e.g. the first three atoms are the N, C-alpha, and C atoms corresponding to the protein backbone). The 3 corresponds to the x, y, and z coordinates of each atom. The atom37 representation of protein structure allows us to use a single format to conveniently represent all amino acids -- **coordinates are only present for the atoms that are present in the amino acid and `nan` otherwise**.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"atom37_positions shape: \", renal_dipep_chain.atom37_positions.shape)\n",
"print(renal_dipep_chain.atom37_positions[:3])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can visualize the protein chain using the `py3Dmol` library\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# First we can create a `py3Dmol` view object\n",
"view = py3Dmol.view(width=500, height=500)\n",
"# py3Dmol requires the atomic coordinates to be in PDB format, so we convert the `ProteinChain` object to a PDB string\n",
"pdb_str = renal_dipep_chain.to_pdb_string()\n",
"# Load the PDB string into the `py3Dmol` view object\n",
"view.addModel(pdb_str, \"pdb\")\n",
"# Set the style of the protein chain\n",
"view.setStyle({\"cartoon\": {\"color\": \"spectrum\"}})\n",
"# Zoom in on the protein chain\n",
"view.zoomTo()\n",
"# Display the protein chain\n",
"view.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, let's try to scaffold a motif from this protein using ESM3 -- we'll prompt the model with the sequence and structure of a helix-coil motif from renal dipeptidase and have the model generate a larger scaffold that includes the motif\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"motif_inds = np.arange(123, 146)\n",
"# `ProteinChain` objects can be indexed like numpy arrays to extract the sequence and atomic coordinates of a subset of residues\n",
"motif_sequence = renal_dipep_chain[motif_inds].sequence\n",
"motif_atom37_positions = renal_dipep_chain[motif_inds].atom37_positions\n",
"print(\"Motif sequence: \", motif_sequence)\n",
"print(\"Motif atom37_positions shape: \", motif_atom37_positions.shape)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can also visualize the motif in the original chain using `py3Dmol`. We'll color the original chain in grey and the motif in blue\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"view = py3Dmol.view(width=500, height=500)\n",
"view.addModel(pdb_str, \"pdb\")\n",
"view.setStyle({\"cartoon\": {\"color\": \"lightgrey\"}})\n",
"motif_res_inds = (\n",
" motif_inds + 1\n",
").tolist() # residue indices are 1-indexed in PDB files, so we add 1 to the indices\n",
"view.addStyle({\"resi\": motif_res_inds}, {\"cartoon\": {\"color\": \"cyan\"}})\n",
"view.zoomTo()\n",
"view.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we can use the `ESMProtein` class to construct a prompt that will instruct ESM3 to scaffold the motif\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"prompt_length = 200\n",
"# First, we can construct a sequence prompt of all masks\n",
"sequence_prompt = [\"_\"] * prompt_length\n",
"# Then, we can randomly insert the motif sequence into the prompt (we randomly choose 72 here)\n",
"sequence_prompt[72 : 72 + len(motif_sequence)] = list(motif_sequence)\n",
"sequence_prompt = \"\".join(sequence_prompt)\n",
"print(\"Sequence prompt: \", sequence_prompt)\n",
"print(\"Length of sequence prompt: \", len(sequence_prompt))\n",
"\n",
"# Next, we can construct a structure prompt of all nan coordinates\n",
"structure_prompt = torch.full((prompt_length, 37, 3), np.nan)\n",
"# Then, we can insert the motif atomic coordinates into the prompt, starting at index 72\n",
"structure_prompt[72 : 72 + len(motif_atom37_positions)] = torch.tensor(\n",
" motif_atom37_positions\n",
")\n",
"print(\"Structure prompt shape: \", structure_prompt.shape)\n",
"print(\n",
" \"Indices with structure conditioning: \",\n",
" torch.where(~torch.isnan(structure_prompt).any(dim=-1).all(dim=-1))[0].tolist(),\n",
")\n",
"\n",
"# Finally, we can use the ESMProtein class to compose the sequence and structure prompts into a single prompt that can be passed to ESM3\n",
"protein_prompt = ESMProtein(sequence=sequence_prompt, coordinates=structure_prompt)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we can use the `generate` method of the model to iteratively sample a protein sequence based on the prompt. Under the hood, the model performs num_steps forward passes and samples a set of tokens at each step until the chosen track being generated is fully unmasked.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# We'll have to first construct a `GenerationConfig` object that specifies the decoding parameters that we want to use\n",
"sequence_generation_config = GenerationConfig(\n",
" track=\"sequence\", # We want ESM3 to generate tokens for the sequence track\n",
" num_steps=sequence_prompt.count(\"_\")\n",
" // 2, # We'll use num(mask tokens) // 2 steps to decode the sequence\n",
" temperature=0.5, # We'll use a temperature of 0.5 to control the randomness of the decoding process\n",
")\n",
"\n",
"# Now, we can use the `generate` method of the model to decode the sequence\n",
"sequence_generation = model.generate(protein_prompt, sequence_generation_config)\n",
"print(\"Sequence Prompt:\\n\\t\", protein_prompt.sequence)\n",
"print(\"Generated sequence:\\n\\t\", sequence_generation.sequence)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can also use the `generate` method to predict the structure of the generated sequence by iteratively sampling structure tokens.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"structure_prediction_config = GenerationConfig(\n",
" track=\"structure\", # We want ESM3 to generate tokens for the structure track\n",
" num_steps=len(sequence_generation) // 8,\n",
" temperature=0.7,\n",
")\n",
"structure_prediction_prompt = ESMProtein(sequence=sequence_generation.sequence)\n",
"structure_prediction = model.generate(\n",
" structure_prediction_prompt, structure_prediction_config\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we can visualize the generated structure using `py3Dmol`. We'll visualize the generated structure (right, green) alongside the original structure (left, grey) from which the motif was drawn. The motif residues are colored in cyan.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Convert the generated structure to a back into a ProteinChain object\n",
"structure_prediction_chain = structure_prediction.to_protein_chain()\n",
"# Align the generated structure to the original structure using the motif residues\n",
"motif_inds_in_generation = np.arange(72, 72 + len(motif_sequence))\n",
"structure_prediction_chain.align(\n",
" renal_dipep_chain, mobile_inds=motif_inds_in_generation, target_inds=motif_inds\n",
")\n",
"crmsd = structure_prediction_chain.rmsd(\n",
" renal_dipep_chain, mobile_inds=motif_inds_in_generation, target_inds=motif_inds\n",
")\n",
"print(\n",
" \"cRMSD of the motif in the generated structure vs the original structure: \", crmsd\n",
")\n",
"\n",
"view = py3Dmol.view(width=1000, height=500, viewergrid=(1, 2))\n",
"view.addModel(pdb_str, \"pdb\", viewer=(0, 0))\n",
"view.addModel(structure_prediction_chain.to_pdb_string(), \"pdb\", viewer=(0, 1))\n",
"view.setStyle({\"cartoon\": {\"color\": \"lightgrey\"}}, viewer=(0, 0))\n",
"view.setStyle({\"cartoon\": {\"color\": \"lightgreen\"}}, viewer=(0, 1))\n",
"view.addStyle({\"resi\": motif_res_inds}, {\"cartoon\": {\"color\": \"cyan\"}}, viewer=(0, 0))\n",
"view.addStyle(\n",
" {\"resi\": (motif_inds_in_generation + 1).tolist()},\n",
" {\"cartoon\": {\"color\": \"cyan\"}},\n",
" viewer=(0, 1),\n",
")\n",
"view.zoomTo()\n",
"view.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Secondary Structure Editing Example: Helix Shortening\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we can try another generation task with ESM3. We'll use the secondary structure track, along with the sequence track, to shorten a helix-coil-helix region (residues 39-111) in a protein structure (colored in blue below)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"helix_shortening_chain = ProteinChain.from_rcsb(\"7XBQ\", \"A\")\n",
"view = py3Dmol.view(width=500, height=500)\n",
"view.addModel(helix_shortening_chain.to_pdb_string(), \"pdb\")\n",
"view.setStyle({\"cartoon\": {\"color\": \"lightgrey\"}})\n",
"helix_region = np.arange(38, 111) # zero-indexed\n",
"view.addStyle(\n",
" {\"resi\": (helix_region + 1).tolist()}, {\"cartoon\": {\"color\": \"lightblue\"}}\n",
")\n",
"view.zoomTo()\n",
"view.show()\n",
"helix_shortening_ss8 = \"CCCSHHHHHHHHHHHTTCHHHHHHHHHHHHHTCSSCCCCHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHTTCHHHHHHHHHHHHHHHHHHHHHHHHHHHHIIIIIGGGCCSHHHHHHHHHHHHHHHHHHHHHCCHHHHHHHHHHHHHHHHHHHHHHHHHSCTTCHHHHHHHHHHHHHIIIIICCHHHHHHHHHHHHHHHHTTCTTCCSSHHHHHHHHHHHHHHHHHHHC\"\n",
"print(\n",
" \"Secondary structure of protein: (H: Alpha Helix, E: Beta Strand, C: Coil) \\n\\t\",\n",
" helix_shortening_ss8,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The helix-coil-helix region in the original protein is 73 residues long. We will try to shorten it to 45 residues by prompting the model with partial sequence and secondary structure\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"shortened_region_length = 45\n",
"\n",
"# We'll construct a sequence prompt that masks the (shortened) helix-coil-helix region, but leaves the flanking regions unmasked\n",
"sequence_prompt = (\n",
" helix_shortening_chain.sequence[: helix_region[0]]\n",
" + \"_\" * shortened_region_length\n",
" + helix_shortening_chain.sequence[helix_region[-1] + 1 :]\n",
")\n",
"print(\"Sequence prompt:\\n\\t\", sequence_prompt)\n",
"\n",
"# We'll construct a secondary structure prompt that retains the secondary structure of the flanking regions, and shortens the lengths of helices in the helix-coil-helix region\n",
"ss8_prompt = (\n",
" helix_shortening_ss8[: helix_region[0]]\n",
" + (\n",
" ((shortened_region_length - 3) // 2) * \"H\"\n",
" + \"C\" * 3\n",
" + ((shortened_region_length - 3) // 2) * \"H\"\n",
" )\n",
" + helix_shortening_ss8[helix_region[-1] + 1 :]\n",
")\n",
"print(\"SS8 prompt:\\n\\t\", ss8_prompt)\n",
"print(\n",
" \"Proposed SS8 for shortened helix-coil-helix region:\\n\\t\",\n",
" \" \" * helix_region[0] + ss8_prompt[helix_region[0] : helix_region[0] + 45],\n",
")\n",
"\n",
"print(\"\")\n",
"print(\"Original sequence:\\n\\t\", helix_shortening_chain.sequence)\n",
"print(\"Original SS8:\\n\\t\", helix_shortening_ss8)\n",
"print(\n",
" \"Original SS8 for helix-coil-helix region:\\n\\t\",\n",
" \" \" * helix_region[0]\n",
" + helix_shortening_ss8[helix_region[0] : helix_region[-1] + 1],\n",
")\n",
"\n",
"\n",
"# We can again use the ESMProtein class to compose the sequence and secondary structure prompts into a single prompt that can be passed to ESM3\n",
"protein_prompt = ESMProtein(sequence=sequence_prompt, secondary_structure=ss8_prompt)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can again use the `generate` method of the model to iteratively decode a protein sequence based on the prompt\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"Generating protein sequence...\")\n",
"sequence_generation = model.generate(\n",
" protein_prompt,\n",
" GenerationConfig(\n",
" track=\"sequence\",\n",
" num_steps=protein_prompt.sequence.count(\"_\") // 2,\n",
" temperature=0.5,\n",
" ),\n",
")\n",
"print(\"Folding protein...\")\n",
"structure_prediction = model.generate(\n",
" ESMProtein(sequence=sequence_generation.sequence),\n",
" GenerationConfig(\n",
" track=\"structure\", num_steps=len(protein_prompt) // 4, temperature=0\n",
" ),\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we can visualize the generated structure using `py3Dmol`. We'll visualize the generated structure (right) alongside the original structure (left) from which the motif was drawn. The helix-coil-helix region in the original structure is colored in blue and the shortened region in the generated structure is colored in pink.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"predicted_chain = structure_prediction.to_protein_chain()\n",
"predicted_chain = predicted_chain.align(\n",
" helix_shortening_chain,\n",
" mobile_inds=np.arange(len(predicted_chain) - 120, len(predicted_chain)),\n",
" target_inds=np.arange(\n",
" len(helix_shortening_chain) - 120, len(helix_shortening_chain)\n",
" ),\n",
")\n",
"view = py3Dmol.view(width=1000, height=500, viewergrid=(1, 2))\n",
"view.addModel(helix_shortening_chain.to_pdb_string(), \"pdb\", viewer=(0, 0))\n",
"view.addModel(predicted_chain.to_pdb_string(), \"pdb\", viewer=(0, 1))\n",
"view.setStyle({\"cartoon\": {\"color\": \"lightgrey\"}})\n",
"view.addStyle(\n",
" {\"resi\": (helix_region + 1).tolist()},\n",
" {\"cartoon\": {\"color\": \"lightblue\"}},\n",
" viewer=(0, 0),\n",
")\n",
"view.addStyle(\n",
" {\"resi\": (np.arange(helix_region[0], helix_region[0] + 45) + 1).tolist()},\n",
" {\"cartoon\": {\"color\": \"pink\"}},\n",
" viewer=(0, 1),\n",
")\n",
"view.zoomTo()\n",
"view.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# SASA Editing Example: Exposing a buried helix\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's grab 1LBS from the PDB and visualize it using `py3Dmol`. 1LBS has an alternating alpha-beta sandwich fold, with a buried helix in the center, highlighted in red\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"lipase_chain = ProteinChain.from_rcsb(\"1LBS\", \"A\")\n",
"span_start = 105\n",
"span_end = 116\n",
"view = py3Dmol.view(width=500, height=500)\n",
"view.addModel(lipase_chain.to_pdb_string(), \"pdb\")\n",
"view.setStyle({\"cartoon\": {\"color\": \"lightgrey\"}})\n",
"view.addStyle(\n",
" {\"resi\": (np.arange(span_start, span_end) + 1).tolist()},\n",
" {\"cartoon\": {\"color\": \"red\"}},\n",
")\n",
"view.zoomTo()\n",
"view.show()\n",
"lipase_ss8 = \"CCSSCCCCSSCHHHHHHTEEETTBBTTBCSSEEEEECCTTCCHHHHHTTTHHHHHHHTTCEEEEECCTTTTCSCHHHHHHHHHHHHHHHHHHTTSCCEEEEEETHHHHHHHHHHHHCGGGGGTEEEEEEESCCTTCBGGGHHHHHTTCBCHHHHHTBTTCHHHHHHHHTTTTBCSSCEEEEECTTCSSSCCCCSSSTTSTTCCBTSEEEEHHHHHCTTCCCCSHHHHHBHHHHHHHHHHHHCTTSSCCGGGCCSTTCCCSBCTTSCHHHHHHHHSTHHHHHHHHHHSCCBSSCCCCCGGGGGGSTTCEETTEECCC\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can construct a multimodal prompt for ESM3 to instruct it to expose the buried helix as follows:\n",
"\n",
"1. Prompt with the **structure** of the buried helix highlighted in red -- this will prompt ESM3 to generate a protein that contains that same helix\n",
"2. Prompt with high **SASA** values for the residues in the buried helix -- this will prompt ESM3 to expose the helix to the surface of the protein\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"structure_prompt = torch.full((len(lipase_chain), 37, 3), torch.nan)\n",
"structure_prompt[span_start:span_end] = torch.tensor(\n",
" lipase_chain[span_start:span_end].atom37_positions, dtype=torch.float32\n",
")\n",
"\n",
"sasa_prompt = [None] * len(lipase_chain)\n",
"sasa_prompt[span_start:span_end] = [40.0] * (span_end - span_start)\n",
"\n",
"print(\"SASA prompt (just for buried region): \", sasa_prompt[span_start:span_end])\n",
"\n",
"protein_prompt = ESMProtein(\n",
" sequence=\"_\" * len(lipase_chain), coordinates=structure_prompt, sasa=sasa_prompt\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This is a more difficult task, so you may need to sample more generations from ESM before you find a solution. We'll sample 32 here and sort by the generations with the highest predicted TM-score (pTM) by ESM3.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generated_proteins = []\n",
"N_SAMPLES = 16\n",
"for i in range(N_SAMPLES):\n",
" print(\"Generating protein sequence...\")\n",
" sequence_generation = model.generate(\n",
" protein_prompt,\n",
" GenerationConfig(\n",
" track=\"sequence\", num_steps=len(protein_prompt) // 8, temperature=0.7\n",
" ),\n",
" )\n",
" print(\"Folding protein...\")\n",
" structure_prediction = model.generate(\n",
" ESMProtein(sequence=sequence_generation.sequence),\n",
" GenerationConfig(track=\"structure\", num_steps=len(protein_prompt) // 32),\n",
" )\n",
" generated_proteins.append(structure_prediction)\n",
"\n",
"# Sort generations by ptm\n",
"generated_proteins = sorted(\n",
" generated_proteins, key=lambda x: x.ptm.item(), reverse=True\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's visualize the top 4 generations by pTM, alongside with the original protein (on the left)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"N_SAMPLES_TO_SHOW = 4\n",
"view = py3Dmol.view(width=1000, height=500, viewergrid=(1, N_SAMPLES_TO_SHOW + 1))\n",
"view.addModel(lipase_chain.to_pdb_string(), \"pdb\", viewer=(0, 0))\n",
"for i in range(N_SAMPLES_TO_SHOW):\n",
" print(\n",
" \"PTM of generated protein {}: {:.2f}\".format(\n",
" i + 1, generated_proteins[i].ptm.item()\n",
" )\n",
" )\n",
" view.addModel(\n",
" generated_proteins[i].to_protein_chain().to_pdb_string(),\n",
" \"pdb\",\n",
" viewer=(0, i + 1),\n",
" )\n",
"view.setStyle({\"cartoon\": {\"color\": \"lightgrey\"}})\n",
"view.addStyle(\n",
" {\"resi\": (np.arange(span_start, span_end) + 1).tolist()},\n",
" {\"cartoon\": {\"color\": \"red\"}},\n",
")\n",
"view.zoomTo()\n",
"view.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
================================================
FILE: cookbook/local/raw_forwards.py
================================================
import random
import torch
import torch.nn.functional as F
from esm.pretrained import (
ESM3_function_decoder_v0,
ESM3_sm_open_v0,
ESM3_structure_decoder_v0,
ESM3_structure_encoder_v0,
)
from esm.tokenization import get_esm3_model_tokenizers
from esm.tokenization.function_tokenizer import (
InterProQuantizedTokenizer as EsmFunctionTokenizer,
)
from esm.tokenization.sequence_tokenizer import EsmSequenceTokenizer
from esm.utils.structure.protein_chain import ProteinChain
from esm.utils.types import FunctionAnnotation
@torch.no_grad()
def inverse_folding_example():
tokenizer = EsmSequenceTokenizer()
encoder = ESM3_structure_encoder_v0("cuda")
model = ESM3_sm_open_v0("cuda")
chain = ProteinChain.from_rcsb("1utn", "A")
coords, plddt, residue_index = chain.to_structure_encoder_inputs()
coords = coords.cuda()
plddt = plddt.cuda()
residue_index = residue_index.cuda()
_, structure_tokens = encoder.encode(coords, residue_index=residue_index)
# Add BOS/EOS padding
coords = F.pad(coords, (0, 0, 0, 0, 1, 1), value=torch.inf)
plddt = F.pad(plddt, (1, 1), value=0)
structure_tokens = F.pad(structure_tokens, (1, 1), value=0)
structure_tokens[:, 0] = 4098
structure_tokens[:, -1] = 4097
output = model.forward(
structure_coords=coords, per_res_plddt=plddt, structure_tokens=structure_tokens
)
sequence_tokens = torch.argmax(output.sequence_logits, dim=-1)
sequence = tokenizer.decode(sequence_tokens[0])
print(sequence)
@torch.no_grad()
def conditioned_prediction_example():
tokenizers = get_esm3_model_tokenizers()
model = ESM3_sm_open_v0("cuda")
# PDB 1UTN
sequence = "MKTFIFLALLGAAVAFPVDDDDKIVGGYTCGANTVPYQVSLNSGYHFCGGSLINSQWVVSAAHCYKSGIQVRLGEDNINVVEGNEQFISASKSIVHPSYNSNTLNNDIMLIKLKSAASLNSRVASISLPTSCASAGTQCLISGWGNTKSSGTSYPDVLKCLKAPILSDSSCKSAYPGQITSNMFCAGYLEGGKDSCQGDSGGPVVCSGKLQGIVSWGSGCAQKNKPGVYTKVCNYVSWIKQTIASN"
tokens = tokenizers.sequence.encode(sequence)
# Calculate the number of tokens to replace, excluding the first and last token
num_to_replace = int((len(tokens) - 2) * 0.75)
# Randomly select indices to replace, excluding the first and last index
indices_to_replace = random.sample(range(1, len(tokens) - 1), num_to_replace)
# Replace selected indices with 32
assert tokenizers.sequence.mask_token_id is not None
for idx in indices_to_replace:
tokens[idx] = tokenizers.sequence.mask_token_id
sequence_tokens = torch.tensor(tokens, dtype=torch.int64)
function_annotations = [
# Peptidase S1A, chymotrypsin family
FunctionAnnotation(label="peptidase", start=100, end=114),
FunctionAnnotation(label="chymotrypsin", start=190, end=202),
]
function_tokens = tokenizers.function.tokenize(function_annotations, len(sequence))
function_tokens = tokenizers.function.encode(function_tokens)
function_tokens = function_tokens.cuda().unsqueeze(0)
sequence_tokens = sequence_tokens.cuda().unsqueeze(0)
output = model.forward(
sequence_tokens=sequence_tokens, function_tokens=function_tokens
)
return sequence, output, sequence_tokens
@torch.no_grad()
def decode(sequence, output, sequence_tokens):
# To save on VRAM, we load these in separate functions
decoder = ESM3_structure_decoder_v0("cuda")
function_decoder = ESM3_function_decoder_v0("cuda")
function_tokenizer = EsmFunctionTokenizer()
# Generally not recommended to just argmax the logits, decode iteratively!
# For quick demonstration only:
structure_tokens = torch.argmax(output.structure_logits, dim=-1)
structure_tokens = (
structure_tokens.where(sequence_tokens != 0, 4098) # BOS
.where(sequence_tokens != 2, 4097) # EOS
.where(sequence_tokens != 31, 4100) # Chainbreak
)
bb_coords = (
decoder.decode(
structure_tokens,
torch.ones_like(sequence_tokens),
torch.zeros_like(sequence_tokens),
)["bb_pred"]
.detach()
.cpu()
)
chain = ProteinChain.from_backbone_atom_coordinates(
bb_coords, sequence="X" + sequence + "X"
)
chain.infer_oxygen().to_pdb("hello.pdb")
# Function prediction
p_none_threshold = 0.05
log_p = F.log_softmax(output.function_logits[:, 1:-1, :], dim=3).squeeze(0)
# Choose which positions have no predicted function.
log_p_nones = log_p[:, :, function_tokenizer.vocab_to_index["<none>"]]
p_none = torch.exp(log_p_nones).mean(dim=1) # "Ensemble of <none> predictions"
where_none = p_none > p_none_threshold # (length,)
log_p[~where_none, :, function_tokenizer.vocab_to_index["<none>"]] = -torch.inf
function_token_ids = torch.argmax(log_p, dim=2)
function_token_ids[where_none, :] = function_tokenizer.vocab_to_index["<none>"]
predicted_function = function_decoder.decode(
function_token_ids,
tokenizer=function_tokenizer,
annotation_threshold=0.1,
annotation_min_length=5,
annotation_gap_merge_max=3,
)
print("function prediction:")
print(predicted_function["interpro_preds"].nonzero())
print(predicted_function["function_keywords"])
if __name__ == "__main__":
inverse_folding_example()
sequence, output, sequence_tokens = conditioned_prediction_example()
torch.cuda.empty_cache()
# And then decode from tokenized representation to outputs:
decode(sequence, output, sequence_tokens)
================================================
FILE: cookbook/snippets/README.md
================================================
Snippets of ESM3 usage that you can copy and paste directly into your scripts.
================================================
FILE: cookbook/snippets/esm3.py
================================================
import os
from esm.models.esm3 import ESM3
from esm.sdk import client
from esm.sdk.api import (
ESM3InferenceClient,
ESMProtein,
ESMProteinError,
ESMProteinTensor,
GenerationConfig,
LogitsConfig,
LogitsOutput,
SamplingConfig,
SamplingTrackConfig,
)
from esm.utils.structure.protein_chain import ProteinChain
from esm.utils.structure.protein_complex import ProteinComplex
from esm.utils.types import FunctionAnnotation
def get_sample_protein() -> ESMProtein:
protein = ProteinChain.from_rcsb("1utn")
protein = ESMProtein.from_protein_chain(protein)
protein.function_annotations = [
# Peptidase S1A, chymotrypsin family: https://www.ebi.ac.uk/interpro/structure/PDB/1utn/
FunctionAnnotation(label="peptidase", start=100, end=114),
FunctionAnnotation(label="chymotrypsin", start=190, end=202),
]
return protein
def get_sample_protein_complex() -> ESMProtein:
protein = ProteinComplex.from_rcsb("7a3w")
protein = ESMProtein.from_protein_complex(protein)
return protein
def main(client: ESM3InferenceClient):
# Single step decoding
protein = get_sample_protein()
protein.function_annotations = None
protein = client.encode(protein)
single_step_protein = client.forward_and_sample(
protein, SamplingConfig(structure=SamplingTrackConfig(topk_logprobs=2))
)
single_step_protein.protein_tensor.sequence = protein.sequence
single_step_protein = client.decode(single_step_protein.protein_tensor)
# Generate with partial sequence.
prompt = (
"___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLK"
"GGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIK"
"TKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________"
)
protein = ESMProtein(sequence=prompt)
protein = client.generate(
protein, GenerationConfig(track="sequence", num_steps=8, temperature=0.7)
)
assert isinstance(protein, ESMProtein), f"ESMProtein was expected but got {protein}"
# Folding
protein = get_sample_protein()
sequence_length = len(protein.sequence) # type: ignore
num_steps = int(sequence_length / 16)
protein.coordinates = None
protein.function_annotations = None
protein.sasa = None
folded_protein = client.generate(
protein,
GenerationConfig(track="structure", schedule="cosine", num_steps=num_steps),
)
assert isinstance(folded_protein, ESMProtein), (
f"ESMProtein was expected but got {protein}"
)
folded_protein.to_pdb("./sample_folded.pdb")
# Inverse folding
protein = get_sample_protein()
protein.sequence = None
protein.sasa = None
protein.function_annotations = None
inv_folded_protein = client.generate(
protein,
GenerationConfig(track="sequence", schedule="cosine", num_steps=num_steps),
)
assert isinstance(inv_folded_protein, ESMProtein)
inv_folded_protein.to_pdb("./sample_inv_folded.pdb")
# Function prediction
protein = get_sample_protein()
protein.function_annotations = None
protein_with_function = client.generate(
protein,
GenerationConfig(track="function", schedule="cosine", num_steps=num_steps),
)
assert isinstance(protein_with_function, ESMProtein), (
f"{protein_with_function} is not an ESMProtein"
)
# Logits
protein = get_sample_protein()
protein.coordinates = None
protein.function_annotations = None
protein.sasa = None
protein_tensor = client.encode(protein)
logits_output = client.logits(
protein_tensor, LogitsConfig(sequence=True, return_embeddings=True)
)
assert isinstance(logits_output, LogitsOutput), (
f"LogitsOutput was expected but got {logits_output}"
)
assert (
logits_output.logits is not None
and logits_output.logits.sequence is not None
and logits_output.embeddings is not None
)
# Chain of Thought (Function -> Secondary Structure -> Structure -> Sequence)
cot_protein = get_sample_protein()
cot_protein.sequence = "_" * len(cot_protein.sequence) # type: ignore
cot_protein.coordinates = None
cot_protein.sasa = None
cot_protein_tensor = client.encode(cot_protein)
for cot_track in ["secondary_structure", "structure", "sequence"]:
cot_protein_tensor = client.generate(
cot_protein_tensor,
GenerationConfig(track=cot_track, schedule="cosine", num_steps=10),
)
assert isinstance(cot_protein_tensor, ESMProteinTensor), (
f"ESMProteinTensor was expected but got {cot_protein_tensor}"
)
cot_protein = client.decode(cot_protein_tensor)
assert isinstance(cot_protein, ESMProtein), (
f"ESMProtein was expected but got {cot_protein}"
)
cot_protein.to_pdb("./sample_cot.pdb")
# Protein Complex
protein = get_sample_protein_complex()
sequence_length = len(protein.sequence) # type: ignore
num_steps = 1
folded_protein = client.generate(
protein,
GenerationConfig(
track="structure", schedule="cosine", num_steps=num_steps, temperature=0.0
),
)
assert isinstance(folded_protein, ESMProtein), (
f"ESMProtein was expected but got {protein}"
)
folded_protein.to_pdb("./sample_folded_complex.pdb")
# Batch examples.
# Batch generation.
prompts = [ESMProtein(sequence=("_" * (10 + 2 * i))) for i in range(5)]
configs = [
GenerationConfig(track="sequence", schedule="cosine", num_steps=(i + 1))
for i in range(5)
]
proteins = client.batch_generate(prompts, configs)
# Batch folding.
# Take the list of proteins batch generated from last step.
configs = [
GenerationConfig(track="structure", schedule="cosine", num_steps=(i + 1))
for i in range(5)
]
# Generate again for the structure track.
proteins = client.batch_generate(proteins, configs)
# Now write sequence and structure to PDB files.
for i, p in enumerate(proteins):
assert isinstance(p, ESMProtein), f"ESMProtein was expected but got {p}"
p.to_pdb(f"./batch_gen_{i}.pdb")
# Batch generation returns ESMProteinError for specific prompts that have issues.
prompts = [ESMProtein(sequence=("_" * (10 + 2 * i))) for i in range(5)]
# Mock error situation. The third prompt has no masks to be sampled.
prompts[2].sequence = "ANTVPYQ"
configs = [
GenerationConfig(track="sequence", schedule="cosine", num_steps=(i + 1))
for i in range(5)
]
proteins = client.batch_generate(prompts, configs)
# Should still get results. But third result is a ESMProteinError.
for i, p in enumerate(proteins):
if i == 2:
assert isinstance(p, ESMProteinError), (
f"ESMProteinError was expected but got {p}"
)
else:
assert isinstance(p, ESMProtein), f"ESMProtein was expected but got {p}"
if __name__ == "__main__":
if os.environ.get("ESM_API_KEY", ""):
print("ESM_API_KEY found. Trying to use model from Forge/Biohub Platform...")
main(client())
else:
print("No ESM_API_KEY found. Trying to load model locally...")
print(
"To try this script with a Forge/Biohub Platform API, please run ESM_API_KEY=your_api_key python esm3.py"
)
main(ESM3.from_pretrained("esm3_sm_open_v1"))
================================================
FILE: cookbook/snippets/esmc.py
================================================
import math
import os
import torch
from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer
from esm.sdk import esmc_client, parallel_executor
from esm.sdk.api import (
ESMCInferenceClient,
ESMProtein,
ESMProteinError,
ESMProteinTensor,
LogitsConfig,
LogitsOutput,
)
from esm.sdk.forge import ESM3ForgeInferenceClient, ESMCForgeInferenceClient
from esm.tokenization import get_esmc_model_tokenizers
def main(client: ESMCInferenceClient | ESM3ForgeInferenceClient):
# ================================================================
# Example usage: one single protein
# ================================================================
protein = ESMProtein(sequence="AAAAA")
# Use logits endpoint. Using bf16 for inference optimization
protein_tensor = client.encode(protein)
assert isinstance(protein_tensor, ESMProteinTensor), (
f"Expected ESMProteinTensor but got error: {protein_tensor}"
)
output = client.logits(
protein_tensor,
LogitsConfig(sequence=True, return_embeddings=True, return_hidden_states=True),
)
assert isinstance(output, LogitsOutput), (
f"LogitsOutput was expected but got error: {output}"
)
assert output.logits is not None and output.logits.sequence is not None
assert output.embeddings is not None
assert output.hidden_states is not None
print(
f"Client returned logits with shape: {output.logits.sequence.shape}, embeddings with shape: {output.embeddings.shape}, and hidden states with shape {output.hidden_states.shape}"
)
# request a specific hidden layer.
assert isinstance(protein_tensor, ESMProteinTensor), (
f"Expected ESMProteinTensor but got error: {protein_tensor}"
)
output = client.logits(
protein_tensor, LogitsConfig(return_hidden_states=True, ith_hidden_layer=1)
)
assert isinstance(output, LogitsOutput), (
f"LogitsOutput was expected but got error: {output}"
)
assert output.hidden_states is not None
print(f"Client returned hidden states with shape {output.hidden_states.shape}")
def raw_forward(model: EsmcForMaskedLM):
protein = ESMProtein(sequence="AAAAA")
assert protein.sequence is not None
sequences = [protein.sequence, protein.sequence]
# ================================================================
# Example usage: directly use the model
# ================================================================
# EsmcModel / EsmcForMaskedLM are the canonical raw ESMC models. They expose
# a Hugging Face style ``forward`` (not the SDK ``encode``/``logits``
# inference API - use ``esmc_client()`` for that). Tokenize with the ESMC
# tokenizer and pass ``input_ids`` directly.
tokenizer = EsmcTokenizer()
encoded = tokenizer(sequences, return_tensors="pt", padding=True)
output = model(
input_ids=encoded["input_ids"],
attention_mask=encoded["attention_mask"],
output_hidden_states=True,
)
logits, embeddings, hiddens = (
output.logits,
output.last_hidden_state,
output.hidden_states,
)
print(
f"Raw model returned logits with shape: {logits.shape}, embeddings with shape: {embeddings.shape} and hidden states with shape {hiddens.shape}"
)
def compute_pseudoperplexity(
forge_client: ESMCForgeInferenceClient, sequence: str
) -> float:
"""Compute L-pass pseudoperplexity for a protein sequence via Forge/Biohub Platform.
Masks each position one at a time, retrieves logits from Forge/Biohub Platform, and returns
exp(-mean(log_prob_true_aa)). Uses parallel_executor for parallel requests.
Example::
forge_client = ESMCForgeInferenceClient(
model="esmc-6b-2024-12",
url="https://biohub.ai",
token=os.environ["ESM_API_KEY"],
)
pppl = compute_pseudoperplexity(forge_client, "MKTLLILAVL...")
"""
L = len(sequence)
masked_sequences = [sequence[:i] + "_" + sequence[i + 1 :] for i in range(L)]
def _get_logits(client: ESMCForgeInferenceClient, sequence: str) -> LogitsOutput:
protein = ESMProtein(sequence=sequence)
protein_tensor = client.encode(protein)
if isinstance(protein_tensor, ESMProteinError):
raise protein_tensor
output = client.logits(protein_tensor, LogitsConfig(sequence=True))
if isinstance(output, ESMProteinError):
raise output
return output
with parallel_executor() as executor:
logit_outputs = executor.execute_batch(
_get_logits, client=forge_client, sequence=masked_sequences
)
# Build vocab from the tokenizer to map amino acid characters to token indices
vocab: dict[str, int] = get_esmc_model_tokenizers().get_vocab()
log_probs = []
for i in range(L):
output = logit_outputs[i]
if isinstance(output, Exception):
raise output
logits = output.logits.sequence # shape: (L+2, V)
position_logits = logits[i + 1] # +1 for BOS token
log_softmax = torch.log_softmax(position_logits, dim=-1)
true_aa_idx = vocab[sequence[i]]
log_probs.append(log_softmax[true_aa_idx].item())
return math.exp(-sum(log_probs) / L)
if __name__ == "__main__":
if os.environ.get("ESM_API_KEY", ""):
print("ESM_API_KEY found. Trying to use model from Forge/Biohub Platform...")
main(esmc_client(model="esmc-300m-2024-12"))
else:
print("No ESM_API_KEY found. Trying to load the model locally...")
print(
"To use the SDK inference API (encode/logits), run "
"ESM_API_KEY=your_api_key python esmc.py"
)
# The local raw model uses the Hugging Face style forward API rather
# than the SDK inference client.
model = EsmcForMaskedLM.from_pretrained("biohub/ESMC-300M")
raw_forward(model)
================================================
FILE: cookbook/snippets/fold_invfold.py
================================================
import os
from typing import cast
import numpy as np
from esm.sdk.api import (
ESM3InferenceClient,
ESMProtein,
GenerationConfig,
InverseFoldingConfig,
)
from esm.sdk.forge import (
ESM3ForgeInferenceClient,
SequenceStructureForgeInferenceClient,
)
from esm.utils.structure.protein_chain import ProteinChain
from esm.utils.types import FunctionAnnotation
def get_sample_protein() -> ESMProtein:
protein = ProteinChain.from_rcsb("1utn")
protein = ESMProtein.from_protein_chain(protein)
protein.function_annotations = [
# Peptidase S1A, chymotrypsin family: https://www.ebi.ac.uk/interpro/structure/PDB/1utn/
FunctionAnnotation(label="peptidase", start=100, end=114),
FunctionAnnotation(label="chymotrypsin", start=190, end=202),
]
return protein
def convert_none_to_nan(data):
"""Recursively convert None values in any deeply nested structure (e.g., list of lists of lists) to np.nan."""
if isinstance(data, list):
return [convert_none_to_nan(x) for x in data]
elif data is None:
return np.nan
else:
return data
def fold(
sequence_structure_client: SequenceStructureForgeInferenceClient,
esm3_client: ESM3InferenceClient,
):
protein = get_sample_protein()
protein.coordinates = None
protein.function_annotations = None
protein.sasa = None
assert protein.sequence is not None, "Protein sequence must be set to fold"
# Folding with esm3 client
config = GenerationConfig(track="structure", num_steps=1, temperature=0)
esm3_client_folded_protein = esm3_client.generate(protein, config)
assert isinstance(esm3_client_folded_protein, ESMProtein), (
f"Using ESM3 client, ESMProtein was expected but got {esm3_client_folded_protein}"
)
# Folding with folding client
sequence_structure_client_folded_protein = sequence_structure_client.fold(
protein.sequence, potential_sequence_of_concern=False
)
assert isinstance(sequence_structure_client_folded_protein, ESMProtein), (
f"Using sequence_structure client, ESMProtein was expected but got {sequence_structure_client_folded_protein}"
)
sequence_structure_client_folded_protein.to_pdb("folded_protein.pdb")
print("Saving folded protein to folded_protein.pdb")
def inverse_fold(
sequence_structure_client: SequenceStructureForgeInferenceClient,
esm3_client: ESM3InferenceClient,
):
protein = get_sample_protein()
protein.sequence = None
protein.sasa = None
protein.function_annotations = None
assert protein.coordinates is not None, (
"Protein coordinates must be set to inverse fold"
)
# Inverse Folding with esm3 client
config = GenerationConfig("sequence", num_steps=1, temperature=0.1)
esm3_client_inv_folded_protein = cast(
ESMProtein, esm3_client.generate(protein, config)
)
assert isinstance(esm3_client_inv_folded_protein, ESMProtein), (
f"Using ESM3 client, ESMProtein was expected but got {esm3_client_inv_folded_protein}"
)
# Inverse Folding with inverse folding client
sequence_structure_client_inv_folded_protein = (
sequence_structure_client.inverse_fold(
protein.coordinates,
config=InverseFoldingConfig(temperature=0.1),
potential_sequence_of_concern=False,
)
)
assert isinstance(sequence_structure_client_inv_folded_protein, ESMProtein), (
f"Using sequence_structure client, ESMProtein was expected but got {sequence_structure_client_inv_folded_protein}"
)
print(
f"Inverse folded protein: {sequence_structure_client_inv_folded_protein.sequence}"
)
if __name__ == "__main__":
if not os.environ.get("ESM_API_KEY", ""):
print(
"Please export your Forge/Biohub Platform API key as ESM_API_KEY environment variable."
)
client = SequenceStructureForgeInferenceClient(token=os.environ["ESM_API_KEY"])
esm3_client = ESM3ForgeInferenceClient(
model="esm3-medium-2024-08", token=os.environ["ESM_API_KEY"]
)
fold(client, esm3_client)
inverse_fold(client, esm3_client)
================================================
FILE: cookbook/snippets/sae.py
================================================
import numpy as np
import torch
from cookbook.snippets.sparse_utils import max_pool, remove_indexes
from esm.sdk import parallel_executor
from esm.sdk.api import ESMProtein, ESMProteinError, LogitsConfig, SAEConfig
from esm.sdk.forge import ESMCForgeInferenceClient
def get_sae_features_single(
client: ESMCForgeInferenceClient,
sae_config: SAEConfig,
sequence: str,
pool: bool = True,
) -> torch.Tensor:
protein = ESMProtein(sequence=sequence)
protein_tensor = client.encode(protein)
if isinstance(protein_tensor, ESMProteinError):
raise ValueError(
f"Error encoding sequence {sequence}: {protein_tensor.error_msg}"
)
# We wrap the SAEConfig in the LogitsConfig, which is normally used to return embeddings and hidden states.
output = client.logits(
protein_tensor, config=LogitsConfig(sae_config=sae_config), return_bytes=False
)
if isinstance(output, ESMProteinError):
raise ValueError(
f"Error getting SAE features for sequence {sequence}: {output.error_msg}"
)
if output.sae_outputs is None:
raise ValueError(f"SAE outputs missing for sequence {sequence}: {output}")
sae_tensor = output.sae_outputs[sae_config.models[0]]
if pool:
# Remove BOS / EOS tokens before pooling.
sae_features = remove_indexes(sae_tensor, {0, -1})
pooled_sae_features = max_pool(sae_features, axis=0)
return pooled_sae_features
else:
return sae_tensor
def get_sae_features(
client: ESMCForgeInferenceClient,
sae_config: SAEConfig,
sequences: list[str],
pool: bool = True,
) -> list[np.ndarray]:
with parallel_executor() as executor:
results = executor.execute_batch(
user_func=get_sae_features_single,
client=client,
sae_config=sae_config,
sequence=sequences,
pool=pool,
)
# Re-raise any errors from the batch
for result in results:
if isinstance(result, Exception):
raise result
return results
================================================
FILE: cookbook/snippets/sae_example.py
================================================
import os
from cookbook.snippets.sae import get_sae_features, get_sae_features_single
from cookbook.snippets.sparse_utils import remove_indexes
from esm.sdk.api import SAEConfig
from esm.sdk.forge import ESMCForgeInferenceClient
# Create ESMC 600M client
client = ESMCForgeInferenceClient(
model="esmc-600m-2024-12", url="https://biohub.ai", token=os.environ["ESM_API_KEY"]
)
# normalize feature activations by TF-IDF. Upweights activations
# of more highly specific features
sae_config = SAEConfig(
models=["esmc-600m-2024-12_k64_codebook16384_layer27"], normalize_features=True
)
# Create a protein
sequence = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQATHVDQWDWEWAGIKATEAFLPDYPDLDA"
sequences = [sequence] * 10
# get unpooled features for a single sequence
unpooled_features = get_sae_features_single(client, sae_config, sequence, pool=False)
print(f"Got unpooled SAE features with shape {unpooled_features.shape}")
print(f"is_sparse: {unpooled_features.is_sparse}")
print(f"layout: {unpooled_features.layout}")
# To remove bos/eos tokens efficiently from sparse tensors, we use a utility
unpooled_features = remove_indexes(unpooled_features, {0, -1})
print(
f"Unpooled SAE features after removing BOS/EOS have shape {unpooled_features.shape}"
)
# get pooled features for a batch
# this function pools by default to save memory.
features = get_sae_features(client, sae_config, sequences)
print(
f"Got SAE features for {len(features)} sequences, each with shape {features[0].shape}"
)
================================================
FILE: cookbook/snippets/sparse_utils.py
================================================
from typing import Iterable
import torch
def remove_indexes(
sparse_coo_tensor: torch.Tensor, indexes_to_remove: Iterable[int]
) -> torch.Tensor:
"""Remove entries at specified position indexes from sparse features.
This function removes positions and remaps the remaining indices to be contiguous.
For example, if we remove position 1 from a tensor with positions [0, 1, 2, 3],
the result will have positions [0, 1, 2] (where old position 2 becomes new position 1).
For example, remove_indexes(x, [0, -1]) will return the equivalent of tensor.to_dense().numpy()[1:-1]
Args:
sparse_coo_tensor: A sparse COO tensor of shape (num_positions, num_features)
indexes_to_remove: Iterable of position indexes to remove (supports negative indexing)
Returns:
A new sparse COO tensor with the specified positions removed and indices remapped
"""
if not sparse_coo_tensor.is_sparse or sparse_coo_tensor.layout != torch.sparse_coo:
raise TypeError("sparse_coo_tensor must be a torch.sparse_coo_tensor.")
if sparse_coo_tensor.dim() != 2:
raise ValueError(
f"sparse tensors with more than 2 dimensions are not supported, got {sparse_coo_tensor.dim()} dimensions"
)
indices = sparse_coo_tensor.indices()
values = sparse_coo_tensor.values()
num_positions = sparse_coo_tensor.size(0)
num_features = sparse_coo_tensor.size(1)
# Convert negative indices to positive and create sorted list
indexes_to_remove_list = []
for idx in indexes_to_remove:
if idx < 0:
idx = num_positions + idx
indexes_to_remove_list.append(idx)
indexes_to_remove_set = set(indexes_to_remove_list)
if max(indexes_to_remove_set) > num_positions - 1:
raise ValueError(
f"Index to remove {max(indexes_to_remove_set)} is out of bounds for tensor with size {num_positions}"
)
position_indices = indices[0]
mask = ~torch.isin(
position_indices,
torch.tensor(list(indexes_to_remove_set), device=position_indices.device),
)
filtered_indices = indices[:, mask]
new_values = values[mask]
# Create mapping from old positions to new positions
# new position = old position - count(removed positions < old position)
old_positions = filtered_indices[0]
sorted_removed = sorted(indexes_to_remove_set)
position_mapping = torch.zeros(
num_positions, dtype=torch.long, device=position_indices.device
)
removed_count = 0
removed_idx = 0
for pos in range(num_positions):
while removed_idx < len(sorted_removed) and sorted_removed[removed_idx] < pos:
removed_count += 1
removed_idx += 1
position_mapping[pos] = pos - removed_count
# Apply mapping to position indices
new_position_indices = position_mapping[old_positions]
# Construct new indices with remapped positions
new_indices = torch.stack([new_position_indices, filtered_indices[1]], dim=0)
new_num_positions = num_positions - len(indexes_to_remove_set)
return torch.sparse_coo_tensor(
new_indices, new_values, size=(new_num_positions, num_features)
).coalesce()
def max_pool(sparse_coo_tensor: torch.Tensor, axis: int) -> torch.Tensor:
"""Max pool sparse features along the specified axis.
Args:
sparse_coo_tensor: A sparse COO tensor of shape (num_positions, num_features)
axis: The axis to pool over (0 for positions, 1 for features)
Returns:
Max-pooled tensor.
"""
if not sparse_coo_tensor.is_sparse or sparse_coo_tensor.layout != torch.sparse_coo:
raise TypeError("sparse_coo_tensor must be a torch.sparse_coo_tensor.")
if sparse_coo_tensor.dim() != 2:
raise ValueError(
f"sparse tensors with more than 2 dimensions are not supported, got {sparse_coo_tensor.dim()} dimensions"
)
if axis not in (0, 1):
raise ValueError(f"axis must be 0 or 1, got {axis}")
indices = sparse_coo_tensor.indices()
values = sparse_coo_tensor.values()
if axis == 0:
# Pool over positions (axis 0), return max per feature
output_size = sparse_coo_tensor.size(1)
scatter_indices = indices[1] # feature indices
else: # axis == 1
# Pool over features (axis 1), return max per position
output_size = sparse_coo_tensor.size(0)
scatter_indices = indices[0] # position indices
result = torch.zeros(output_size, dtype=values.dtype, device=values.device)
result.scatter_reduce_(
0, scatter_indices, values, reduce="amax", include_self=False
)
return result
================================================
FILE: cookbook/tutorials/README.md
================================================
# **ESM Tutorial Notebooks**
Tutorial notebooks are the best way to get hands-on with ESM models\! Use the notebooks to explore model capabilities, learn workflows that can be applied to your own data, and learn how to interpret model outputs.
**ESMC**
ESMC is a protein language model that embeds sequences into rich numerical representations. Use it for analyzing, classifying, comparing, and interpreting proteins.
| Notebook | Colab Notebook | Description |
| :---- | :---- | :---- |
| Embedding sequences with ESMC | `embed.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/embed.ipynb) | Embed protein sequences and explore how different transformer layers encode structural and functional information. |
| Zero-shot entropy and mutation analysis | `esmc_mutation_scoring.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esmc_mutation_scoring.ipynb) | Compute per-position entropy and log-likelihood ratios to identify constrained vs. mutation-tolerant sites. |
| Layer sweep for enzyme function classification | `esmc_layer_sweep.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esmc_layer_sweep.ipynb) | Learn how to sweep all layers to find which one is best using enzyme classification as a task. |
| Fine-tuning ESMC | `esmc_finetune.ipynb`<br> [](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esmc_finetune.ipynb) | Fine-tune a classification or regression head for your dataset on top of ESMC using Parameter Efficient Fine-tuning (PEFT) |
## **Interpretable features through Sparse Autoencoders (SAEs)**
| Notebook | Colab Notebook | Description |
| :---- | :---- | :---- |
| Understanding proteins with SAE features |`esmc_sae_feature_interpretation.ipynb`<br> [](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esmc_sae_feature_interpretation.ipynb) |Extract and visualize sparse autoencoder features, rank by peak activation and prevalence, and map activations onto 3D structure. |
## **ESMFold2**
ESMFold2 predicts 3D protein structure from sequence, including DNA/RNA and small molecules.
| Notebook | Colab Notebook | Description |
| :---- | :---- | :---- |
| Folding with ESMFold2 | `esmfold2.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esmfold2.ipynb) | Fold proteins in combination with DNA, RNA and small-molecule ligands. |
| Binder design | `binder_design.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/binder_design.ipynb) | Design antibodies and minibinders with high hit rates. Implements the protocol featured in our paper, which produced binders exhibiting nanomolar affinity, target specificity, and functional activity in laboratory assays. |
## **ESM3**
ESM3 is a generative model that reasons jointly over protein sequence, structure, and function. Use it for designing new proteins or editing existing ones.
| Notebook | Colab Notebook | Description |
| :---- | :---- | :---- |
| Understanding the ESMProtein class | `esmprotein.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esmprotein.ipynb) | Get familiar with how ESM3 represents proteins. |
| Generating proteins with ESM3 | `esm3_generate.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esm3_generate.ipynb) | Learn how to scaffold a functional motif, edit secondary structure, and guide design using solvent exposure. |
| Designing a novel GFP with ESM3 | `gfp_design.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/gfp_design.ipynb) | Walk through the exact prompting strategy used to design a novel fluorescent protein with no close natural relatives. |
| Guided generation with ESM3 | `esm3_guided_generation.ipynb`<br>[](https://colab.research.google.com/github/biohub/esm/blob/main/cookbook/tutorials/esm3_guided_generation.ipynb) | Add scoring functions into the generation process, such as structural quality, sequence constraints, or other properties. |
================================================
FILE: cookbook/tutorials/binder_design.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"id": "b5b44288",
"metadata": {},
"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://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",
"\n",
"**You'll need:**\n",
"- A target protein sequence (or pick one of the built-in presets)\n",
"- A [Modal](https://modal.com/) account and token (free tier works to get started)\n",
"- The notebook contains cells launching independent trajectories and a small sweep (N=256). With Modal H100 pricing, this will cost ≈$2 and ≈$150, respectively\n",
"\n",
"**You'll get:** a file of top-ranked designed binder sequences, plus 3D structures of the predicted complexes that you can view in the notebook.\n",
"\n",
"**Workflow:**\n",
"1. **Setup** (one-time): install dependencies, get a Modal token, deploy the design app\n",
"2. **Try one job**: pick a target and binder type, run a single design end-to-end as a sanity check\n",
"3. **Run a sweep**: launch many parallel jobs to produce real candidates\n",
"4. **Pick the designs to order**: filter and rank, save the shortlist\n",
"\n",
"\n",
"> **Why does this notebook need Modal?** Each design job needs a GPU and a few minutes of compute, and a real campaign means launching hundreds of these in parallel. [Modal](https://modal.com/) is a cloud service that lets this notebook spin up remote GPUs on demand, run each job on one, and return results to you. You don't manage any servers or containers yourself. When you \"deploy\" the design app to Modal (in section 1), you're uploading a Python file that defines the job; Modal then runs that code for you whenever the notebook calls `app.design.spawn(...)`."
]
},
{
"cell_type": "markdown",
"id": "4421cdaa",
"metadata": {},
"source": [
"### 1. Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "99783792",
"metadata": {},
"outputs": [],
"source": [
"# Environment\n",
"! pip install esm\n",
"! pip install modal py3dmol pyarrow"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d3de7d00",
"metadata": {},
"outputs": [],
"source": [
"# Confirm you have a modal token, or make one\n",
"! modal token info # Check\n",
"# ! modal token new # Create"
]
},
{
"cell_type": "markdown",
"id": "c88594a9",
"metadata": {},
"source": [
"### Deploy the design app to Modal\n",
"\n",
"The file `binder_design.py` lives in the same folder as this notebook. It defines the GPU job, the model, and the design loop.\n",
"\n",
"You only need to deploy once. Re-run this cell only if the underlying `.py` file changes."
]
},
{
"cell_type": "markdown",
"id": "8e88a943",
"metadata": {},
"source": [
"### If you're running on Colab\n",
"\n",
"Colab only pulls the notebook itself when you open it, not the surrounding files from the repo. Run the cell below to grab `binder_design.py` into your Colab workspace. If you're running locally, skip it."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "af900bae",
"metadata": {},
"outputs": [],
"source": [
"# Colab only: download binder_design.py into the working directory\n",
"! wget -q https://raw.githubusercontent.com/Biohub/esm/main/cookbook/tutorials/binder_design.py\n",
"! ls binder_design.py "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2ee89767",
"metadata": {},
"outputs": [],
"source": [
"# Deploy (or redeploy after changing binder_design.py).\n",
"# This only needs to be run a single time, unless code in binder_design.py changes.\n",
"! modal deploy binder_design.py"
]
},
{
"cell_type": "markdown",
"id": "dc8456da",
"metadata": {},
"source": [
"### Imports"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "37c03b59",
"metadata": {},
"outputs": [],
"source": [
"from itertools import product\n",
"from pathlib import Path\n",
"\n",
"import modal\n",
"import pandas as pd\n",
"import py3Dmol\n",
"from Bio.SeqUtils.ProtParam import ProteinAnalysis\n",
"from tqdm.auto import tqdm"
]
},
{
"cell_type": "markdown",
"id": "1fe9a141",
"metadata": {},
"source": [
"### App setup\n",
"\n",
"`modal.Cls.from_name(...)` grabs a handle to the design app you just deployed, without rerunning anything on Modal yet. Instantiating it gives you `app`, which is what you'll call `app.design.spawn(...)` on to launch design jobs.\n",
"\n",
"`use_scaling_critics=True` is the default. Setting it to `False` skips the additional critic models from the paper, reducing compute per job at the cost of selection quality."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "291d8bfe",
"metadata": {},
"outputs": [],
"source": [
"ESMFold2Design = modal.Cls.from_name(\"esmfold2-design\", \"ESMFold2DesignModal\")\n",
"# Set 'use_scaling_critics' to evaluate with the additional critics.\n",
"# On by default.\n",
"app = ESMFold2Design(use_scaling_critics=True)"
]
},
{
"cell_type": "markdown",
"id": "159b63df",
"metadata": {},
"source": [
"## 2. Try one design job\n",
"\n",
"Run a single job end-to-end before launching a sweep. This is a sanity check that everything is wired up and that the target/scaffold combo you've chosen produces a sensible complex.\n",
"\n",
"Pick **one** of the two options below and run only that cell. (They both define a variable called `future`, so running both back-to-back overwrites the first one.)\n",
"\n",
"**Option 1** uses a built-in target and binder scaffold. Available targets: `ctla4`, `egfr`, `pdgfrb`, `pd-l1`, `cd45`. Available binder types: `minibinder`, `trastuzumab_framework_vhvl` (an antibody scaffold). Easiest if your target is one of the built-ins.\n",
"\n",
"**Option 2** takes your own target sequence and binder scaffold. Pass a display `target_name` and `binder_name` (these do not need to be preset keys) along with the sequences. In the binder scaffold, `#` means \"design this position\" and any amino acid letter means \"keep this position fixed.\" For example:\n",
"- `\"#\" * 60` designs a fully free 60-residue minibinder\n",
"- A trastuzumab-style antibody scaffold (shown below) fixes the framework regions and lets the model design the CDR loops \n",
"\n",
"If you're designing an antibody, pass `is_antibody=True` so the selection step later uses the antibody-appropriate scoring.\n",
"\n",
"Jobs run on Modal in the background. The `dashboard_url` link the cell prints is a clickable link to live progress."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "826c88d1",
"metadata": {},
"outputs": [],
"source": [
"# ---- Option 1: Use presets. ----\n",
"# Relies on the registry in binder_design.py::{TARGET_SEQUENCES,BINDER_PROMPT_FACTORIES}, which can be modified.\n",
"future = app.design.spawn(target_name=\"ctla4\", binder_name=\"minibinder\")\n",
"future.get_dashboard_url() # A clickable link to Modal dashboard"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c1bda1a6",
"metadata": {},
"outputs": [],
"source": [
"# ---- Option 2: Provide your own sequences. ----\n",
"# Our pd-l1 sequence crop.\n",
"pdl1_sequence = \"AFTVTVPKDLYVVEYGSNMTIECKFPVEKQLDLAALIVYWEMEDKNIIQFVHGEEDLKVQHSSYRQRARLLKDQLSLGNAALQITDVKLQDAGVYRCMISYGGADYKRITVKVNA\"\n",
"# A sample of 'trastuzumab_framework_vhvl' template. From binder_design.py::BINDER_PROMPT_FACTORIES.\n",
"trastuzumab_framework_vhvl = \"EVQLVESGGGLVQPGGSLRLSCAAS#######YIHWVRQAPGKGLEWVARI#####TRYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCSR###########WGQGTLVTVSSGGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC###########WYQQKPGKAPKLLIY#######GVPSRFSGSRSGTDFTLTISSLQPEDFATYYC#########FGQGTKVEIK\"\n",
"future2 = app.design.spawn(\n",
" target_name=\"pd-l1-custom\",\n",
" target_sequence=pdl1_sequence,\n",
" binder_name=\"trastuzumab-custom\",\n",
" binder_sequence=trastuzumab_framework_vhvl,\n",
" is_antibody=True,\n",
")\n",
"future2.get_dashboard_url() # A clickable link to Modal dashboard"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e442e64e",
"metadata": {},
"outputs": [],
"source": [
"# ---- Monitor ----\n",
"# Tail a function's output here in jupyter. Interrupt the kernel to stop the tail.\n",
"! modal app logs esmfold2-design -f --function-call {future.object_id}"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "79ea37f3",
"metadata": {},
"outputs": [],
"source": [
"# ---- Load result ----\n",
"best_sequences, trajectory, critic_results = future.get()\n",
"# best_sequences, trajectory, critic_results = future2.get()\n",
"print(\"Best sequences: \", best_sequences)\n",
"df = pd.DataFrame(critic_results)\n",
"df.drop(columns=[\"logits\", \"complex\"])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d80597fa",
"metadata": {},
"outputs": [],
"source": [
"# ---- Visualize ----\n",
"protein_complex = (\n",
" df[df.critic_name.eq(\"ESMFold2-Experimental-Cutoff2025\")].iloc[0].complex\n",
")\n",
"(\n",
" py3Dmol.view(width=600, height=600)\n",
" .addModel(protein_complex.to_pdb_string(), \"pdb\")\n",
" .setStyle({\"chain\": \"A\"}, {\"cartoon\": {\"color\": \"green\"}})\n",
" .setStyle(\n",
" {\"chain\": \"B\"},\n",
" {\n",
" \"cartoon\": {\n",
" \"colorscheme\": {\"prop\": \"b\", \"gradient\": \"rwb\", \"min\": 60, \"max\": 100}\n",
" }\n",
" },\n",
" )\n",
" .addStyle( # B factor coloring for binder\n",
" {\"and\": [{\"chain\": \"B\"}, {\"not\": {\"atom\": [\"N\", \"C\", \"O\"]}}]},\n",
" {\n",
" \"stick\": {\n",
" \"colorscheme\": {\"prop\": \"b\", \"gradient\": \"rwb\", \"min\": 60, \"max\": 100},\n",
" \"radius\": 0.2,\n",
" }\n",
" },\n",
" )\n",
" .addStyle( # Target colored green\n",
" {\"and\": [{\"chain\": \"A\"}, {\"not\": {\"atom\": [\"N\", \"C\", \"O\"]}}]},\n",
" {\"stick\": {\"color\": \"green\", \"radius\": 0.2}},\n",
" )\n",
" .center()\n",
" .zoomTo()\n",
")"
]
},
{
"cell_type": "markdown",
"id": "fc105292",
"metadata": {},
"source": [
"## 3. Run a sweep for real designs\n",
"For real candidates worth ordering, sweep across many seeds (and optionally multiple targets, binder types, or lengths) and select the best.\n",
"\n",
"Edit `targets`, `binders`, and the other sweep axes below. Each target/binder is a `(name, sequence)` tuple: pass `None` for sequence to use a built-in preset (name is the preset key); otherwise name is your display name and sequence is the amino-acid string. The default below sweeps 128 seeds across two binder types against PD-L1. \n",
"\n",
"**Before you click Run on the Launch cell, check the printed shape of the dataframe and confirm it's the number of jobs you intended.** "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac02bbaa",
"metadata": {},
"outputs": [],
"source": [
"# ---- Config ----\n",
"save_dir = Path(\"sweep\")\n",
"save_dir.mkdir(exist_ok=True)\n",
"\n",
"targets = [(\"pd-l1\", None)]\n",
"binders = [(\"minibinder\", None), (\"trastuzumab_framework_vhvl\", None)]\n",
"\n",
"line_sweeps = dict(\n",
" target=targets,\n",
" binder=binders,\n",
" use_scaling_critics=[True],\n",
" seed=list(range(16)),\n",
" # NOTE - reduce if you want lower latency to get results.\n",
" batch_size=[6],\n",
")\n",
"df = pd.DataFrame(product(*line_sweeps.values()), columns=line_sweeps.keys())\n",
"df[\"target_name\"], df[\"target_sequence\"] = zip(*df[\"target\"], strict=True)\n",
"df[\"binder_name\"], df[\"binder_sequence\"] = zip(*df[\"binder\"], strict=True)\n",
"df = df.drop(columns=[\"target\", \"binder\"])\n",
"display(df.head(2))\n",
"df.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9b768813",
"metadata": {},
"outputs": [],
"source": [
"# ---- Launch ----\n",
"df[\"call_id\"] = [\n",
" app.design.spawn(\n",
" target_name=row.target_name,\n",
" target_sequence=row.target_sequence,\n",
" binder_name=row.binder_name,\n",
" binder_sequence=row.binder_sequence,\n",
" seed=row.seed,\n",
" batch_size=row.batch_size,\n",
" ).object_id\n",
" for row in df.itertuples()\n",
"]\n",
"df.to_parquet(save_dir / \"manifest.parquet\", index=False)\n",
"print(\n",
" f\"Spawned {len(df)} jobs. It is safe to close the notebook.\"\n",
" \"The next cell will resume from call_id's, saved by Modal for up to 7 days.\"\n",
")"
]
},
{
"cell_type": "markdown",
"id": "13f0542e",
"metadata": {},
"source": [
"### Coming back later (important)\n",
"\n",
"Your jobs are now running on Modal's GPUs. **You do not need to keep this notebook open.** Modal continues running the jobs on its own and holds the results for up to 7 days.\n",
"\n",
"**If you're running on Colab:** the runtime is wiped when you disconnect. To resume a sweep on Colab, mount Google Drive **before** running the Launch cell and point `save_dir` at a Drive path, e.g.:\n",
"\n",
"\\`\\`\\`python\n",
"from google.colab import drive\n",
"drive.mount('/content/drive')\n",
"save_dir = Path('/content/drive/MyDrive/binder_sweep')\n",
"\\`\\`\\`\n",
"\n",
"When you reopen the notebook later, you'll need to re-install the dependencies, re-authenticate Modal (`modal token new`), re-mount Drive, then resume from the Monitor cell."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b9a637f0",
"metadata": {},
"outputs": [],
"source": [
"# ---- Monitor ----\n",
"from tqdm.contrib.concurrent import thread_map\n",
"\n",
"df = pd.read_parquet(save_dir / \"manifest.parquet\")\n",
"df[\"future\"] = thread_map(modal.FunctionCall.from_id, df.call_id.values)\n",
"df[\"status\"] = thread_map(lambda f: f.get_call_graph()[0].status.name, df.future.values)\n",
"print(\"First task url: \", df.at[0, \"future\"].get_dashboard_url()) # pyright: ignore\n",
"df.status.value_counts()"
]
},
{
"cell_type": "markdown",
"id": "32fbc8d8",
"metadata": {},
"source": [
"The Collect cell waits for all jobs that succeeded and unpacks their results. Jobs that failed or are still running are skipped, so you can re-run Monitor + Collect periodically as more jobs finish."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "973f7d6f",
"metadata": {},
"outputs": [],
"source": [
"# ---- Collect ----\n",
"df_success = df[df.status.eq(\"SUCCESS\")].copy()\n",
"df_success[\"result\"] = thread_map(\n",
" lambda x: x.get(), df_success.future.values, max_workers=32, chunksize=32\n",
") # Blocks until all jobs are complete.\n",
"df_success[\"result_df\"] = [pd.DataFrame(r[2]) for r in tqdm(df_success.result)] # pyright: ignore"
]
},
{
"cell_type": "markdown",
"id": "d53ad844",
"metadata": {},
"source": [
"## 4. Pick the designs to order\n",
"\n",
"This is your final shortlist. The cell below:\n",
"\n",
"1. Combines results from all successful jobs into one dataframe.\n",
"2. Filters minibinders to isoelectric point under 6 (helps with solubility and expression). Antibodies pass through unfiltered.\n",
"3. Scores each unique designed sequence by averaging `iptm` (interface predicted TM-score, higher is better) and an `iptm_proxy` term across its trajectories.\n",
"4. Returns the top 84 designs per (target, binder type), saved to `selection.parquet` inside your `save_dir`.\n",
"5. Writes one PDB per selected design to `selected_structures/` inside your `save_dir`.\n",
"\n",
"84 is a plate-friendly number for ordering and screening. The cutoff and the isoelectric-point filter are currently hardcoded inside the cell, so to change them, edit the values directly in the function.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "39fd979c",
"metadata": {},
"outputs": [],
"source": [
"# ---- Select ----\n",
"\n",
"# Join all result_df's, with other fields in df broadcasted as metadata.\n",
"df_result = pd.concat(\n",
" [\n",
" row.result_df.assign(**row.drop([\"result\", \"result_df\"]).to_dict()) # pyright: ignore\n",
" for _, row in df_success.iterrows()\n",
" ],\n",
" ignore_index=True,\n",
" axis=0,\n",
")\n",
"\n",
"# Filter minibinder designs with isoelectric point >= 6.\n",
"df_result[\"binder_sequence\"] = df_result.designed_sequence.str.split(r\"\\|\").str[1]\n",
"df_result[\"isoelectric_point\"] = [\n",
" ProteinAnalysis(seq).isoelectric_point()\n",
" for seq in tqdm(df_result.binder_sequence.values)\n",
"]\n",
"# Isoelectric point filter\n",
"df_filter = df_result[df_result.is_antibody | df_result.isoelectric_point.lt(6)]\n",
"\n",
"\n",
"# Select the top 84 designs from each (target, binder) combination\n",
"SCALING_CHECKPOINT_SUBSTRING = \"ESMFold2-Experimental-Fast-base\"\n",
"\n",
"\n",
"def select(df: pd.DataFrame) -> pd.DataFrame:\n",
" df = df.copy()\n",
" is_scaling = df.critic_name.str.contains(\n",
" SCALING_CHECKPOINT_SUBSTRING, regex=False, na=False\n",
" )\n",
" iptm_proxy = df.distogram_iptm_proxy.where(\n",
" ~df.is_antibody, df.cdr_distogram_iptm_proxy\n",
" )\n",
"\n",
" df[\"iptm_score\"] = df.iptm.where(~is_scaling)\n",
" df[\"iptm_proxy_score\"] = iptm_proxy.where(is_scaling)\n",
" scores = df.groupby(\"designed_sequence\", as_index=False).agg(\n",
" iptm_score=(\"iptm_score\", \"mean\"), iptm_proxy_score=(\"iptm_proxy_score\", \"mean\")\n",
" )\n",
" scores[\"selection_score\"] = 0.5 * scores.iptm_score.fillna(\n",
" 0\n",
" ) + 0.5 * scores.iptm_proxy_score.fillna(0)\n",
" return scores.nlargest(min(len(scores), 84), \"selection_score\")\n",
"\n",
"\n",
"df_select = df_filter.groupby([\"target_name\", \"binder_name\"]).apply(\n",
" select, include_groups=False\n",
")\n",
"df_select.to_parquet(save_dir / \"selection.parquet\", index=False)\n",
"df_select"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5a131c33",
"metadata": {},
"outputs": [],
"source": [
"# ---- Write selected structures ----\n",
"complexes = (\n",
" df_result[df_result.critic_name.eq(\"ESMFold2-Experimental-Cutoff2025\")]\n",
" .drop_duplicates(\"designed_sequence\")\n",
" .set_index(\"designed_sequence\")[\"complex\"]\n",
")\n",
"pdb_dir = save_dir / \"selected_structures\"\n",
"pdb_dir.mkdir(exist_ok=True)\n",
"selected = df_select.reset_index().sort_values(\n",
" [\"target_name\", \"binder_name\", \"selection_score\"], ascending=[True, True, False]\n",
")\n",
"for rank, row in enumerate(selected.itertuples(), start=1):\n",
" pdb_path = (\n",
" pdb_dir\n",
" / f\"{rank:04d}_{row.target_name}_{row.binder_name}_score{row.selection_score:.3f}.pdb\"\n",
" )\n",
" pdb_path.write_text(complexes[row.designed_sequence].to_pdb_string())\n",
"print(f\"Wrote {len(selected)} PDBs to {pdb_dir.resolve()}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b54a71de",
"metadata": {},
"outputs": [],
"source": [
"df_result[df_result.critic_name.eq(\"ESMFold2-Experimental-Cutoff2025\")].drop(\n",
" columns=[\"complex\", \"logits\"]\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "851365c0",
"metadata": {},
"outputs": [],
"source": [
"df_select"
]
},
{
"cell_type": "markdown",
"id": "7e174c1e",
"metadata": {},
"source": [
"## Appendix"
]
},
{
"cell_type": "markdown",
"id": "3f79528d",
"metadata": {},
"source": [
"### Modal Primer"
]
},
{
"cell_type": "markdown",
"id": "4297bc4d",
"metadata": {},
"source": [
"- **info: ephemeral vs deployment** \n",
" Ephemeral = temporary app from `modal run` or `app.run()`, stopped when the client exits. Deployment = persistent named app from `modal deploy`, reused and observable across runs. ([modal.com](https://modal.com/docs/guide/apps?utm_source=openai))\n",
"\n",
"- **info: dashboard** \n",
" Generic dashboard/apps page: [https://modal.com/apps](https://modal.com/apps). Modal also prints app/deployment links during runs/deploys. ([modal.com](https://modal.com/docs/guide/apps?utm_source=openai))\n",
"\n",
"- **cli: ephemeral run** \n",
" ```bash\n",
" modal run path/to/app.py\n",
" ```\n",
"\n",
"- **cli: deploy/redeploy** \n",
" ```bash\n",
" modal deploy path/to/app.py\n",
" ```\n",
" Running this on an existing app name redeploys a new version. ([modal.com](https://modal.com/docs/reference/cli/deploy?utm_source=openai))\n",
"\n",
"- **local: ephemeral from Python** \n",
" ```python\n",
" with modal.enable_output():\n",
" with modal_app.run():\n",
" result = local_modal_obj.remote(...)\n",
" ```\n",
"\n",
"- **local: call a deployment** \n",
" ```python\n",
" Cls = modal.Cls.from_name(\"app-name\", \"ClassName\")\n",
" obj = Cls(...)\n",
" result = obj.method.remote(...)\n",
" ```\n",
" `Cls.from_name` references a class from a deployed app lazily. ([modal.com](https://modal.com/docs/reference/modal.Cls?utm_source=openai))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.12.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
================================================
FILE: cookbook/tutorials/binder_design.py
================================================
# /// script
# requires-python = "<=3.13"
# dependencies = [
# "abnumber",
# "esm@git+https://github.com/Biohub/esm.git@main",
# "modal",
# ]
# ///
"""
Code for binder design with ESMFold2 and ESMC.
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
import math
import os
import random
import string
import time
from dataclasses import dataclass
from functools import cache
from typing import Any
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
import biotite.structure
import modal
import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim
from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer
from esm.models.esmfold2 import (
ELEMENT_NUMBER_TO_SYMBOL,
ProteinInput,
StructurePredictionInput,
load_ccd,
prepare_esmfold2_input,
)
from esm.models.esmfold2.constants import (
MOL_TYPE_NONPOLYMER,
PROTEIN_1TO3,
PROTEIN_3TO1,
RES_TYPE_TO_CCD,
)
from esm.models.esmfold2.experimental import EsmFold2ExperimentalModel
from esm.models.esmfold2.experimental import MSAEncoder as EsmFold2MSAEncoder
from esm.models.esmfold2.layers import (
BACKEND_CUEQ,
BACKEND_FUSED,
CUE_AVAILABLE,
TRITON_KERNELS_AVAILABLE,
PairUpdateBlock,
)
from esm.models.esmfold2.layers import _seed_context as seed_context
from esm.utils.structure.mmcif_parsing import PLDDT_B_FACTOR_SCALE
from esm.utils.structure.protein_chain import ProteinChain
from esm.utils.structure.protein_complex import ProteinComplex
os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
TrajectoryStep = dict[str, torch.Tensor | float]
Trajectory = dict[int, TrajectoryStep]
# ---- Constants ----
# General
TOKENS = ["<pad>", "-"] + [RES_TYPE_TO_CCD[i] for i in range(2, 33)]
ELEMENTS = ["X"] * (max(ELEMENT_NUMBER_TO_SYMBOL) + 1)
ELEMENTS[0] = "<pad>"
for _atomic_num, _symbol in ELEMENT_NUMBER_TO_SYMBOL.items():
ELEMENTS[_atomic_num] = _symbol[:1] + _symbol[1:].lower()
TOKEN_IDS = {token: idx for idx, token in enumerate(TOKENS)}
AA_DIMS = 20
# Cysteine index in the 20-dim AA space (TOKEN_IDS are offset by 2 for <pad> and -)
CYS_IDX = TOKEN_IDS[PROTEIN_1TO3["C"]] - 2
MUTABLE_TOKEN = "#"
# Contains AA chars at fixed positions and MUTABLE_TOKEN at mutable positions.
BinderPromptStr = str
# Design
LOSS_WEIGHTS = {"intra_contact": 0.5, "inter_contact": 0.5, "glob": 0.2, "epitope": 0.5}
STEPS = 150
LOG_INTERVAL = 5
LEARNING_RATE = 0.1
TEMPERATURE_MIN = 1e-2
ESMC_MASK_FRACTION = 0.15
LM_LOSS_BATCH_SIZE = 128
LM_MASK_PASSES = 4
BYTES_PER_GIB = 1024**3
COMPILE = True
# NOTE - This significantly reduces VRAM usage.
# On config (target_name="cd45", binder_name="trastuzumab_framework_vhvl", batch_size=1)
# this reduces VRAM from 51GB -> 27GB. And enables increasing batch size up to 6.
# We are testing this setting in silico, and may change the default to True, in the future.
REUSE_ESMC = True
# ---- Prompts ----
@dataclass(frozen=True)
class PromptFactory:
"""A simple factory for making binder prompt strings."""
name: str
template: str # string with format fields
length_ranges: dict[str, tuple[int, int]] # map from field name tp length range
is_antibody: bool # Used to set LM loss weight for antibodies.
def sample(self, seed: int) -> BinderPromptStr:
random.seed(seed)
return self.template.format(
**{
key: MUTABLE_TOKEN * random.randint(low, high)
for key, (low, high) in self.length_ranges.items()
}
)
# fmt: off
BINDER_PROMPT_FACTORIES = {
"minibinder": PromptFactory(name="minibinder", template="{seq}", length_ranges={"seq": (60, 200)}, is_antibody=False),
"trastuzumab_framework_vhvl": PromptFactory(
name="trastuzumab_framework_vhvl",
template="EVQLVESGGGLVQPGGSLRLSCAAS{hcdr1}YIHWVRQAPGKGLEWVARI{hcdr2}TRYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCSR{hcdr3}WGQGTLVTVSSGGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY{lcdr2}GVPSRFSGSRSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK",
length_ranges = {"hcdr1": (7, 9), "hcdr2": (5, 6), "hcdr3": (9, 15), "lcdr1": (11, 16), "lcdr2": (7, 7), "lcdr3": (9, 9)},
is_antibody=True,
),
"atezolizumab_framework_vhvl": PromptFactory(
name="atezolizumab_framework_vhvl",
template="EVQLVESGGGLVQPGGSLRLSCAAS{hcdr1}WIHWVRQAPGKGLEWVAWI{hcdr2}TYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCAR{hcdr3}WGQGTLVTVSSGGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY{lcdr2}GVPSRFSGSGSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK",
length_ranges = {"hcdr1": (7, 9), "hcdr2": (5, 6), "hcdr3": (9, 15), "lcdr1": (11, 16), "lcdr2": (7, 7), "lcdr3": (9, 9)},
is_antibody=True,
),
"ocankitug_framework_vhvl": PromptFactory(
name="ocankitug_framework_vhvl",
template="QVQLVQSGAEVKKPGSSVKVSCKAS{hcdr1}WMHWVRQAPGQGLEWMGII{hcdr2}TSLNQKFQGRVTITADTSTSTAYMELSSLRSEDTAVYYCAR{hcdr3}WGQGTLVTVSSGGGSGGGSGGGSGGGSDIQMTQSPSSLSASVGDRVTITC{lcdr1}WYQQKPGKAPKLLIY{lcdr2}GVPSRFSGSGSGTDFTLTISSLQPEDFATYYC{lcdr3}FGQGTKVEIK",
length_ranges = {"hcdr1": (7, 9), "hcdr2": (5, 6), "hcdr3": (8, 14), "lcdr1": (11, 16), "lcdr2": (7, 7), "lcdr3": (9, 9)},
is_antibody=True,
)
}
TARGET_SEQUENCES = {
# https://www.uniprot.org/uniprotkb/P08575 389-574
"cd45": "GSPGEPQIIFCRSEAAHQGVITWNPPQRSFHNFTLCYIKETEKDCLNLDKNLIKYDLQNLKPYTKYVLSLHAYIIAKVQRNGSAAMCHFTTKSAPPSQVWNMTVSMTSDNSMHVKCRPPRDRNGPHERYHLEVEAGNTLVRNESHKNCDFRVKDLQYSTDYTFKAYFHNGDYPGEPFILHHSTSY",
# https://www.uniprot.org/uniprotkb/P16410 37-155
"ctla4": "MHVAQPAVVLASSRGIASFVCEYASPGKATEVRVTVLRQADSQVTEVCAATYMMGNELTFLDDSICTGTSSGNQVNLTIQGLRAMDTGLYICKVELMYPPPYYLGIGNGTQIYVIDPE",
# https://www.uniprot.org/uniprotkb/P00533 333-524
"egfr": "RKVCNGIGIGEFKDSLSINATNIKHFKNCTSISGDLHILPVAFRGDSFTHTPPLDPQELDILKTVKEITGFLLIQAWPENRTDLHAFENLEIIRGRTKQHGQFSLAVVSLNITSLGLRSLKEISDGDVIISGNKNLCYANTINWKKLFGTSGQKTKIISNRGENSCKATGQVCHALCSPEGCWGPEPRDCV",
# https://www.uniprot.org/uniprotkb/Q9NZQ7 17-132
"pd-l1": "AFTVTVPKDLYVVEYGSNMTIECKFPVEKQLDLAALIVYWEMEDKNIIQFVHGEEDLKVQHSSYRQRARLLKDQLSLGNAALQITDVKLQDAGVYRCMISYGGADYKRITVKVNA",
# https://www.uniprot.org/uniprotkb/P09619 125-312
"pdgfr": "GFLPNDAEELFIFLTEITEITIPCRVTDPQLVVTLHEKKGDVALPVPYDHQRGFSGIFEDRSYICKTTIGDREVDSDAYYVYRLQVSSINVSVNAVQTVVRQGENITLMCIVIGNEVVNFEWTYPRKESGRLVEPVTDFLLDMPYHIRSILHIPSAELEDSGTYTCNVTESVNDHQDEKAINITVVE",
}
# fmt: on
# ---- Helper functions ----
def build_initial_soft_sequence_logits(sequence: str, batch_size: int) -> torch.Tensor:
"""
Initialize logits with:
- High confidence (10.0) for fixed positions
- Random (~0) for mutable positions
- -1e6 for cysteines
"""
if all(aa == MUTABLE_TOKEN for aa in sequence):
logits = 0.01 * torch.randn([batch_size, len(sequence), AA_DIMS])
logits[:, :, CYS_IDX] = -1e6 # remove cysteines
else:
logits = torch.zeros([batch_size, len(sequence), AA_DIMS])
for i, aa in enumerate(sequence):
if aa == MUTABLE_TOKEN: # mutable position - random
logits[:, i, :] = 0.01 * torch.randn(batch_size, AA_DIMS)
logits[:, i, CYS_IDX] = -1e6
else: # fixed position
assert aa in PROTEIN_1TO3, aa
token_id = TOKEN_IDS[PROTEIN_1TO3[aa]]
logits[:, i, token_id - 2] = 10.0
return logits.requires_grad_(True)
def build_gradient_mask(sequence: str, batch_size: int) -> torch.Tensor:
"""
Build gradient mask [B, L, V]:
- 0 for fixed (all amino acids)
- 0 for cysteine at all positions
- 1 for non-cysteine amino acids at mutable positions
"""
mask = torch.ones([batch_size, len(sequence), AA_DIMS])
fixed_positions = [i for i, aa in enumerate(sequence) if aa != MUTABLE_TOKEN]
mask[:, fixed_positions, :] = 0.0
mask[:, :, CYS_IDX] = 0.0
return mask
def sequence_to_one_hot(sequence: str, device="cuda") -> torch.Tensor:
"""Convert target string to one-hot tensor [1, L_target, num_tokens]."""
const_dict = {token: i for i, token in enumerate(TOKENS)}
target_index = [const_dict[PROTEIN_1TO3[letter]] for letter in sequence]
one_hot = F.one_hot(torch.tensor(target_index), num_classes=len(TOKENS))
return one_hot.to(device).unsqueeze(0).float()
def get_mid_points() -> torch.Tensor:
"""128 distance bin midpoints (2p-52 Angstrom range)."""
boundaries = torch.linspace(2, 52.0, 127)
lower = torch.tensor([1.0])
upper = torch.tensor([52.0 + 5.0])
exp_boundaries = torch.cat((lower, boundaries, upper))
return (exp_boundaries[:-1] + exp_boundaries[1:]) / 2
def binned_entropy(
dgram: torch.Tensor, bin_distance: torch.Tensor, cutoff: float
) -> torch.Tensor:
"""Entropy of distance distribution within cutoff (design losses only)."""
bin_mask = ~(bin_distance < cutoff)
masked_dgram = dgram - (1e7 * bin_mask)
px = torch.softmax(masked_dgram, dim=-1)
log_px = torch.log_softmax(dgram, dim=-1)
return -(px * log_px).sum(-1)
def masked_min_k(x: torch.Tensor, mask: torch.Tensor, k: int) -> torch.Tensor:
"""Mean of the smallest k values in x under mask along the last dimension."""
mask = mask.bool()
y = torch.sort(torch.where(mask, x, float("nan")))[0]
k_mask = (torch.arange(y.shape[-1]).to(y.device) < k) & (~torch.isnan(y))
return torch.where(k_mask, y, 0).sum(-1) / (k_mask.sum(-1) + 1e-8)
def masked_average(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
"""Masked mean along last axis."""
mask = mask.bool()
return torch.where(mask, x, 0).sum(-1) / (torch.where(mask, 1, 0).sum(-1) + 1e-8)
# ---- Loss functions ----
def compute_contact_loss(
distogram_logits: torch.Tensor,
bin_distance: torch.Tensor,
num_contacts: int,
min_sep: int,
cutoff: float,
chain_mask: torch.Tensor,
binder_mask: torch.Tensor,
) -> torch.Tensor:
"""Algorithm 12 Contact Losses.
Entropy-based contact loss with sequence separation constraint."""
con_loss = binned_entropy(distogram_logits, bin_distance, cutoff)
position = torch.arange(distogram_logits.shape[1])
p_dist = position[:, None] - position[None, :]
if min_sep > 0:
separation_mask = (torch.abs(p_dist) >= min_sep).to(distogram_logits.device)
binder_mask = torch.logical_and(separation_mask, binder_mask)
per_residue = masked_min_k(con_loss, mask=binder_mask, k=num_contacts).to(
distogram_logits.device
)
return masked_average(per_residue, mask=chain_mask).to(distogram_logits.device)
def compute_intra_contact_loss(
distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor
) -> torch.Tensor:
"""Binder internal contacts (k=2, min_sep=9, cutoff=14A)."""
full_len = distogram_logits.shape[1]
is_binder = torch.ones(full_len, device=distogram_logits.device)
is_binder[:-binder_length] *= 0.0
return compute_contact_loss(
distogram_logits,
bin_distance,
num_contacts=2,
min_sep=9,
cutoff=14.0,
chain_mask=is_binder,
binder_mask=is_binder,
)
def compute_inter_contact_loss(
distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor
) -> torch.Tensor:
"""Binder-target interface (k=1, min_sep=0, cutoff=22A)."""
full_len = distogram_logits.shape[1]
is_binder = torch.ones(full_len, device=distogram_logits.device)
is_binder[:-binder_length] *= 0.0
return compute_contact_loss(
distogram_logits,
bin_distance,
num_contacts=1,
min_sep=0,
cutoff=22.0,
chain_mask=1 - is_binder,
binder_mask=is_binder,
)
def compute_globularity_loss(
distogram_logits: torch.Tensor, binder_length: int, bin_distance: torch.Tensor
) -> torch.Tensor:
"""Algorithm 13 Globularity Loss.
Radius of gyration vs theoretical packed protein."""
binder_disto = distogram_logits[:, -binder_length:, -binder_length:, :]
n = binder_disto.shape[1]
disto_probs = torch.softmax(binder_disto, dim=-1)
bin_distance = bin_distance.clamp(max=27)
e_sq_dist = torch.sum(disto_probs * torch.square(bin_distance), dim=-1)
sum_sq_dist = torch.sum(torch.tril(e_sq_dist, diagonal=-1), dim=(1, 2))
rg_term = torch.sqrt(sum_sq_dist / (n * n))
rg_th = 2.38 * (n**0.365)
return F.elu(rg_term - rg_th)
def get_epitope_mask(
target_sequence: str,
binder_length: int,
target_hotspot_ids: list[str],
batch_size: int,
device: torch.device,
) -> torch.Tensor:
"""Full-complex mask for sequence-local target hotspot residue IDs."""
target_length = len(target_sequence)
epitope_mask = torch.zeros(
target_length + binder_length, dtype=torch.bool, device=device
)
for hotspot_id in target_hotspot_ids:
assert len(hotspot_id) >= 2, (
f"Hotspot residue ID {hotspot_id!r} must have form <residue><idx>, "
"for example L150."
)
residue = hotspot_id[0]
hotspot_idx_text = hotspot_id[1:]
assert residue in PROTEIN_1TO3, (
f"Hotspot residue ID {hotspot_id!r} starts with unknown residue "
f"{residue!r}."
)
assert hotspot_idx_text.isdecimal(), (
f"Hotspot residue ID {hotspot_id!r} must have form <residue><idx>, "
"for example L150."
)
hotspot_idx = int(hotspot_idx_text)
assert 1 <= hotspot_idx <= target_length, (
f"Hotspot residue {hotspot_id!r} is out of range for target sequence "
f"length {target_length}. Hotspots are 1-indexed within the provided "
"target sequence."
)
target_residue = target_sequence[hotspot_idx - 1]
assert target_residue == residue, (
f"Hotspot residue ID {hotspot_id!r} does not match target sequence: "
f"target residue at 1-indexed position {hotspot_idx} is "
f"{target_residue!r}."
)
epitope_mask[hotspot_idx - 1] = True
return epitope_mask[None].repeat(batch_size, 1)
def compute_epitope_loss(
distogram_logits: torch.Tensor,
epitope_mask: torch.Tensor,
target_length: int,
binder_length: int,
bin_distance: torch.Tensor,
epitope_contact_distance: float = 12.0,
) -> torch.Tensor:
"""Encourage binder contacts to the requested target hotspot residues.
Converts full-complex distogram logits into contact probabilities below
``epitope_contact_distance``, then penalizes each hotspot by the best few
binder contacts. The target is assumed to occupy the first ``target_length``
positions of the complex and the binder the final ``binder_length`` positions.
Parameters
----------
distogram_logits
Full-complex distogram logits with shape ``(B, L, L, num_bins)``.
epitope_mask
Boolean full-complex mask with shape ``(B, L)`` and true values at
target hotspot positions.
target_length
Number of target residues at the start of the complex sequence.
binder_length
Number of binder residues at the end of the complex sequence.
bin_distance
Distance represented by each distogram bin.
epitope_contact_distance
Distance threshold used to define a hotspot-binder contact.
Returns
-------
torch.Tensor
Per-example epitope losses with shape ``(B,)``.
"""
B, L, L2, _ = distogram_logits.shape
assert L == L2, (L, L2)
assert L == target_length + binder_length, (L, target_length, binder_length)
binder_mask = torch.zeros((B, L), dtype=torch.bool, device=distogram_logits.device)
binder_mask[:, target_length:] = True
positive_mask = bin_distance < epitope_contact_distance
assert binder_mask.shape == (B, L), f"Binder mask: {binder_mask.shape} != {B, L}"
assert epitope_mask.shape == (B, L), f"Epitope mask: {epitope_mask.shape} != {B, L}"
# NOTE - this could be made bidirectional.
epitope_binder_mask = epitope_mask[:, :, None] & binder_mask[:, None, :]
probabilities = torch.softmax(distogram_logits, dim=-1)
positive_probs = torch.sum(probabilities * positive_mask, dim=-1)
contact_measure = -torch.log(positive_probs + 1e-8)
epitope_loss = masked_min_k(contact_measure, mask=epitope_binder_mask, k=5)
return masked_average(epitope_loss, mask=epitope_mask)
def compute_structure_losses(
distogram_logits: torch.Tensor,
binder_length: int,
target_sequence: str | None = None,
target_hotspot_ids: list[str] | None = None,
epitope_contact_distance: float = 12.0,
) -> dict[str, torch.Tensor]:
"""Compute structural losses and a weighted total."""
bin_distance = get_mid_points().to(distogram_logits.device)
losses: dict[str, torch.Tensor] = {}
losses["intra_contact_loss"] = compute_intra_contact_loss(
distogram_logits, binder_length, bin_distance
)
losses["inter_contact_loss"] = compute_inter_contact_loss(
distogram_logits, binder_length, bin_distance
)
losses["glob_loss"] = compute_globularity_loss(
distogram_logits, binder_length, bin_distance
)
B = distogram_logits.size(0)
total = torch.tensor([0.0] * B, device=distogram_logits.device, requires_grad=True)
total = total + LOSS_WEIGHTS["intra_contact"] * losses["intra_contact_loss"]
total = total + LOSS_WEIGHTS["inter_contact"] * losses["inter_contact_loss"]
total = total + LOSS_WEIGHTS["glob"] * losses["glob_loss"]
if target_hotspot_ids is not None:
target_length = distogram_logits.shape[1] - binder_length
assert target_sequence is not None, (
"target_sequence is required when target_hotspot_ids is provided."
)
assert "|" not in target_sequence, (
"Epitope loss only supports one target chain."
)
assert len(target_sequence) == target_length, (
f"Target sequence length {len(target_sequence)} does not match distogram "
f"target length {target_length}."
)
epitope_mask = get_epitope_mask(
target_sequence=target_sequence,
binder_length=binder_length,
target_hotspot_ids=target_hotspot_ids,
batch_size=B,
device=distogram_logits.device,
)
losses["epitope_loss"] = compute_epitope_loss(
distogram_logits=distogram_logits,
epitope_mask=epitope_mask,
target_length=target_length,
binder_length=binder_length,
bin_distance=bin_distance,
epitope_contact_distance=epitope_contact_distance,
)
total = total + LOSS_WEIGHTS["epitope"] * losses["epitope_loss"]
losses["total_loss"] = total
return losses
# ---- Distogram iptm proxy ----
def _binding_confidence_entropy(
dgram: torch.Tensor, bin_distance: torch.Tensor, cutoff: float
) -> torch.Tensor:
"""Pair entropy within cutoff."""
probs = torch.softmax(dgram, dim=-1)
cutoff_mask = bin_distance < cutoff
p_cut = probs[..., cutoff_mask]
p_cut = p_cut / (p_cut.sum(-1, keepdim=True) + 1e-8)
return -(p_cut * torch.log(p_cut + 1e-10)).sum(-1)
def _entropy_to_confidence(mean_entropy: float) -> float:
"""Map mean pair entropy to [0, 1]; lower entropy → higher score."""
return float(max(0.0, min(1.0, 1.0 - mean_entropy / math.log(51))))
def _is_antibody_sequence(binder_sequence: str) -> bool:
"""Return True if ANARCI recognizes binder_sequence as antibody variable domain(s)."""
from abnumber.common import _anarci_align
sequence = binder_sequence.replace(MUTABLE_TOKEN, "A")
result = _anarci_align(
sequences=[sequence], scheme="chothia", allowed_species=None
)[0]
if not result:
return False
valid_chain_types = {"H", "K", "L"}
return all(chain_type in valid_chain_types for _, chain_type, *_ in result)
def _cdr_indices(binder_sequence: str) -> list[int]:
"""0-based binder indices for all Chothia CDRs."""
from abnumber import Chain
from abnumber.common import _anarci_align
result = _anarci_align(
sequences=[binder_sequence], scheme="chothia", allowed_species=None
)[0]
chains = [
Chain("".join(result[i][0].values()), scheme="chothia")
for i in range(len(result))
]
if len(chains) == 2 and not chains[0].is_heavy_chain():
chains.reverse()
indices: list[int] = []
for chain in chains:
for cdr in (chain.cdr1_seq, chain.cdr2_seq, chain.cdr3_seq):
start = binder_sequence.find(cdr)
assert start >= 0
indices.extend(range(start, start + len(cdr)))
return indices
def compute_distogram_iptm_proxy(
distogram_logits: torch.Tensor,
target_length: int,
binder_sequence: str,
is_antibody: bool,
) -> dict[str, float]:
"""Algorithm 15 Distogram ipTM Proxy.
Distogram iptm proxy for a target|binder complex (binder at suffix).
Returns distogram_iptm_proxy for all designs and
cdr_distogram_iptm_proxy when the binder can be annotated as an
antibody; otherwise the CDR score is NaN.
"""
if distogram_logits.ndim == 4:
distogram_logits = distogram_logits[0]
binder_length = len(binder_sequence)
assert distogram_logits.shape[0] == target_length + binder_length
bin_distance = get_mid_points().to(distogram_logits.device)
binder_start = target_length
def _mean_lowest_k(entropies: torch.Tensor, k: int) -> float:
sorted_entropies, _ = torch.sort(entropies.reshape(-1))
k = min(k, sorted_entropies.numel())
return float(sorted_entropies[:k].mean())
binder_to_target_entropy = _binding_confidence_entropy(
distogram_logits[binder_start:, :target_length, :], bin_distance, cutoff=22.0
)
distogram_iptm_proxy = _entropy_to_confidence(
_mean_lowest_k(binder_to_target_entropy, k=binder_length)
)
if not is_antibody:
cdr_distogram_iptm_proxy = float("nan")
else:
cdr_indices = _cdr_indices(binder_sequence)
cdr_rows = [binder_start + i for i in cdr_indices]
cdr_to_target_entropy = _binding_confidence_entropy(
distogram_logits[cdr_rows, :target_length, :], bin_distance, cutoff=22.0
)
cdr_distogram_iptm_proxy = _entropy_to_confidence(
_mean_lowest_k(cdr_to_target_entropy, k=len(cdr_indices))
)
return {
"distogram_iptm_proxy": distogram_iptm_proxy,
"cdr_distogram_iptm_proxy": cdr_distogram_iptm_proxy,
}
# ---- Folding ----
def _resize_tensor(tensor: torch.Tensor, *, dim: int, size: int) -> torch.Tensor:
current = tensor.shape[dim]
if current >= size:
return tensor.narrow(dim, 0, size)
pad_shape = list(tensor.shape)
pad_shape[dim] = size - current
pad = torch.zeros(pad_shape, dtype=tensor.dtype, device=tensor.device)
return torch.cat((tensor, pad), dim=dim)
_ATOM_FEATURE_DIMS = {
"ref_pos": 0,
"ref_element": 0,
"ref_charge": 0,
"ref_atom_name_chars": 0,
"ref_space_uid": 0,
"atom_attention_mask": 0,
"atom_to_token": 0,
"is_resolved": 0,
"gt_coords": 1,
}
@cache
def _ensure_ccd_loaded() -> None:
load_ccd()
def prepare_esmfold2_tensors(
input: StructurePredictionInput,
max_tokens: int | None = None,
max_atoms: int | None = None,
max_seqs: int = 16384,
pad_to_max_seqs: bool = False,
seed: int | None = None,
use_vectorized_msa_assembly: bool = True,
) -> dict[str, torch.Tensor]:
del max_tokens, max_seqs, pad_to_max_seqs, use_vectorized_msa_assembly
_ensure_ccd_loaded()
features, _ = prepare_esmfold2_input(input, seed=seed)
if max_atoms is not None:
for key, dim in _ATOM_FEATURE_DIMS.items():
if key in features:
features[key] = _resize_tensor(features[key], dim=dim, size=max_atoms)
return features
def fold_and_get_distogram(
model: EsmFold2ExperimentalModel,
target_seq: str,
target_one_hot: torch.Tensor,
design: torch.Tensor,
num_loops: int = 0,
num_sampling_steps: int = 1,
calculate_confidence: bool = False,
seed: int | None = None,
) -> dict:
"""Prepare inputs, run model forward, return distogram_logits + raw output."""
padding = (2, 11)
padded_design = F.pad(design, padding, mode="constant", value=0)
# Argmax to get the designed sequence string.
token_lists = torch.argmax(padded_design, dim=-1)
designed_seq = [
[PROTEIN_3TO1[TOKENS[int(tkn.item())]] for tkn in token_list]
for token_list in token_lists
]
seq_list = [target_seq + "|" + "".join(seq) for seq in designed_seq]
max_atoms = None if len(seq_list) == 1 else ((len(seq_list[0]) - 1) * 14) // 32 * 32
inputs_list = []
for seq in seq_list:
sequences = {
sequence: [str(idx)] for idx, sequence in enumerate(seq.split("|"))
}
inputs_raw = StructurePredictionInput(
sequences=[
ProteinInput(id=chain_id, sequence=sequence, msa=None)
for sequence, chain_id in sequences.items()
]
)
inputs_list.append(prepare_esmfold2_tensors(inputs_raw, max_atoms=max_atoms))
inputs = {
key: torch.stack([inp[key] for inp in inputs_list], dim=0).cuda()
for key in inputs_list[0]
}
inputs["res_type_soft"] = torch.cat(
(target_one_hot.repeat(design.size(0), 1, 1), padded_design), dim=1
)
with seed_context(seed):
output = model(
**inputs,
num_diffusion_samples=1,
num_sampling_steps=num_sampling_steps,
num_loops=num_loops,
calculate_confidence=calculate_confidence,
seed=seed,
)
result: dict = {
"distogram_logits": output["distogram_logits"],
"inputs": inputs,
"inputs_list": inputs_list,
"output": output,
"seq_list": seq_list,
}
if calculate_confidence:
result.update(
{
"ptm": output.get("ptm"),
"iptm": output.get("iptm"),
"plddt": output.get("plddt"),
}
)
return result
_CHAIN_ID_ALPHABET = string.ascii_uppercase + string.ascii_lowercase + string.digits
def _asym_id_to_chain_label(asym_id: int) -> str:
if asym_id < 0:
raise ValueError(f"asym_id must be >= 0, got {asym_id}")
label = ""
n = len(_CHAIN_ID_ALPHABET)
while True:
label = _CHAIN_ID_ALPHABET[asym_id % n] + label
asym_id = asym_id // n - 1
if asym_id < 0:
return label
def to_atom_array(
coords: np.ndarray,
atom_to_token: np.ndarray,
res_type: np.ndarray,
residue_index: np.ndarray,
asym_id: np.ndarray,
mol_type: np.ndarray,
ref_atom_name_chars: np.ndarray,
ref_element: np.ndarray,
atom_attention_mask: np.ndarray,
plddt_per_atom: np.ndarray | None = None,
) -> biotite.structure.AtomArray:
atoms = []
for atom_i, (
atom_coord,
token_idx,
atom_name_chars,
element_idx,
is_not_pad,
) in enumerate(
zip(
coords, atom_to_token, ref_atom_name_chars, ref_element, atom_attention_mask
)
):
if not is_not_pad:
continue
atoms.append(
biotite.structure.Atom(
coord=atom_coord,
chain_id=_asym_id_to_chain_label(int(asym_id[token_idx])),
res_id=residue_index[token_idx] + 1,
res_name=TOKENS[res_type[token_idx]],
atom_name="".join(chr(c + 32) for c in atom_name_chars if c != 0),
element=ELEMENTS[element_idx],
ins_code=" ",
hetero=mol_type[token_idx] == MOL_TYPE_NONPOLYMER,
b_factor=float(plddt_per_atom[atom_i])
if plddt_per_atom is not None
else 0.0,
)
)
return biotite.structure.array(atoms)
def build_complex(
inputs: dict[str, torch.Tensor], output: dict[str, Any]
) -> ProteinComplex:
"""Build ProteinComplex from model output."""
plddt_per_atom = output.get("plddt_per_atom")
if plddt_per_atom is not None:
plddt_per_atom = plddt_per_atom[0].cpu().numpy() * PLDDT_B_FACTOR_SCALE
atom_arr = to_atom_array(
coords=output["sample_atom_coords"][0].cpu().numpy(),
atom_to_token=inputs["atom_to_token"][0].cpu().numpy(),
res_type=inputs["res_type"][0].cpu().numpy(),
residue_index=inputs["token_index"][0].cpu().numpy(),
asym_id=inputs["asym_id"][0].cpu().numpy(),
mol_type=inputs["mol_type"][0].cpu().numpy(),
ref_atom_name_chars=inputs["ref_atom_name_chars"][0].cpu().numpy(),
ref_element=inputs["ref_element"][0].cpu().numpy(),
atom_attention_mask=inputs["atom_attention_mask"][0].cpu().numpy(),
plddt_per_atom=plddt_per_atom,
)
return ProteinComplex.from_chains(
[
ProteinChain.from_atomarray(a, is_predicted=True)
for a in biotite.structure.chain_iter(atom_arr)
]
)
# ---- LM loss ----
@cache
def _folding_trunk_to_lm_aa_vocab_matrix(device: torch.device) -> torch.Tensor:
"""Build a matrix of shape [ft_aas=20, lm_aas=20]."""
three_to_one_map = {v: k for k, v in PROTEIN_1TO3.items()}
ft_aas = [three_to_one_map[tok_3letter] for tok_3letter in TOKENS[2:22]]
lm_vocab = sorted(EsmcTokenizer().vocab.items(), key=lambda x: x[1])
lm_aas = [lm_vocab[i][0] for i in range(4, 24)]
ft_to_lm_aa_matrix = torch.zeros(20, 20)
for ft_idx, ft_aa in enumerate(ft_aas):
lm_idx = lm_aas.index(ft_aa)
ft_to_lm_aa_matrix[ft_idx, lm_idx] = 1
return ft_to_lm_aa_matrix.to(device=device)
def _one_hot_from_probs(probs: torch.Tensor) -> torch.Tensor:
return F.one_hot(torch.argmax(probs, dim=-1), num_classes=probs.size(-1)).to(
probs.dtype
)
def _straight_through(discrete: torch.Tensor, continuous: torch.Tensor) -> torch.Tensor:
return continuous + (discrete - continuous).detach()
def compute_esmc_pseudoperplexity_nll(
esmc_model: EsmcForMaskedLM,
binder_design: torch.Tensor,
score_mask: torch.Tensor,
batch_size: int = 4,
n_passes: int = 4,
) -> torch.Tensor:
"""Algorithm 14 ESMC Pseudo-perplexity Sequence Regularization.
Approximate pseudoperplexity NLL via multiple sampled masks."""
device = binder_design.device
lm_vocab_size = esmc_model.config.vocab_size
model_dtype = esmc_model.esmc.embed.weight.dtype
target_esm = binder_design @ _folding_trunk_to_lm_aa_vocab_matrix(device)
input_esm = _straight_through(_one_hot_from_probs(target_esm), target_esm)
input_ids = torch.zeros(
(binder_design.size(0), binder_design.size(1) + 2, lm_vocab_size),
dtype=model_dtype,
device=device,
)
tokenizer = EsmcTokenizer()
input_ids[:, 0, tokenizer.cls_token_id] = 1 # pyright: ignore
input_ids[:, -1, tokenizer.eos_token_id] = 1 # pyright: ignore
input_ids[:, 1:-1, 4:24] = input_esm.to(model_dtype)
if score_mask.ndim == 1:
score_mask = score_mask.unsqueeze(0).expand(binder_design.size(0), -1)
elif score_mask.shape != binder_design.shape[:2]:
raise ValueError(
f"Expected score_mask with shape {(binder_design.size(0), binder_design.size(1))}, "
f"got {tuple(score_mask.shape)}"
)
score_mask = score_mask.to(device=device, dtype=torch.bool)
mask_token = torch.zeros(lm_vocab_size, dtype=model_dtype, device=device)
mask_token[esmc_model.config.mask_token_id] = 1
esmc = esmc_model.esmc
all_masked_sequences = []
all_pass_masks = []
for batch_idx in range(binder_design.size(0)):
position_indices = score_mask[batch_idx].nonzero(as_tuple=False).flatten()
num_positions = int(position_indices.numel())
if num_positions == 0:
raise ValueError(
"ESMC pseudoperplexity score mask selected zero positions."
)
num_masked = max(1, math.ceil(ESMC_MASK_FRACTION * num_positions))
random_scores = torch.rand((n_passes, num_positions), device=device)
masked_offsets = random_scores.topk(num_masked, dim=-1, largest=False).indices
pass_masks = torch.zeros(
(n_passes, binder_design.size(1)), dtype=torch.bool, device=device
)
pass_masks[
torch.arange(n_passes, device=device)[:, None],
position_indices[masked_offsets],
] = True
masked_sequences = input_ids[batch_idx : batch_idx + 1].repeat(n_passes, 1, 1)
mask_rows, mask_cols = pass_masks.nonzero(as_tuple=True)
masked_sequences[mask_rows, mask_cols + 1] = mask_token
all_masked_sequences.append(masked_sequences)
all_pass_masks.append(pass_masks)
logit_chunks = []
masked_sequence_rows = torch.cat(all_masked_sequences, dim=0)
for start in range(0, masked_sequence_rows.size(0), batch_size):
chunk = masked_sequence_rows[start : start + batch_size]
with torch.autocast(
device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"
):
hidden, *_ = esmc.transformer(
chunk @ esmc.embed.weight.to(chunk.dtype),
sequence_id=None,
layers_to_collect=[],
output_attentions=False,
)
logit_chunks.append(esmc_model.lm_head(hidden))
logits = torch.cat(logit_chunks, dim=0)
losses = []
for batch_idx, pass_masks in enumerate(all_pass_masks):
start = batch_idx * n_passes
stop = start + n_passes
target_weights = target_esm[batch_idx]
log_probs = logits[start:stop].log_softmax(dim=-1)[:, 1:-1, 4:24]
nlls = -(log_probs * target_weights.to(log_probs.dtype).unsqueeze(0)).sum(
dim=-1
)
losses.append(nlls[pass_masks].mean())
return torch.stack(losses, dim=0)
# ---- Design ----
def normalized_gradient_tensor(
grad: torch.Tensor, gradient_mask: torch.Tensor
) -> torch.Tensor:
masked_grad = grad * gradient_mask
index_has_nonzero_grad = torch.square(masked_grad).sum(-1) > 0 # (B, L)
eff_L = index_has_nonzero_grad.sum(-1) # (B,)
grad_norm = torch.linalg.norm(masked_grad, axis=(-1, -2)) # (B,)
normalized_grad = (masked_grad / (grad_norm[:, None, None] + 1e-7)) * torch.sqrt(
eff_L[:, None, None]
)
return normalized_grad * gradient_mask
def design_binder(
inversion_models: dict[str, EsmFold2ExperimentalModel],
hf_critic_models: dict[str, EsmFold2ExperimentalModel],
esmc_model: EsmcForMaskedLM,
scaling_critic_names: list[str] | None,
target_name: str,
target_sequence: str | None,
binder_name: str,
binder_sequence: str | None,
is_antibody: bool | None,
seed: int,
batch_size: int = 1,
target_hotspot_ids: list[str] | None = None,
epitope_contact_distance: float = 12.0,
) -> tuple[list[str], Trajectory, list[dict]]:
"""
Algorithm 11 Gradient-Guided Binder Sequence Optimization.
Run the full optimization loop.
Returns dict with designed_sequence, complex, and trajectory.
Every critic is folded once on the best designed sequence via HF ESMFold2.
Hero critics expose iPTM; scaling critics contribute distogram scores only.
``distogram_binding_confidence`` / ``cdr_distogram_binding_confidence`` come
from the distogram in all cases.
``target_hotspot_ids`` enables epitope loss and is interpreted as 1-indexed
residue IDs within the provided target sequence, for example ``["L150"]``.
"""
# Setup
device = "cuda"
if target_name in TARGET_SEQUENCES:
if target_sequence is not None:
raise ValueError(
f"{target_name!r} is a preset target; omit target_sequence."
)
target_sequence = TARGET_SEQUENCES[target_name]
elif target_sequence is None:
raise ValueError(
f"{target_name!r} is not a preset target; provide target_sequence."
)
target_one_hot = sequence_to_one_hot(target_sequence, device=device)
if binder_name in BINDER_PROMPT_FACTORIES:
if binder_sequence is not None:
raise ValueError(
f"{binder_name!r} is a preset binder; omit binder_sequence."
)
binder_prompt_factor = BINDER_PROMPT_FACTORIES[binder_name]
if is_antibody is not None:
assert binder_prompt_factor.is_antibody == is_antibody, (
"Conflict in is_antibody settings."
)
is_antibody = binder_prompt_factor.is_antibody
binder_sequence = binder_prompt_factor.sample(seed=seed)
elif binder_sequence is None:
raise ValueError(
f"{binder_name!r} is not a preset binder; provide binder_sequence."
)
elif is_antibody is None:
is_antibody = _is_antibody_sequence(binder_sequence)
binder_length = len(binder_sequence)
# By default, we only support single binder and target chains.
# To support this case, remove the asserts below and check that losses
# and selection metrics are appropriate for your multi-chain case.
assert "|" not in target_sequence
assert "|" not in binder_sequence
with seed_context(seed), torch.device(device):
logits = build_initial_soft_sequence_logits(
binder_sequence, batch_size=batch_size
)
gradient_mask = build_gradient_mask(binder_sequence, batch_size=batch_size)
# step -> {loss_name: [B] tensor on CPU, metric_name: float}
trajectory: Trajectory = {}
global_step = 0
def run_step(
logits: torch.Tensor,
optimizer: optim.Optimizer,
temperature: float,
calculate_confidence: bool,
) -> tuple[torch.Tensor, list[str], list[float] | None]:
nonlocal global_step
torch.cuda.reset_peak_memory_stats()
start = time.time()
optimizer.zero_grad()
random.seed(seed + global_step)
replicate_choice = random.randint(0, len(inversion_models) - 1)
inversion_model = list(inversion_models.values())[replicate_choice]
design = F.softmax(logits / temperature, dim=-1)
fold_result = fold_and_get_distogram(
inversion_model,
target_sequence,
target_one_hot,
design,
num_loops=1,
num_sampling_steps=50 if calculate_confidence else 1,
calculate_confidence=calculate_confidence,
seed=seed + global_step,
)
sequences: list[str] = fold_result["seq_list"]
losses = compute_structure_losses(
fold_result["distogram_logits"],
binder_length,
target_sequence=target_sequence,
target_hotspot_ids=target_hotspot_ids,
epitope_contact_distance=epitope_contact_distance,
)
structure_loss = losses["total_loss"]
structure_grad = torch.autograd.grad(structure_loss.mean(), logits)[0]
# Recompute the logits -> design transform for a fresh graph.
design = F.softmax(logits / temperature, dim=-1)
score_mask = gradient_mask.sum(dim=-1) > 0
with seed_context(seed + global_step):
plm_loss = compute_esmc_pseudoperplexity_nll(
esmc_model=esmc_model,
binder_design=design,
score_mask=score_mask,
batch_size=LM_LOSS_BATCH_SIZE,
n_passes=LM_MASK_PASSES,
)
plm_grad = torch.autograd.grad(plm_loss.mean(), logits)[0]
logits.grad = normalized_gradient_tensor(structure_grad, gradient_mask) + (
0.05 if is_antibody else 0.15
) * normalized_gradient_tensor(plm_grad, gradient_mask)
for g in optimizer.param_groups:
g["lr"] = LEARNING_RATE * temperature
optimizer.step()
step = global_step
step_losses: TrajectoryStep = {k: v.detach().cpu() for k, v in losses.items()}
step_losses["plm_loss"] = plm_loss.detach().cpu()
step_losses["total_loss"] = (structure_loss + plm_loss).detach().cpu()
step_losses["time"] = time.time() - start
step_losses["peak_allocated_gib"] = (
torch.cuda.max_memory_allocated() / BYTES_PER_GIB
)
step_losses["peak_reserved_gib"] = (
torch.cuda.max_memory_reserved() / BYTES_PER_GIB
)
trajectory[step] = step_losses
loss_str = " ".join(
f"{k}={v.mean().item() if torch.is_tensor(v) else v:.4f}"
for k, v in step_losses.items()
)
if step % LOG_INTERVAL == 0:
logger.info(f" step {step:3d} | {loss_str} T={temperature:.4f}")
global_step += 1
return logits, sequences, fold_result.get("iptm", None)
# Optimize
optimizer = optim.SGD([logits], lr=LEARNING_RATE)
best_iptm: list[float] = [-1.0] * batch_size
best_sequences: list[str] = [""] * batch_size
for step in range(STEPS):
# Cosine schedule
t = (step + 1) / STEPS
remaining = 0.5 * (1 + math.cos(math.pi * t))
temperature = TEMPERATURE_MIN + (1 - TEMPERATURE_MIN) * remaining
logits, sequences, iptm = run_step(
logits,
optimizer,
temperature=temperature,
calculate_confidence=temperature < 0.05,
)
if iptm is not None:
for b in range(batch_size):
if iptm[b] is not None and iptm[b] > best_iptm[b]:
best_iptm[b] = iptm[b]
best_sequences[b] = sequences[b]
assert all(seq != "" for seq in best_sequences)
# Score
critic_results: list[dict] = []
target_length = len(target_sequence.replace("|", ""))
final_total_loss = trajectory[global_step - 1]["total_loss"]
assert isinstance(final_total_loss, torch.Tensor)
def score_critic(
critic_name: str,
critic_model: EsmFold2ExperimentalModel,
batch_idx: int,
is_scaling_critic: bool,
) -> None:
best_seq = best_sequences[batch_idx]
binder_seq = best_seq.split("|")[-1]
binder_design = sequence_to_one_hot(binder_seq)[..., 2:22]
final_fold = fold_and_get_distogram(
critic_model,
target_sequence,
target_one_hot,
binder_design,
num_loops=3,
num_sampling_steps=1 if is_scaling_critic else 200,
calculate_confidence=not is_scaling_critic,
seed=seed,
)
pred_complex = (
None
if is_scaling_critic
else build_complex(final_fold["inputs"], final_fold["output"])
)
iptm_proxy_scores = compute_distogram_iptm_proxy(
final_fold["distogram_logits"], target_length, binder_seq, is_antibody
)
iptm_tensor = final_fold.get("iptm")
iptm = iptm_tensor.item() if iptm_tensor is not None else None
critic_results.append(
{
"is_antibody": is_antibody,
"critic_name": critic_name,
"is_scaling_critic": is_scaling_critic,
"batch_idx": batch_idx,
"designed_sequence": best_seq,
"complex": pred_complex,
"final_loss": final_total_loss[batch_idx].item(),
"iptm": iptm,
"logits": logits[batch_idx].detach().cpu(),
**iptm_proxy_scores,
}
)
del final_fold
for critic_name, critic_model in hf_critic_models.items():
for batch_idx in range(batch_size):
score_critic(critic_name, critic_model, batch_idx, is_scaling_critic=False)
for critic_name in scaling_critic_names or []:
critic_model = _load_hf_model(
critic_name, lm_dropout=0.25, cache_esmc=False, device="cuda"
)
for batch_idx in range(batch_size):
score_critic(critic_name, critic_model, batch_idx, is_scaling_critic=True)
del critic_model
torch.cuda.empty_cache()
if not critic_results:
for batch_idx in range(batch_size):
critic_results.append(
{
"is_antibody": is_antibody,
"batch_idx": batch_idx,
"designed_sequence": best_sequences[batch_idx],
"final_loss": final_total_loss[batch_idx].item(),
"logits": logits[batch_idx].detach().cpu(),
}
)
return best_sequences, trajectory, critic_results
# ---- Model Loading ----
_ESMC_CACHE: dict[str, torch.nn.Module] = {}
def _load_hf_model(
critic_name: str, lm_dropout: float, cache_esmc: bool, device: str
) -> Any:
"""Loads ESMFold2 from huggingface. Will cache ESMC by checkpoint ID among
all non-scaling checkpoints, to save on VRAM and load time."""
repo_id = f"biohub/{critic_name}"
model = EsmFold2ExperimentalModel.from_pretrained(repo_id, load_esmc=not cache_esmc)
if cache_esmc:
esmc_id = model.config.esmc_id
if esmc_id not in _ESMC_CACHE:
model.load_esmc(esmc_id)
assert model._esmc is not None
_ESMC_CACHE[esmc_id] = model._esmc
model._esmc = _ESMC_CACHE[esmc_id]
model.configure_lm_dropout(lm_dropout, force_lm_dropout_during_inference=True)
kernel_backend = None
if TRITON_KERNELS_AVAILABLE:
kernel_backend = BACKEND_FUSED
elif CUE_AVAILABLE:
kernel_backend = BACKEND_CUEQ
model.set_kernel_backend(kernel_backend)
return model.to(device=device).eval().requires_grad_(False)
def _apply_torch_compile(model: torch.nn.Module) -> None:
"""A helper for torch compiling the model."""
torch._dynamo.config.cache_size_limit = 512
torch._dynamo.config.accumulated_cache_size_limit = 512
compile_targets = (EsmFold2MSAEncoder, PairUpdateBlock)
def _maybe_compile_module(module: torch.nn.Module) -> None:
if not isinstance(module, compile_targets):
return
module.forward = torch.compile(module.forward) # ty:ignore[invalid-assignment]
model.apply(_maybe_compile_module)
class ESMFold2Design:
lm_name = "biohub/ESMC-6B"
inversion_model_names: list[str] = [
"ESMFold2-Experimental-Fast",
"ESMFold2-Experimental-Fast-Cutoff2025",
]
hero_critic_hf_paths: list[str] = [
"ESMFold2-Experimental-Fast",
"ESMFold2-Experimental-Fast-Cutoff2025",
"ESMFold2-Experimental",
"ESMFold2-Experimental-Cutoff2025",
]
scaling_critic_hf_paths: list[str] = []
def load(self, use_scaling_critics: bool):
self.scaling_critic_hf_paths = []
if use_scaling_critics:
self.scaling_critic_hf_paths = [
f"ESMFold2-Experimental-Fast-base{size}-step{step}k"
for size in ("300M", "600M", "6B")
for step in ("250", "500", "750", "1000", "1500")
]
self.inversion_models = {
model_name: _load_hf_model(
model_name, lm_dropout=0.5, cache_esmc=True, device="cuda"
)
for model_name in self.inversion_model_names
}
if COMPILE:
for model in self.inversion_models.values():
_apply_torch_compile(model)
self.hf_critic_models: dict[str, Any] = {}
for name in self.hero_critic_hf_paths:
self.hf_critic_models[name] = _load_hf_model(
name, lm_dropout=0.25, cache_esmc=True, device="cuda"
)
self.esmc_model = EsmcForMaskedLM.from_pretrained(
self.lm_name, dtype=torch.float32
)
if REUSE_ESMC:
reusable_esmc_model = self.inversion_models["ESMFold2-Experimental-Fast"]
assert reusable_esmc_model.config.esmc_id == self.lm_name, (
f"Cannot reuse ESMC trunk from {reusable_esmc_model.config.esmc_id!r} "
f"with LM head from {self.lm_name!r}."
)
assert reusable_esmc_model._esmc is not None
del self.esmc_model.esmc
torch.cuda.empty_cache()
self.esmc_model.esmc = reusable_esmc_model._esmc
self.esmc_model = self.esmc_model.cuda().eval().requires_grad_(False)
def design(
self,
target_name: str,
binder_name: str,
target_sequence: str | None = None,
binder_sequence: str | None = None,
is_antibody: bool | None = None,
seed: int = 0,
batch_size: int = 1,
target_hotspot_ids: list[str] | None = None,
epitope_contact_distance: float = 12.0,
) -> tuple[list[str], Trajectory, list[dict]]:
return design_binder(
self.inversion_models,
self.hf_critic_models,
self.esmc_model,
self.scaling_critic_hf_paths,
target_name=target_name,
target_sequence=target_sequence,
binder_name=binder_name,
binder_sequence=binder_sequence,
is_antibody=is_antibody,
seed=seed,
batch_size=batch_size,
target_hotspot_ids=target_hotspot_ids,
epitope_contact_distance=epitope_contact_distance,
)
# ---- Modal ----
def get_base_image():
return (
modal.Image.micromamba(python_version="3.12")
.run_commands("apt update && apt install -y git build-essential")
.micromamba_install(
"anarci>=2020.04.03",
"hmmer=3.4",
"cuda-version=12.8",
"cuda-libraries-dev=12.8",
"cuda-nvcc=12.8",
"cmake",
"ninja",
channels=["conda-forge", "bioconda"],
)
.pip_install(
"torch==2.8.0",
"triton==3.4.0",
index_url="https://download.pytorch.org/whl/cu128",
)
.pip_install(
"flash-attn==2.8.3",
"transformer-engine[core-cu12,pytorch]==2.13.0",
"xformers==0.0.32.post1",
extra_options="--no-build-isolation",
env={
"CPATH": (
"/opt/conda/lib/python3.12/site-packages/nvidia/cudnn/include:"
"/opt/conda/lib/python3.12/site-packages/nvidia/nccl/include:"
"/opt/conda/lib/python3.12/site-packages/nvidia/nvtx/include"
),
"LIBRARY_PATH": (
"/opt/conda/lib/python3.12/site-packages/nvidia/cudnn/lib:"
"/opt/conda/lib/python3.12/site-packages/nvidia/nccl/lib:"
"/opt/conda/lib/python3.12/site-packages/nvidia/nvtx/lib"
),
"MAX_JOBS": "8",
"NVTE_FRAMEWORK": "pytorch",
},
)
.pip_install(
"abnumber", "esm@git+https://github.com/Biohub/esm.git@main", "modal"
)
.env(
{
"HF_HOME": "/models",
"HF_XET_HIGH_PERFORMANCE": "1",
"XFORMERS_IGNORE_FLASH_VERSION_CHECK": "1",
}
)
)
app = modal.App(
name="esmfold2-design",
image=get_base_image(),
volumes={
"/models": modal.Volume.from_name("esmfold2-models", create_if_missing=True)
},
)
# NOTE - Currently the memory usage is quite high, inflating costs.
# In an update coming soon, scaling critics will be loaded on demand
# to avoid needing this large amount of RAM.
@app.cls(gpu="H100", timeout=60 * 60, cpu=16, memory=80 * 1024)
class ESMFold2DesignModal(ESMFold2Design):
"""Modal entrypoint. Hero critics are HF experimental exports with
confidence heads. Set ``use_scaling_critics=True`` to also load the
15-checkpoint scaling-experiment ensemble (distogram binding confidence only).
"""
use_scaling_critics: bool = modal.parameter(default=True)
deterministic: bool = modal.parameter(default=False)
@modal.enter()
def load(self):
if self.deterministic:
torch.use_deterministic_algorithms(True, warn_only=True)
return super().load(self.use_scaling_critics)
@modal.method()
def design(self, *args, **kws):
return super().design(*args, **kws)
@app.local_entrypoint()
def main(
target_name: str,
binder_name: str,
target_sequence: str | None = None,
binder_sequence: str | None = None,
use_scaling_critics: bool = True,
is_antibody: bool | None = None,
local: bool = False,
seed: int = 0,
batch_size: int = 1,
target_hotspot_ids: list[str] | None = None,
epitope_contact_distance: float = 12.0,
):
if local:
assert not use_scaling_critics, (
"'abnumber' will fail if running this script with uv run. "
"It requires conda packages. To be addressed soon."
)
app = ESMFold2Design()
app.load(use_scaling_critics)
run_fn = app.design
else:
app = ESMFold2DesignModal(
use_scaling_critics=use_scaling_critics # ty:ignore[unknown-argument]
)
run_fn = app.design.remote
seq, trajectory, results = run_fn(
target_name=target_name,
target_sequence=target_sequence,
binder_name=binder_name,
binder_sequence=binder_sequence,
is_antibody=is_antibody,
seed=seed,
batch_size=batch_size,
target_hotspot_ids=target_hotspot_ids,
epitope_contact_distance=epitope_contact_distance,
)
avg_final_loss = sum(r["final_loss"] for r in results) / len(results)
logger.info(f"\nDesigned sequence: {seq}")
logger.info(f"Trajectory length: {len(trajectory)} steps")
logger.info(f"Average final loss: {avg_final_loss:.4f}")
if __name__ == "__main__":
# Run a single local design.
main(
# Example case 1
target_name="pd-l1",
binder_name="minibinder",
is_antibody=False,
# Example case 2
# target_name="cd45",
# binder_name="trastuzumab_framework_vhvl",
# is_antibody=True,
# Common settings
seed=0,
batch_size=1,
local=True,
use_scaling_critics=True,
)
================================================
FILE: cookbook/tutorials/embed.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# [Tutorial](https://github.com/biohub/esm/tree/main/cookbook/tutorials): Embedding with ESMC\n",
"\n",
"In this notebook we will see how to embed a batch of sequences using ESMC, as well as explore its different layers"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Imports"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# If you are working in colab, uncomment these lines to install dependencies\n",
"#! pip install esm\n",
"#! pi
gitextract_5y54ehgk/
├── .github/
│ ├── scripts/
│ │ └── airtable_issue_sync.py
│ └── workflows/
│ ├── airtable-issue-sync.yaml
│ └── ci.yml
├── .gitignore
├── .pre-commit-config.yaml
├── CONTRIBUTIONS.md
├── LICENSE.md
├── README.md
├── THIRD_PARTY_NOTICE.md
├── _assets/
│ └── ESM3_README.md
├── cookbook/
│ ├── local/
│ │ ├── README.md
│ │ ├── open_generate.ipynb
│ │ └── raw_forwards.py
│ ├── snippets/
│ │ ├── README.md
│ │ ├── esm3.py
│ │ ├── esmc.py
│ │ ├── fold_invfold.py
│ │ ├── sae.py
│ │ ├── sae_example.py
│ │ └── sparse_utils.py
│ └── tutorials/
│ ├── README.md
│ ├── binder_design.ipynb
│ ├── binder_design.py
│ ├── embed.ipynb
│ ├── esm3_generate.ipynb
│ ├── esm3_guided_generation.ipynb
│ ├── esmc_finetune.ipynb
│ ├── esmc_layer_sweep.ipynb
│ ├── esmc_mutation_scoring.ipynb
│ ├── esmc_sae_feature_interpretation.ipynb
│ ├── esmfold2.ipynb
│ ├── esmfold2_local_applesilicon.ipynb
│ ├── esmfold2_local_gpu.ipynb
│ ├── esmprotein.ipynb
│ ├── g3l5_chainA.a3m
│ ├── g3l5_chainB.a3m
│ └── gfp_design.ipynb
├── esm/
│ ├── __init__.py
│ ├── data/
│ │ ├── ParentChildTreeFile.txt
│ │ ├── entry_list_safety_29026.list
│ │ ├── interpro_29026_to_keywords_58641.csv
│ │ ├── keyword_idf_safety_filtered_58641.npy
│ │ └── keyword_vocabulary_safety_filtered_58641.txt
│ ├── layers/
│ │ ├── attention.py
│ │ ├── blocks.py
│ │ ├── codebook.py
│ │ ├── ffn.py
│ │ ├── geom_attention.py
│ │ ├── regression_head.py
│ │ ├── rotary.py
│ │ ├── structure_proj.py
│ │ └── transformer_stack.py
│ ├── models/
│ │ ├── esm3.py
│ │ ├── esmc/
│ │ │ ├── __init__.py
│ │ │ ├── checkpoint_layout.py
│ │ │ ├── compatibility.py
│ │ │ ├── config.py
│ │ │ ├── kernels.py
│ │ │ ├── layers.py
│ │ │ ├── model.py
│ │ │ ├── sae.py
│ │ │ └── tokenizer.py
│ │ ├── esmfold2/
│ │ │ ├── __init__.py
│ │ │ ├── config.py
│ │ │ ├── conformers.py
│ │ │ ├── constants.py
│ │ │ ├── experimental.py
│ │ │ ├── hf_adapter.py
│ │ │ ├── hf_checkpoint.py
│ │ │ ├── kernels/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── fused_attention_pair_bias.py
│ │ │ │ ├── fused_dropout_residual.py
│ │ │ │ ├── fused_dual_gemm.py
│ │ │ │ ├── fused_ln_residual.py
│ │ │ │ ├── fused_lnlin_swiglu.py
│ │ │ │ ├── trimul_einsum_triton.py
│ │ │ │ └── trimul_with_residual.py
│ │ │ ├── layers.py
│ │ │ ├── model.py
│ │ │ ├── output.py
│ │ │ ├── paired_msa.py
│ │ │ ├── prepare_input.py
│ │ │ ├── processor.py
│ │ │ ├── protein_utils.py
│ │ │ └── types.py
│ │ ├── function_decoder.py
│ │ ├── hub.py
│ │ └── vqvae.py
│ ├── pretrained.py
│ ├── sdk/
│ │ ├── __init__.py
│ │ ├── api.py
│ │ ├── base_forge_client.py
│ │ ├── experimental/
│ │ │ ├── __init__.py
│ │ │ ├── constrained_generation.py
│ │ │ └── guided_generation.py
│ │ ├── forge.py
│ │ ├── retry.py
│ │ ├── sagemaker.py
│ │ └── validation.py
│ ├── tokenization/
│ │ ├── __init__.py
│ │ ├── function_tokenizer.py
│ │ ├── residue_tokenizer.py
│ │ ├── sasa_tokenizer.py
│ │ ├── sequence_tokenizer.py
│ │ ├── ss_tokenizer.py
│ │ ├── structure_tokenizer.py
│ │ └── tokenizer_base.py
│ ├── utils/
│ │ ├── constants/
│ │ │ ├── api.py
│ │ │ ├── esm3.py
│ │ │ ├── models.py
│ │ │ └── physics.py
│ │ ├── decoding.py
│ │ ├── encoding.py
│ │ ├── forge_context_manager.py
│ │ ├── function/
│ │ │ ├── encode_decode.py
│ │ │ ├── interpro.py
│ │ │ ├── lsh.py
│ │ │ └── tfidf.py
│ │ ├── generation.py
│ │ ├── misc.py
│ │ ├── msa/
│ │ │ ├── __init__.py
│ │ │ ├── filter_sequences.py
│ │ │ └── msa.py
│ │ ├── noise_schedules.py
│ │ ├── parsing.py
│ │ ├── residue_constants.py
│ │ ├── sampling.py
│ │ ├── sequential_dataclass.py
│ │ ├── structure/
│ │ │ ├── affine3d.py
│ │ │ ├── aligner.py
│ │ │ ├── atom_indexer.py
│ │ │ ├── input_builder.py
│ │ │ ├── metrics.py
│ │ │ ├── mmcif_parsing.py
│ │ │ ├── molecular_complex.py
│ │ │ ├── normalize_coordinates.py
│ │ │ ├── predicted_aligned_error.py
│ │ │ ├── protein_chain.py
│ │ │ ├── protein_complex.py
│ │ │ └── protein_structure.py
│ │ ├── system.py
│ │ └── types.py
│ └── widgets/
│ ├── components/
│ │ ├── function_annotator.py
│ │ ├── results_visualizer.py
│ │ ├── sasa_prompt_selector.py
│ │ ├── secondary_structure_prompt_selector.py
│ │ ├── sequence_prompt_selector.py
│ │ └── structure_prompt_selector.py
│ ├── utils/
│ │ ├── clients.py
│ │ ├── drawing/
│ │ │ ├── colors.py
│ │ │ ├── draw_category_array.py
│ │ │ ├── draw_function_annotations.py
│ │ │ └── draw_protein_structure.py
│ │ ├── indexing.py
│ │ ├── parsing.py
│ │ ├── printing.py
│ │ ├── prompting.py
│ │ ├── protein_import.py
│ │ ├── serialization.py
│ │ └── types.py
│ └── views/
│ ├── esm3_generation_launcher.py
│ ├── esm3_prompt_preview.py
│ ├── esm3_prompt_selector.py
│ ├── generation.py
│ ├── inverse_folding.py
│ ├── login.py
│ └── prediction.py
├── pyproject.toml
├── tests/
│ ├── Makefile
│ ├── __init__.py
│ ├── compatibility/
│ │ ├── __init__.py
│ │ ├── compatibility_test.py
│ │ ├── esmc_legacy_contract_test.py
│ │ ├── esmfold2_hf_adapter_test.py
│ │ └── esmfold2_hf_checkpoint_test.py
│ ├── conftest.py
│ ├── models/
│ │ ├── __init__.py
│ │ ├── esmc_test.py
│ │ ├── esmfold2_api_test.py
│ │ ├── esmfold2_builds_test.py
│ │ ├── esmfold2_cpu_only_test.py
│ │ ├── esmfold2_execution_test.py
│ │ ├── esmfold2_inputs_test.py
│ │ ├── esmfold2_msa_test.py
│ │ ├── esmfold2_sampler_test.py
│ │ ├── esmfold2_test.py
│ │ └── prepare_input_test.py
│ ├── oss_pytests/
│ │ ├── Dockerfile
│ │ ├── requirements.txt
│ │ ├── test_oss_client.py
│ │ ├── test_output_attentions.py
│ │ └── test_placeholder.py
│ ├── regenerate_reference.py
│ ├── sdk/
│ │ ├── __init__.py
│ │ └── forge_context_manager_test.py
│ └── utils/
│ ├── __init__.py
│ ├── input_builder_test.py
│ ├── misc_test.py
│ ├── molecular_complex_test.py
│ ├── msa_test.py
│ └── sampling_test.py
└── tools/
└── README.md
Copy disabled (too large)
Condensed preview — 201 files, each showing path, character count, and a content snippet. Download the .json file for the full structured content (18,715K chars).
[
{
"path": ".github/scripts/airtable_issue_sync.py",
"chars": 13400,
"preview": "#!/usr/bin/env python3\n\"\"\"Sync a GitHub issue into an Airtable base.\n\nEnv: AIRTABLE_TOKEN, AIRTABLE_BASE, AIRTABLE_TABLE..."
},
{
"path": ".github/workflows/airtable-issue-sync.yaml",
"chars": 905,
"preview": "name: Sync issues to Airtable\n\non:\n issues:\n types: [opened, edited, reopened, closed, labeled, unlabeled]\n\npermissi..."
},
{
"path": ".github/workflows/ci.yml",
"chars": 3063,
"preview": "name: ESM Tests\n\non:\n pull_request:\n branches:\n - \"**\"\n workflow_dispatch:\n merge_group:\n types: [checks_r..."
},
{
"path": ".gitignore",
"chars": 150,
"preview": "esm.egg-info\n# pixi environments\n.pixi\n*.egg-info\n*.pyc\n\n# pytest --cov artifacts\n.coverage\n.coverage.*\ncoverage.xml\nhtm..."
},
{
"path": ".pre-commit-config.yaml",
"chars": 1275,
"preview": "# See https://pre-commit.com for more information\n# See https://pre-commit.com/hooks.html for more hooks\nexclude: (fasta..."
},
{
"path": "CONTRIBUTIONS.md",
"chars": 1377,
"preview": "We welcome community contributions to help make this package better!\n\n## Contributing\n\n_to be written_\n\n## Testing\n\nImpr..."
},
{
"path": "LICENSE.md",
"chars": 1087,
"preview": "**License (MIT)**\n\nCopyright 2026 Chan Zuckerberg Biohub, Inc.\n\nPermission is hereby granted, free of charge, to any per..."
},
{
"path": "README.md",
"chars": 19280,
"preview": "<div align=\"center\">\n <img src=\"_assets/header.png\" style=\"width: 60%; height: auto;\" />\n\n# A world model of protein bi..."
},
{
"path": "THIRD_PARTY_NOTICE.md",
"chars": 948,
"preview": "The code in this repository depends on the following third-party libraries:\n\n| Library | License | Link |\n|----------|--..."
},
{
"path": "_assets/ESM3_README.md",
"chars": 6442,
"preview": "# ESM3 README\n\n[ESM3](https://www.science.org/doi/10.1126/science.ads0018) is a frontier generative model for biology, a..."
},
{
"path": "cookbook/local/README.md",
"chars": 47,
"preview": "Examples utilizing the open model run locally.\n"
},
{
"path": "cookbook/local/open_generate.ipynb",
"chars": 26049,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# ESM3\\n\",\n \"\\n\",\n \"ESM3 is a..."
},
{
"path": "cookbook/local/raw_forwards.py",
"chars": 5514,
"preview": "import random\n\nimport torch\nimport torch.nn.functional as F\n\nfrom esm.pretrained import (\n ESM3_function_decoder_v0,..."
},
{
"path": "cookbook/snippets/README.md",
"chars": 79,
"preview": "Snippets of ESM3 usage that you can copy and paste directly into your scripts.\n"
},
{
"path": "cookbook/snippets/esm3.py",
"chars": 7559,
"preview": "import os\n\nfrom esm.models.esm3 import ESM3\nfrom esm.sdk import client\nfrom esm.sdk.api import (\n ESM3InferenceClient..."
},
{
"path": "cookbook/snippets/esmc.py",
"chars": 5955,
"preview": "import math\nimport os\n\nimport torch\n\nfrom esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer\nfrom esm.sdk import esmc..."
},
{
"path": "cookbook/snippets/fold_invfold.py",
"chars": 4173,
"preview": "import os\nfrom typing import cast\n\nimport numpy as np\n\nfrom esm.sdk.api import (\n ESM3InferenceClient,\n ESMProtein..."
},
{
"path": "cookbook/snippets/sae.py",
"chars": 2082,
"preview": "import numpy as np\nimport torch\n\nfrom cookbook.snippets.sparse_utils import max_pool, remove_indexes\nfrom esm.sdk import..."
},
{
"path": "cookbook/snippets/sae_example.py",
"chars": 1501,
"preview": "import os\n\nfrom cookbook.snippets.sae import get_sae_features, get_sae_features_single\nfrom cookbook.snippets.sparse_uti..."
},
{
"path": "cookbook/snippets/sparse_utils.py",
"chars": 4684,
"preview": "from typing import Iterable\n\nimport torch\n\n\ndef remove_indexes(\n sparse_coo_tensor: torch.Tensor, indexes_to_remove:..."
},
{
"path": "cookbook/tutorials/README.md",
"chars": 5030,
"preview": "# **ESM Tutorial Notebooks**\n\nTutorial notebooks are the best way to get hands-on with ESM models\\! Use the notebooks to..."
},
{
"path": "cookbook/tutorials/binder_design.ipynb",
"chars": 23703,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"b5b44288\",\n \"metadata\": {},\n \"source\": [\n \"## [Tutorial](..."
},
{
"path": "cookbook/tutorials/binder_design.py",
"chars": 54787,
"preview": "# /// script\n# requires-python = \"<=3.13\"\n# dependencies = [\n# \"abnumber\",\n# \"esm@git+https://github.com/Biohub/..."
},
{
"path": "cookbook/tutorials/embed.ipynb",
"chars": 115336,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# [Tutorial](https://github.com/bio..."
},
{
"path": "cookbook/tutorials/esm3_generate.ipynb",
"chars": 26727,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# [Tutorial](https://github.com/bio..."
},
{
"path": "cookbook/tutorials/esm3_guided_generation.ipynb",
"chars": 14361,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# [Tutorial](https://github.com/bio..."
},
{
"path": "cookbook/tutorials/esmc_finetune.ipynb",
"chars": 22795,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"6b74e995\",\n \"metadata\": {},\n \"source\": [\n \"# Finetuning E..."
},
{
"path": "cookbook/tutorials/esmc_layer_sweep.ipynb",
"chars": 252902,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"725805fc\",\n \"metadata\": {},\n \"source\": [\n \"# [Tutorial](h..."
},
{
"path": "cookbook/tutorials/esmc_mutation_scoring.ipynb",
"chars": 208508,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"TAZqskyDHKdz\",\n \"metadata\": {\n \"id\": \"TAZqskyDHKdz\"\n },..."
},
{
"path": "cookbook/tutorials/esmc_sae_feature_interpretation.ipynb",
"chars": 52258,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"cell-1\",\n \"metadata\": {\n \"id\": \"cell-1\"\n },\n \"source\":..."
},
{
"path": "cookbook/tutorials/esmfold2.ipynb",
"chars": 2416327,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"JN3w1zjabgRM\"\n },\n \"source\": [\n \"# [Tut..."
},
{
"path": "cookbook/tutorials/esmfold2_local_applesilicon.ipynb",
"chars": 18860,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"bd417770\",\n \"metadata\": {},\n \"source\": [\n \"# Folding prot..."
},
{
"path": "cookbook/tutorials/esmfold2_local_gpu.ipynb",
"chars": 16996,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"12326c8b\",\n \"metadata\": {},\n \"source\": [\n \"# Folding prot..."
},
{
"path": "cookbook/tutorials/esmprotein.ipynb",
"chars": 29016,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# [Tutorial](https://github.com/bio..."
},
{
"path": "cookbook/tutorials/g3l5_chainA.a3m",
"chars": 10116,
"preview": ">key=0\nGPYYPTNKLQAAVMETDRENAIIRQRNDEIPTRTLDTAIFTDASTVASAQIHLYYNSNIGKIIMSLNGKKHTFNLYDDNDIRTLLPILLLSK\n>key=1\n--YYPTNKLQAAV..."
},
{
"path": "cookbook/tutorials/g3l5_chainB.a3m",
"chars": 8571,
"preview": ">key=0\nGPNMFFMPKRKIPDPIDRLRRANLACEDDKLMIYGLPWMTTQTSALSINSKPIVYKDCAKLLRSINGSQPVSLNDVLRR\n>key=1\n--NMFFMPKRKIPDPIDRLRRANLAC..."
},
{
"path": "cookbook/tutorials/gfp_design.ipynb",
"chars": 22246,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"zWXOAcBB8h3z\"\n },\n \"source\": [\n \"# [Tut..."
},
{
"path": "esm/__init__.py",
"chars": 22,
"preview": "__version__ = \"3.4.0\"\n"
},
{
"path": "esm/data/ParentChildTreeFile.txt",
"chars": 594786,
"preview": "IPR000008::C2 domain::\n--IPR014705::Synaptotagmin-17, C2B domain::\n--IPR028692::Synaptotagmin-13, C2B domain::\n--IPR0305..."
},
{
"path": "esm/data/entry_list_safety_29026.list",
"chars": 1595245,
"preview": "ENTRY_AC\tENTRY_TYPE\tENTRY_NAME\nIPR000001\tDomain\tKringle\nIPR000003\tFamily\tRetinoid X receptor/HNF4\nIPR000006\tFamily\tMetal..."
},
{
"path": "esm/data/interpro_29026_to_keywords_58641.csv",
"chars": 10061298,
"preview": "interpro_id,keywords\nIPR000001,kringle\nIPR000003,\"retinoid,receptorhnf4,retinoid x,x receptorhnf4,nuclear,steroid,nuclea..."
},
{
"path": "esm/data/keyword_vocabulary_safety_filtered_58641.txt",
"chars": 787696,
"preview": "0 gamma\n0 group\n0 metabolic\n1 16\n1 2\n1 2mata\n1 3\n1 32\n1 4\n1 4zasp\n1 5\n1 5idn2\n1 5phosphorelay\n1 6\n1 6sfr4\n1 8\n1 acid\n1 a..."
},
{
"path": "esm/layers/attention.py",
"chars": 5205,
"preview": "import functools\n\nimport einops\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\n\nfrom esm.layers.rotar..."
},
{
"path": "esm/layers/blocks.py",
"chars": 6198,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom esm.layers.attention import FlashMultiHeadAtten..."
},
{
"path": "esm/layers/codebook.py",
"chars": 2983,
"preview": "import numpy as np\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\nimport torch.nn.functional as F..."
},
{
"path": "esm/layers/ffn.py",
"chars": 688,
"preview": "import torch.nn as nn\nimport torch.nn.functional as F\nfrom torch import Tensor\n\n# NOT CURRENTLY USED\n\n\nclass SwiGLU(nn.M..."
},
{
"path": "esm/layers/geom_attention.py",
"chars": 5879,
"preview": "from math import sqrt\n\nimport torch\nfrom einops import rearrange\nfrom torch import nn\nfrom torch.nn import functional as..."
},
{
"path": "esm/layers/regression_head.py",
"chars": 633,
"preview": "import torch.nn as nn\n\n\ndef RegressionHead(\n d_model: int, output_dim: int, hidden_dim: int | None = None\n) -> nn.Mod..."
},
{
"path": "esm/layers/rotary.py",
"chars": 10797,
"preview": "# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's G..."
},
{
"path": "esm/layers/structure_proj.py",
"chars": 2510,
"preview": "import torch\nimport torch.nn as nn\n\nfrom esm.utils.constants.physics import BB_COORDINATES\nfrom esm.utils.structure.affi..."
},
{
"path": "esm/layers/transformer_stack.py",
"chars": 4626,
"preview": "import math\n\nimport torch\nimport torch.nn as nn\n\nfrom esm.layers.blocks import UnifiedTransformerBlock\nfrom esm.utils.st..."
},
{
"path": "esm/models/esm3.py",
"chars": 23417,
"preview": "from __future__ import annotations\n\nimport contextlib\nfrom functools import partial\nfrom typing import Callable\n\nimport..."
},
{
"path": "esm/models/esmc/__init__.py",
"chars": 1076,
"preview": "from esm.models.esmc.compatibility import ESMC, ESMCOutput\nfrom esm.models.esmc.config import EsmcConfig\nfrom esm.models..."
},
{
"path": "esm/models/esmc/checkpoint_layout.py",
"chars": 7297,
"preview": "\"\"\"Translation between the published ESMC checkpoint layout and this\nimplementation's in-memory layout.\n\nPublished check..."
},
{
"path": "esm/models/esmc/compatibility.py",
"chars": 12086,
"preview": "\"\"\"Deprecated ``ESMC`` surface, kept so existing code keeps running.\n\n``ESMC`` used to be one class holding an encoder,..."
},
{
"path": "esm/models/esmc/config.py",
"chars": 8487,
"preview": "\"\"\"Configuration for the ESMC models.\"\"\"\n\nimport json\nimport os\nimport warnings\nfrom dataclasses import asdict, dataclas..."
},
{
"path": "esm/models/esmc/kernels.py",
"chars": 3258,
"preview": "\"\"\"Optional accelerated kernels used by the ESMC layers.\n\nESMC runs on pure PyTorch by default, but selects fused kernel..."
},
{
"path": "esm/models/esmc/layers.py",
"chars": 27427,
"preview": "\"\"\"ESMC-specific transformer layers.\n\nRoPE, QK-LayerNorm attention, SwiGLU feed-forward and the transformer stack used\nb..."
},
{
"path": "esm/models/esmc/model.py",
"chars": 32846,
"preview": "\"\"\"ESMC encoder and task heads.\n\nESMC is a protein language model trained by EvolutionaryScale with a\nmasked-token objec..."
},
{
"path": "esm/models/esmc/sae.py",
"chars": 14466,
"preview": "\"\"\"ESMC sparse autoencoder (SAE) model.\n\n* :class:`EsmcSaeModel` - the published container, one repo per\n ``(backbone,..."
},
{
"path": "esm/models/esmc/tokenizer.py",
"chars": 5319,
"preview": "\"\"\"Tokenization for ESMC.\"\"\"\n\nfrom tokenizers import AddedToken, Tokenizer\nfrom tokenizers.models import BPE\nfrom tokeni..."
},
{
"path": "esm/models/esmfold2/__init__.py",
"chars": 1526,
"preview": "from esm.models.esmfold2.config import (\n ESMFOLD2_EXPERIMENTAL_HF_REPO,\n ESMFOLD2_HF_REPO,\n EsmFold2Config,\n)..."
},
{
"path": "esm/models/esmfold2/config.py",
"chars": 22552,
"preview": "# Copyright 2026 Biohub. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you..."
},
{
"path": "esm/models/esmfold2/conformers.py",
"chars": 9127,
"preview": "\"\"\"CCD conformer loading utilities.\n\nLoads idealized conformer coordinates from a CCD pickle file containing RDKit molec..."
},
{
"path": "esm/models/esmfold2/constants.py",
"chars": 10383,
"preview": "\"\"\"Constants for the ESMFold2 input pipeline.\n\nIncludes molecule types, residue types, vocabularies, atom lists, and ele..."
},
{
"path": "esm/models/esmfold2/experimental.py",
"chars": 46538,
"preview": "# coding=utf-8\n# Copyright 2026 Biohub. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Li..."
},
{
"path": "esm/models/esmfold2/hf_adapter.py",
"chars": 15535,
"preview": "\"\"\"Present this package's ESMFold2 API on top of the upstream HuggingFace port.\n\n:class:`EsmFold2HFAdapter` wraps ``tran..."
},
{
"path": "esm/models/esmfold2/hf_checkpoint.py",
"chars": 11324,
"preview": "\"\"\"Read a HuggingFace-layout ESMFold2 checkpoint into this implementation.\n\nThe upstream ``transformers`` port renames m..."
},
{
"path": "esm/models/esmfold2/kernels/__init__.py",
"chars": 840,
"preview": "# coding=utf-8\n# Copyright 2026 Biohub. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Li..."
},
{
"path": "esm/models/esmfold2/kernels/fused_attention_pair_bias.py",
"chars": 24735,
"preview": "\"\"\"Fused Triton kernel for AttentionPairBias (forward + backward).\n\nInspired by cuequivariance's ``attention_pair_bias``..."
},
{
"path": "esm/models/esmfold2/kernels/fused_dropout_residual.py",
"chars": 8451,
"preview": "\"\"\"Fused row-shared-dropout + residual-add kernel for the pair stream.\n\nFor the pairformer pattern:\n\n pair_new = pair..."
},
{
"path": "esm/models/esmfold2/kernels/fused_dual_gemm.py",
"chars": 22163,
"preview": "\"\"\"Native Triton ``fused_sigmoid_gated_dual_gemm`` for TriMul stage 2.\n\nComputes ``sigmoid(x @ w1.T) * (x @ w2.T)`` with..."
},
{
"path": "esm/models/esmfold2/kernels/fused_ln_residual.py",
"chars": 14468,
"preview": "\"\"\"Triton LayerNorm (bf16 IO, fp32 stats) with optional fused residual-add in\nthe *backward* pass.\n\nInspired by cuequiva..."
},
{
"path": "esm/models/esmfold2/kernels/fused_lnlin_swiglu.py",
"chars": 13597,
"preview": "\"\"\"Fused LayerNorm + Linear(d, 2*d_inner) + SwiGLU(silu(x1) * x2) kernel.\n\nCollapses the standard LayerNorm -> Linear ->..."
},
{
"path": "esm/models/esmfold2/kernels/trimul_einsum_triton.py",
"chars": 8381,
"preview": "\"\"\"Triton kernel for trimul stage-3 batched einsum in native (D, B, L, L) layout.\n\n outgoing: ``out[d,b,i,j] = sum_k..."
},
{
"path": "esm/models/esmfold2/kernels/trimul_with_residual.py",
"chars": 21666,
"preview": "\"\"\"TriMul with output residual + dropout-mask epilogue fused into the final GEMM.\n\nInspired by cuequivariance's ``triang..."
},
{
"path": "esm/models/esmfold2/layers.py",
"chars": 104313,
"preview": "# coding=utf-8\n# Copyright 2026 Biohub. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Li..."
},
{
"path": "esm/models/esmfold2/model.py",
"chars": 53256,
"preview": "\"\"\"PyTorch ESMFold2 model — the standard released architecture.\n\nQuickstart::\n\n from transformers import EsmFold2Mode..."
},
{
"path": "esm/models/esmfold2/output.py",
"chars": 8353,
"preview": "from itertools import groupby\nfrom typing import Any\n\nimport numpy as np\nimport torch\n\nfrom esm.models.esmfold2.constant..."
},
{
"path": "esm/models/esmfold2/paired_msa.py",
"chars": 9376,
"preview": "\"\"\"Taxonomy-paired MSA construction for ESMFold2 inference.\n\nTaxonomy IDs are read from FASTA headers as ``key=N`` token..."
},
{
"path": "esm/models/esmfold2/prepare_input.py",
"chars": 53177,
"preview": "\"\"\"Prepare ESMFold2 model inputs from sequence-level StructurePredictionInput.\n\nThis module converts StructurePrediction..."
},
{
"path": "esm/models/esmfold2/processor.py",
"chars": 16862,
"preview": "import random\nimport warnings\nfrom contextlib import contextmanager, nullcontext\nfrom pathlib import Path\nfrom typing im..."
},
{
"path": "esm/models/esmfold2/protein_utils.py",
"chars": 26186,
"preview": "# coding=utf-8\n# Copyright 2026 Biohub. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Li..."
},
{
"path": "esm/models/esmfold2/types.py",
"chars": 807,
"preview": "\"\"\"Re-exports of the canonical SPI dataclasses from input_builder.\n\nThis module exists so the HF processor and downstrea..."
},
{
"path": "esm/models/function_decoder.py",
"chars": 13277,
"preview": "\"\"\"Function Token Decoder.\"\"\"\n\nfrom collections import defaultdict\nfrom dataclasses import dataclass, field\n\nimport nump..."
},
{
"path": "esm/models/hub.py",
"chars": 8195,
"preview": "\"\"\"Shared HuggingFace Hub loading for the opensource model packages.\n\nBoth ESMC and ESMFold2 publish plain safetensors c..."
},
{
"path": "esm/models/vqvae.py",
"chars": 15895,
"preview": "import torch\nimport torch.nn as nn\n\nfrom esm.layers.blocks import UnifiedTransformerBlock\nfrom esm.layers.codebook impor..."
},
{
"path": "esm/pretrained.py",
"chars": 5860,
"preview": "import inspect\nimport warnings\nfrom typing import Callable\n\nimport torch\nimport torch.nn as nn\nfrom accelerate import in..."
},
{
"path": "esm/sdk/__init__.py",
"chars": 3305,
"preview": "import os\nimport warnings\n\nfrom esm.sdk.api import ESM3InferenceClient, ESMCInferenceClient\nfrom esm.sdk.forge import (..."
},
{
"path": "esm/sdk/api.py",
"chars": 31958,
"preview": "from __future__ import annotations\n\nimport warnings\nfrom abc import ABC\nfrom copy import deepcopy\nfrom typing import Seq..."
},
{
"path": "esm/sdk/base_forge_client.py",
"chars": 15961,
"preview": "import asyncio\nimport time\nfrom abc import ABC, abstractmethod\nfrom contextlib import suppress\nfrom typing import Any, G..."
},
{
"path": "esm/sdk/experimental/__init__.py",
"chars": 375,
"preview": "from .constrained_generation import (\n ConstraintType,\n ESM3GuidedDecodingWithConstraints,\n GenerationConstrain..."
},
{
"path": "esm/sdk/experimental/constrained_generation.py",
"chars": 11350,
"preview": "from __future__ import annotations\n\nfrom dataclasses import dataclass, field\nfrom enum import Enum\nfrom typing import Li..."
},
{
"path": "esm/sdk/experimental/guided_generation.py",
"chars": 9339,
"preview": "from abc import ABC, abstractmethod\nfrom typing import Tuple\n\nimport attr\nimport torch\nfrom tqdm import tqdm\n\nfrom esm.m..."
},
{
"path": "esm/sdk/forge.py",
"chars": 54602,
"preview": "from __future__ import annotations\n\nimport asyncio\nimport base64\nimport pickle\nimport warnings\nfrom concurrent.futures i..."
},
{
"path": "esm/sdk/retry.py",
"chars": 2901,
"preview": "import inspect\nfrom contextvars import ContextVar\nfrom functools import wraps\n\nfrom tenacity import (\n retry,\n ret..."
},
{
"path": "esm/sdk/sagemaker.py",
"chars": 4249,
"preview": "import json\n\nimport boto3\n\nfrom esm.sdk.forge import (\n ESM3ForgeInferenceClient,\n SequenceStructureForgeInference..."
},
{
"path": "esm/sdk/validation.py",
"chars": 818,
"preview": "from esm.utils.structure.input_builder import ProteinInput, StructurePredictionInput\n\n\ndef validate_fold_max_accuracy_in..."
},
{
"path": "esm/tokenization/__init__.py",
"chars": 2253,
"preview": "from dataclasses import dataclass\nfrom typing import Protocol\n\nfrom esm.utils.constants.models import ESM3_OPEN_SMALL, n..."
},
{
"path": "esm/tokenization/function_tokenizer.py",
"chars": 14996,
"preview": "\"\"\"Tokenizes annotations of protein function.\"\"\"\n\nimport re\nimport string\nfrom functools import cache, cached_property,..."
},
{
"path": "esm/tokenization/residue_tokenizer.py",
"chars": 7911,
"preview": "from functools import cached_property\nfrom typing import Any\n\nimport pandas as pd\nimport torch\nimport torch.nn.functiona..."
},
{
"path": "esm/tokenization/sasa_tokenizer.py",
"chars": 4887,
"preview": "from functools import cached_property\n\nimport torch\n\nfrom esm.tokenization.tokenizer_base import EsmTokenizerBase\nfrom e..."
},
{
"path": "esm/tokenization/sequence_tokenizer.py",
"chars": 4098,
"preview": "from tokenizers import Tokenizer\nfrom tokenizers.models import BPE\nfrom tokenizers.processors import TemplateProcessing..."
},
{
"path": "esm/tokenization/ss_tokenizer.py",
"chars": 3644,
"preview": "from functools import cached_property\nfrom typing import Sequence\n\nimport torch\n\nfrom esm.tokenization.tokenizer_base im..."
},
{
"path": "esm/tokenization/structure_tokenizer.py",
"chars": 2697,
"preview": "from esm.tokenization.tokenizer_base import EsmTokenizerBase\nfrom esm.utils.constants import esm3 as C\n\n\nclass Structure..."
},
{
"path": "esm/tokenization/tokenizer_base.py",
"chars": 513,
"preview": "from typing import Protocol, runtime_checkable\n\n\n@runtime_checkable\nclass EsmTokenizerBase(Protocol):\n mask_token: st..."
},
{
"path": "esm/utils/constants/api.py",
"chars": 247,
"preview": "MAX_TOPK_SEQUENCE = 32\nMAX_TOPK_STRUCTURE = MAX_TOPK_SEQUENCE\nMAX_TOPK_SECONDARY_STRUCTURE = MAX_TOPK_SEQUENCE\nMAX_TOPK_..."
},
{
"path": "esm/utils/constants/esm3.py",
"chars": 3373,
"preview": "import os\nfrom functools import cache\nfrom pathlib import Path\n\nfrom huggingface_hub import snapshot_download\n\nSEQUENCE_..."
},
{
"path": "esm/utils/constants/models.py",
"chars": 1134,
"preview": "# Model names\nESM3_OPEN_SMALL = \"esm3_sm_open_v1\"\nESM3_OPEN_SMALL_ALIAS_1 = \"esm3-open-2024-03\"\nESM3_OPEN_SMALL_ALIAS_2..."
},
{
"path": "esm/utils/constants/physics.py",
"chars": 112,
"preview": "BB_COORDINATES = [\n [0.5256, 1.3612, 0.0000],\n [0.0000, 0.0000, 0.0000],\n [-1.5251, 0.0000, 0.0000],\n]\n"
},
{
"path": "esm/utils/decoding.py",
"chars": 9344,
"preview": "import pickle\nimport warnings\nfrom typing import Any, Mapping, cast\n\nimport attr\nimport torch\nfrom requests import Respo..."
},
{
"path": "esm/utils/encoding.py",
"chars": 8317,
"preview": "from typing import Sequence\n\nimport torch\nimport torch.nn.functional as F\n\nfrom esm.models.vqvae import StructureTokenEn..."
},
{
"path": "esm/utils/forge_context_manager.py",
"chars": 6201,
"preview": "import threading\nfrom collections import deque\nfrom concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait..."
},
{
"path": "esm/utils/function/encode_decode.py",
"chars": 6959,
"preview": "import re\nfrom typing import Sequence\n\nimport torch\n\nfrom esm.models.function_decoder import FunctionTokenDecoder, merge..."
},
{
"path": "esm/utils/function/interpro.py",
"chars": 5798,
"preview": "\"\"\"Utilities for interacting with InterPro.\"\"\"\n\nimport itertools\nimport re\nfrom dataclasses import dataclass\nfrom enum i..."
},
{
"path": "esm/utils/function/lsh.py",
"chars": 3402,
"preview": "import numpy as np\nfrom cloudpathlib import AnyPath\n\nfrom esm.utils.types import PathLike\n\n\nclass LSHTable:\n def __in..."
},
{
"path": "esm/utils/function/tfidf.py",
"chars": 1937,
"preview": "\"\"\"Term-Frequency / Inverse Document Frequency (TF-IDF) model.\"\"\"\n\nfrom collections import Counter\nfrom functools import..."
},
{
"path": "esm/utils/generation.py",
"chars": 30774,
"preview": "import os\nfrom typing import Any, Callable, Sequence\nfrom warnings import warn\n\nimport attr\nimport torch\nfrom tqdm impor..."
},
{
"path": "esm/utils/misc.py",
"chars": 16216,
"preview": "from __future__ import annotations\n\nimport os\nfrom collections import defaultdict\nfrom contextlib import nullcontext\nfro..."
},
{
"path": "esm/utils/msa/__init__.py",
"chars": 141,
"preview": "from esm.utils.msa.msa import MSA, FastMSA, remove_insertions_from_sequence\n\n__all__ = [\"MSA\", \"FastMSA\", \"remove_insert..."
},
{
"path": "esm/utils/msa/filter_sequences.py",
"chars": 2816,
"preview": "import os\nimport tempfile\nfrom pathlib import Path\n\nimport numpy as np\nfrom scipy.spatial.distance import cdist\n\nfrom es..."
},
{
"path": "esm/utils/msa/msa.py",
"chars": 23790,
"preview": "from __future__ import annotations\n\nimport dataclasses\nimport string\nfrom dataclasses import dataclass\nfrom functools im..."
},
{
"path": "esm/utils/noise_schedules.py",
"chars": 628,
"preview": "import math\n\nimport torch\n\n\ndef cosine_schedule(t: torch.Tensor):\n # t is a tensor of size (batch_size,) with values..."
},
{
"path": "esm/utils/parsing.py",
"chars": 3533,
"preview": "import io\nfrom pathlib import Path\nfrom typing import Generator, Iterable, NamedTuple\n\nPathOrBuffer = str | Path | io.Te..."
},
{
"path": "esm/utils/residue_constants.py",
"chars": 41457,
"preview": "# Copyright 2025 EvolutionaryScale\n# Copyright 2021 AlQuraishi Laboratory\n# Copyright 2021 DeepMind Technologies Limited..."
},
{
"path": "esm/utils/sampling.py",
"chars": 12442,
"preview": "import warnings\nfrom typing import Literal\n\nimport attr\nimport torch\nimport torch.nn.functional as F\n\nfrom esm.sdk.api i..."
},
{
"path": "esm/utils/sequential_dataclass.py",
"chars": 7073,
"preview": "from abc import ABC, abstractmethod\nfrom dataclasses import dataclass, fields, replace\nfrom typing import TypeVar\n\nimpor..."
},
{
"path": "esm/utils/structure/affine3d.py",
"chars": 19947,
"preview": "from __future__ import annotations\n\nimport typing as T\nfrom abc import ABC\nfrom dataclasses import dataclass\n\nimport tor..."
},
{
"path": "esm/utils/structure/aligner.py",
"chars": 3145,
"preview": "from __future__ import annotations\n\nfrom dataclasses import Field, replace\nfrom typing import Any, ClassVar, Protocol, T..."
},
{
"path": "esm/utils/structure/atom_indexer.py",
"chars": 450,
"preview": "import numpy as np\n\nfrom esm.utils.structure.protein_structure import index_by_atom_name\n\n\nclass AtomIndexer:\n def __..."
},
{
"path": "esm/utils/structure/input_builder.py",
"chars": 9579,
"preview": "from dataclasses import dataclass\nfrom typing import Any, Sequence, TypeAlias, Union\n\nimport numpy as np\n\nfrom esm.utils..."
},
{
"path": "esm/utils/structure/metrics.py",
"chars": 15816,
"preview": "import numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom torch import Tensor\nfr..."
},
{
"path": "esm/utils/structure/mmcif_parsing.py",
"chars": 20817,
"preview": "from __future__ import annotations\n\nimport functools\nimport io\nimport os\nfrom dataclasses import dataclass\nfrom datetime..."
},
{
"path": "esm/utils/structure/molecular_complex.py",
"chars": 57947,
"preview": "from __future__ import annotations\n\nimport io\nimport os\nimport re\nfrom dataclasses import asdict, dataclass\nfrom pathlib..."
},
{
"path": "esm/utils/structure/normalize_coordinates.py",
"chars": 2791,
"preview": "from typing import TypeVar\n\nimport numpy as np\nimport torch\nfrom torch import Tensor\n\nfrom esm.utils import residue_cons..."
},
{
"path": "esm/utils/structure/predicted_aligned_error.py",
"chars": 3324,
"preview": "import torch\nimport torch.nn.functional as F\n\nfrom esm.utils.structure.affine3d import Affine3D\n\n\ndef masked_mean(\n m..."
},
{
"path": "esm/utils/structure/protein_chain.py",
"chars": 54238,
"preview": "from __future__ import annotations\n\nimport io\nimport warnings\nfrom dataclasses import asdict, dataclass, replace\nfrom fu..."
},
{
"path": "esm/utils/structure/protein_complex.py",
"chars": 49241,
"preview": "from __future__ import annotations\n\nimport io\nimport itertools\nimport random\nimport re\nimport warnings\nfrom dataclasses..."
},
{
"path": "esm/utils/structure/protein_structure.py",
"chars": 11376,
"preview": "from __future__ import annotations\n\nfrom typing import Tuple, TypeVar\n\nimport numpy as np\nimport torch\nimport torch.nn.f..."
},
{
"path": "esm/utils/system.py",
"chars": 1201,
"preview": "import io\nimport subprocess\nimport typing as T\nfrom pathlib import Path\n\nPathLike = T.Union[str, Path]\nPathOrBuffer = T...."
},
{
"path": "esm/utils/types.py",
"chars": 893,
"preview": "from __future__ import annotations\n\nimport io\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing imp..."
},
{
"path": "esm/widgets/components/function_annotator.py",
"chars": 4463,
"preview": "from typing import Callable\n\nimport pygtrie\nfrom ipywidgets import widgets\n\nfrom esm.sdk.api import FunctionAnnotation\nf..."
},
{
"path": "esm/widgets/components/results_visualizer.py",
"chars": 13499,
"preview": "from datetime import datetime\nfrom functools import partial\nfrom typing import Any, Callable, Literal\n\nimport ipywidgets..."
},
{
"path": "esm/widgets/components/sasa_prompt_selector.py",
"chars": 4660,
"preview": "from typing import Any, Callable, Sequence\n\nimport ipywidgets as widgets\n\nfrom esm.utils.structure.protein_chain import..."
},
{
"path": "esm/widgets/components/secondary_structure_prompt_selector.py",
"chars": 5671,
"preview": "from typing import Any, Callable, Sequence\n\nimport ipywidgets as widgets\nimport pydssp\n\nfrom esm.utils.structure.protein..."
},
{
"path": "esm/widgets/components/sequence_prompt_selector.py",
"chars": 6182,
"preview": "from typing import Callable\n\nimport ipywidgets as widgets\n\nfrom esm.widgets.utils.drawing.colors import (\n hex_to_rgb..."
},
{
"path": "esm/widgets/components/structure_prompt_selector.py",
"chars": 14353,
"preview": "from functools import partial\nfrom typing import Callable\n\nimport ipywidgets as widgets\nimport matplotlib.pyplot as plt..."
},
{
"path": "esm/widgets/utils/clients.py",
"chars": 867,
"preview": "import os\n\nimport huggingface_hub\nimport huggingface_hub.errors\nimport torch\n\nfrom esm.models.esm3 import ESM3\nfrom esm...."
},
{
"path": "esm/widgets/utils/drawing/colors.py",
"chars": 1048,
"preview": "def hex_to_rgba_tuple(hex_color, alpha=1.0):\n hex_color = hex_color.lstrip(\"#\")\n r, g, b = tuple(int(hex_color[i :..."
},
{
"path": "esm/widgets/utils/drawing/draw_category_array.py",
"chars": 4144,
"preview": "import random\nfrom typing import Sequence\n\nimport ipywidgets as widgets\nimport matplotlib.colors as mcolors\nimport matpl..."
},
{
"path": "esm/widgets/utils/drawing/draw_function_annotations.py",
"chars": 2459,
"preview": "import io\nfrom contextlib import contextmanager\n\nimport matplotlib\nimport matplotlib.pyplot as plt\nfrom dna_features_vie..."
},
{
"path": "esm/widgets/utils/drawing/draw_protein_structure.py",
"chars": 757,
"preview": "import py3Dmol\nfrom IPython.display import clear_output\nfrom ipywidgets import widgets\n\nfrom esm.utils.structure.protein..."
},
{
"path": "esm/widgets/utils/indexing.py",
"chars": 1399,
"preview": "import numpy as np\n\nfrom esm.utils.structure.protein_chain import ProteinChain\n\nZERO_INDEX = \"Zero index\"\nPDB_INDEX = \"P..."
},
{
"path": "esm/widgets/utils/parsing.py",
"chars": 506,
"preview": "def convert_range_string_to_list_of_ranges(range_str: str) -> list[tuple[int, int]]:\n def parse_range(range_str: str)..."
},
{
"path": "esm/widgets/utils/printing.py",
"chars": 150,
"preview": "import textwrap\n\n\ndef wrapped_print(text, width=70):\n text = str(text)\n wrapped_text = textwrap.fill(text, width=w..."
},
{
"path": "esm/widgets/utils/prompting.py",
"chars": 15654,
"preview": "from collections import defaultdict\nfrom typing import Any, Callable, Sequence\n\nimport matplotlib.pyplot as plt\nimport t..."
},
{
"path": "esm/widgets/utils/protein_import.py",
"chars": 6190,
"preview": "import codecs\nfrom io import StringIO\nfrom typing import Callable\n\nfrom ipywidgets import widgets\n\nfrom esm.utils.struct..."
},
{
"path": "esm/widgets/utils/serialization.py",
"chars": 2303,
"preview": "import base64\nimport json\nfrom io import StringIO\nfrom typing import Literal\n\nfrom ipywidgets import widgets\n\nfrom esm.s..."
},
{
"path": "esm/widgets/utils/types.py",
"chars": 684,
"preview": "from typing import Any, Callable, Literal, TypedDict\n\nfrom esm.sdk.api import ESM3InferenceClient\n\n\nclass ClientInitCont..."
},
{
"path": "esm/widgets/views/esm3_generation_launcher.py",
"chars": 7353,
"preview": "import datetime\nimport traceback\nfrom typing import Any, Callable, Literal\n\nfrom ipywidgets import widgets\n\nfrom esm.mod..."
},
{
"path": "esm/widgets/views/esm3_prompt_preview.py",
"chars": 5417,
"preview": "import torch\nfrom ipywidgets import widgets\n\nfrom esm.sdk.api import ESMProtein, FunctionAnnotation\nfrom esm.utils.const..."
},
{
"path": "esm/widgets/views/esm3_prompt_selector.py",
"chars": 3785,
"preview": "from ipywidgets import widgets\n\nfrom esm.widgets.components.sasa_prompt_selector import create_sasa_prompt_selector\nfrom..."
},
{
"path": "esm/widgets/views/generation.py",
"chars": 9007,
"preview": "from typing import Any, Literal\n\nfrom ipywidgets import widgets\n\nfrom esm.sdk.api import ESM3InferenceClient, ESMProtein..."
},
{
"path": "esm/widgets/views/inverse_folding.py",
"chars": 3500,
"preview": "from ipywidgets import widgets\n\nfrom esm.sdk.api import (\n ESM3InferenceClient,\n ESMProtein,\n ESMProteinError,..."
},
{
"path": "esm/widgets/views/login.py",
"chars": 6241,
"preview": "import os\nfrom functools import partial\nfrom textwrap import dedent\n\nfrom ipywidgets import widgets\n\nfrom esm.widgets.ut..."
},
{
"path": "esm/widgets/views/prediction.py",
"chars": 6018,
"preview": "from ipywidgets import widgets\n\nfrom esm.sdk.api import (\n ESM3InferenceClient,\n ESMProtein,\n ESMProteinError,..."
},
{
"path": "pyproject.toml",
"chars": 6789,
"preview": "[project]\nname = \"esm\"\nversion = \"3.4.0\"\ndescription = \"EvolutionaryScale open model repository\"\nreadme = \"README.md\"\nre..."
},
{
"path": "tests/Makefile",
"chars": 454,
"preview": "# OSS-specific variables and commands\nDOCKER_TAG ?= dev\nDOCKER_IMAGE_OSS=oss_pytests:${DOCKER_TAG}\nINFRA_PROVIDER ?= AWS..."
},
{
"path": "tests/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "tests/compatibility/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "tests/compatibility/compatibility_test.py",
"chars": 5966,
"preview": "\"\"\"Pins the deprecated ``ESMC`` surface so the compatibility wrapper can't rot.\"\"\"\n\nimport warnings\n\nimport pytest\nimpor..."
},
{
"path": "tests/compatibility/esmc_legacy_contract_test.py",
"chars": 12320,
"preview": "\"\"\"Pins the ESM 3.x public contract for ``ESMC``, as shipped at tag ``v3.2.3``.\n\nEvery assertion comes from ``v3.2.3`` (..."
},
{
"path": "tests/compatibility/esmfold2_hf_adapter_test.py",
"chars": 18131,
"preview": "\"\"\"``EsmFold2HFAdapter``: our ESMFold2 API presented over the upstream HF port.\n\nThe port is not importable here (transf..."
},
{
"path": "tests/compatibility/esmfold2_hf_checkpoint_test.py",
"chars": 10420,
"preview": "\"\"\"Loading an upstream HuggingFace-port ESMFold2 checkpoint into this package.\n\nThe remap is chosen from the tensor key..."
},
{
"path": "tests/conftest.py",
"chars": 18723,
"preview": "\"\"\"Shared fixtures for the ESMC and ESMFold2 test suites.\n\nTiny randomly-initialised models cover the structural contrac..."
},
{
"path": "tests/models/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "tests/models/esmc_test.py",
"chars": 57741,
"preview": "\"\"\"ESMC test suite.\n\nStructural tests run on CPU against tiny randomly-initialised models. Tests that\nneed published wei..."
},
{
"path": "tests/models/esmfold2_api_test.py",
"chars": 13912,
"preview": "\"\"\"ESMFold2 public-API tests — ``fold``, ``infer_protein``, ``infer_protein_as_pdb``.\n\nAll cases run on CPU against a ti..."
},
{
"path": "tests/models/esmfold2_builds_test.py",
"chars": 32682,
"preview": "\"\"\"ESMFold2 build configurations.\n\nThe axes a user can turn without changing the weights: kernel backend, chunk\nsize, fl..."
},
{
"path": "tests/models/esmfold2_cpu_only_test.py",
"chars": 2410,
"preview": "# coding=utf-8\n# Copyright 2026 Biohub. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"Li..."
},
{
"path": "tests/models/esmfold2_execution_test.py",
"chars": 24950,
"preview": "\"\"\"ESMFold2 execution correctness against the stored tiny-model reference.\n\nA randomly-initialised model at a fixed seed..."
},
{
"path": "tests/models/esmfold2_inputs_test.py",
"chars": 36649,
"preview": "\"\"\"ESMFold2 featurizer, multi-chain and decode tests.\n\nEverything here runs on CPU against a tiny randomly-initialised m..."
},
{
"path": "tests/models/esmfold2_msa_test.py",
"chars": 19138,
"preview": "\"\"\"ESMFold2 MSA-path tests.\n\nEvery case runs on CPU against a tiny randomly-initialised model with\n``msa_encoder.enabled..."
},
{
"path": "tests/models/esmfold2_sampler_test.py",
"chars": 26685,
"preview": "\"\"\"ESMFold2 sampler invariances.\n\nThese are the ESMFold2 analogue of ESMC's\n``test_real_sequences_score_better_than_rand..."
},
{
"path": "tests/models/esmfold2_test.py",
"chars": 49513,
"preview": "\"\"\"ESMFold2 test suite.\n\nStructural tests run on CPU against a tiny randomly-initialised model with\nsynthetic LM states,..."
},
{
"path": "tests/models/prepare_input_test.py",
"chars": 3788,
"preview": "\"\"\"Tests for ESMFold2 input preparation (prepare_input).\"\"\"\n\nimport pytest\nfrom rdkit import Chem\n\nfrom esm.models.esmfo..."
},
{
"path": "tests/oss_pytests/Dockerfile",
"chars": 487,
"preview": "# Dockerfile.sdktest\nFROM python:3.12-slim\n\n# Install pip and basic dependencies\nRUN apt-get update && apt-get install -..."
},
{
"path": "tests/oss_pytests/requirements.txt",
"chars": 31,
"preview": "esm >=3.2.1post1,<4.0.0\npytest\n"
},
{
"path": "tests/oss_pytests/test_oss_client.py",
"chars": 2879,
"preview": "import os\n\nimport pytest\nimport torch\n\nfrom esm.sdk import client # pyright: ignore\nfrom esm.sdk.api import ( # pyrigh..."
},
{
"path": "tests/oss_pytests/test_output_attentions.py",
"chars": 4134,
"preview": "import torch\n\nfrom esm.models.esm3 import ESM3, ESMOutput\nfrom esm.models.esmc import ESMC, ESMCOutput\nfrom esm.tokeniza..."
},
{
"path": "tests/oss_pytests/test_placeholder.py",
"chars": 106,
"preview": "import pytest\n\n\n@pytest.mark.skip(reason=\"no other tests in this suite\")\ndef test_placeholder():\n pass\n"
},
{
"path": "tests/regenerate_reference.py",
"chars": 5274,
"preview": "\"\"\"Regenerate the numerical reference values the test suites compare against.\n\nRun from the repo root when a deliberate..."
},
{
"path": "tests/sdk/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "tests/sdk/forge_context_manager_test.py",
"chars": 1030,
"preview": "import pytest\n\nfrom esm.sdk import batch_executor, parallel_executor\nfrom esm.utils.forge_context_manager import ForgeBa..."
},
{
"path": "tests/utils/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "tests/utils/input_builder_test.py",
"chars": 4970,
"preview": "\"\"\"Round-trip tests for serialize/deserialize StructurePredictionInput.\"\"\"\n\nimport json\nfrom dataclasses import asdict..."
},
{
"path": "tests/utils/misc_test.py",
"chars": 1121,
"preview": "\"\"\"Tests for misc.py\"\"\"\n\nfrom esm.utils.misc import merge_annotations\nfrom esm.utils.types import FunctionAnnotation\n\n\nd..."
},
{
"path": "tests/utils/molecular_complex_test.py",
"chars": 11371,
"preview": "\"\"\"Tests for MolecularComplex CIF roundtrip: chain separation and entity info.\n\nVerifies that from_mmcif -> to_blob -> f..."
},
{
"path": "tests/utils/msa_test.py",
"chars": 7217,
"preview": "\"\"\"Tests for MSA.from_a3m deletion handling (a3m lowercase insertions).\"\"\"\n\nimport gzip\n\nimport numpy as np\n\nfrom esm.mo..."
},
{
"path": "tests/utils/sampling_test.py",
"chars": 1074,
"preview": "import pytest\nimport torch\n\nfrom esm.utils.sampling import sample_logits\n\n\ndef test_sample_logits():\n # batched input..."
}
]
// ... and 1 more files (download for full content)
About this extraction
This page contains the full source code of the evolutionaryscale/esm GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 202 files (17.4 MB), approximately 4.6M tokens. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.