### Install fair-esm with ESMFold Support

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Installs the fair-esm package with ESMFold capabilities, including necessary dependencies for OpenFold. Requires Python 3.9 or earlier and PyTorch. Also installs dllogger and a specific version of openfold.

```bash
pip install "fair-esm[esmfold]"
# OpenFold and its remaining dependency
pip install 'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'
```

--------------------------------

### Install Dependencies

Source: https://github.com/facebookresearch/esm/blob/main/examples/lm-design/free_generation.ipynb

Installs necessary dependencies for the project from a requirements file. Ensure 'additional_requirements.txt' is available in the same directory.

```python
# First install additional dependencies
!pip install -r additional_requirements.txt
```

--------------------------------

### Install fair-esm Package

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Installs the latest release of the fair-esm Python package or the bleeding-edge version directly from the GitHub repository.

```bash
pip install fair-esm  # latest release, OR:
pip install git+https://github.com/facebookresearch/esm.git  # bleeding edge, current repo main branch
```

--------------------------------

### Setup Colab Environment Dependencies (Python)

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/notebook_multichain.ipynb

Installs PyTorch, PyTorch Geometric, ESM, and Biotite for a Google Colab environment. It dynamically determines PyTorch and CUDA versions to install compatible PyTorch Geometric packages. Requires a GPU runtime for optimal performance.

```python
# Colab environment setup

# Install the correct version of Pytorch Geometric.
import torch

def format_pytorch_version(version):
  return version.split('+')[0]

TORCH_version = torch.__version__
TORCH = format_pytorch_version(TORCH_version)

def format_cuda_version(version):
  return 'cu' + version.replace('.', '')

CUDA_version = torch.version.cuda
CUDA = format_cuda_version(CUDA_version)

!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html
!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html
!pip install -q torch-cluster -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html
!pip install -q torch-spline-conv -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html
!pip install -q torch-geometric

# Install esm
!pip install -q git+https://github.com/facebookresearch/esm.git

# Install biotite
!pip install -q biotite
```

--------------------------------

### Install ESM Library

Source: https://github.com/facebookresearch/esm/blob/main/examples/esm_structural_dataset.ipynb

Installs the ESM library from its GitHub repository. This is a prerequisite for using the structural split dataset functionalities.

```python
# Install esm:
# !pip install git+https://github.com/facebookresearch/esm.git
```

--------------------------------

### Install ESM and Fetch Data (Colab)

Source: https://github.com/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb

These commands are for use in a Google Colab environment. They install the ESM package, download pre-computed protein embeddings, and fetch the FASTA file containing protein sequences and their effects.

```bash
# !pip install git+https://github.com/facebookresearch/esm.git
# !curl -O https://dl.fbaipublicfiles.com/fair-esm/examples/P62593_reprs.tar.gz
# !tar -xzf P62593_reprs.tar.gz
# !curl -O https://dl.fbaipublicfiles.com/fair-esm/examples/P62593.fasta
# !pwd
# !ls
```

--------------------------------

### Install py3Dmol for Visualization

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/notebook_multichain.ipynb

Installs the py3Dmol library, a Python package used for interactive 3D visualization of molecular structures, often used in bioinformatics and computational chemistry contexts.

```shell
!pip install -q py3Dmol
```

--------------------------------

### Verify PyTorch Geometric Installation (Python)

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/notebook_multichain.ipynb

Verifies the correct installation of PyTorch Geometric and its components. This is crucial for preventing runtime crashes due to version incompatibilities between PyTorch Geometric and PyTorch.

```python
## Verify that pytorch-geometric is correctly installed
import torch_geometric
import torch_sparse
from torch_geometric.nn import MessagePassing
```

--------------------------------

### Setting Up Conda Environment for Inverse Folding

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/README.md

Installs necessary packages for the inverse folding model, including PyTorch, PyTorch Geometric, and the ESM library. It's recommended to create a new conda environment to avoid CUDA compatibility issues.

```bash
conda create -n inverse python=3.9
conda activate inverse
conda install pytorch cudatoolkit=11.3 -c pytorch
conda install pyg -c pyg -c conda-forge
conda install pip
pip install biotite
pip install git+https://github.com/facebookresearch/esm.git
```

--------------------------------

### Load ESMFold Model for Protein Design

Source: https://github.com/facebookresearch/esm/blob/main/examples/protein-programming-language/tutorial.ipynb

Initializes and loads the ESMFold model, which is used to guide protein design. This step is crucial for utilizing ESMFold's capabilities in the protein programming language.

```python
%load_ext autoreload
%autoreload 2
from language import EsmFoldv1

folding_callback = EsmFoldv1()
folding_callback.load(device="cuda:0")
```

--------------------------------

### Download Files with aria2c

Source: https://github.com/facebookresearch/esm/blob/main/scripts/atlas/README.md

Provides an example of using the aria2c command-line tool to download files from a URL list. This is a recommended method for bulk downloading structures from the ESM Metagenomic Atlas.

```shell
aria2c --dir=/path/to/download/to --input-file=url-file-provided.txt
```

--------------------------------

### ESMFold v1 and v0 Model Loading (Python)

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Loads pre-trained ESMFold models for protein structure prediction using the ESM library. `esmfold_v1` is recommended for best performance, while `esmfold_v0` was used in prior research. These functions require the ESM library to be installed.

```python
import esm

# Recommended model for best performance
model_v1, alphabet_v1 = esm.pretrained.esmfold_v1()

# Model used in Lin et al. 2022 experiments
model_v0, alphabet_v0 = esm.pretrained.esmfold_v0()
```

--------------------------------

### Import and Compute Energy for Symmetric Binding Site Scaffolding

Source: https://github.com/facebookresearch/esm/blob/main/examples/protein-programming-language/tutorial.ipynb

This snippet demonstrates how to import a pre-defined protein design program for symmetric binding site scaffolding. It then computes and prints the energy terms for the designed protein using the provided folding callback and energy term functions.

```python
from programs.symmetric_binding import symmetric_binding_il10

program = symmetric_binding_il10()

sequence, residue_indices = program.get_sequence_and_set_residue_index_ranges()

folding_output = folding_callback.fold(sequence, residue_indices)

energy_terms = program.get_energy_term_functions()
for name, weight, energy_fn in energy_terms:
    print(f"{name} = {weight:.1f} * {energy_fn(folding_output):.2f}")
```

--------------------------------

### Download Model Checkpoints

Source: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb

This section details the process of downloading pre-trained model checkpoints for ESM. It shows the status of downloads, including whether they have already completed or were newly downloaded, and provides a summary of the download results with average speeds and paths. Dependencies include the torch library for handling model checkpoints.

```plaintext
04/12 20:46:44 [NOTICE] Downloading 1 item(s)

04/12 20:46:44 [NOTICE] GID#2bb7743048f09594 - Download has already completed: /root/.cache/torch/hub/checkpoints/esm1b_t33_650M_UR50S.pt

04/12 20:46:44 [NOTICE] Download complete: /root/.cache/torch/hub/checkpoints/esm1b_t33_650M_UR50S.pt

Download Results:
gid   |stat|avg speed  |path/URI
======+====+===========+=======================================================
2bb774|OK  |       0B/s|/root/.cache/torch/hub/checkpoints/esm1b_t33_650M_UR50S.pt

Status Legend:
(OK):download completed.

04/12 20:46:44 [NOTICE] Downloading 1 item(s)

04/12 20:46:45 [NOTICE] GID#9acf7bb38fe44e13 - Download has already completed: /root/.cache/torch/hub/checkpoints/esm1b_t33_650M_UR50S-contact-regression.pt

04/12 20:46:45 [NOTICE] Download complete: /root/.cache/torch/hub/checkpoints/esm1b_t33_650M_UR50S-contact-regression.pt

Download Results:
gid   |stat|avg speed  |path/URI
======+====+===========+=======================================================
9acf7b|OK  |       0B/s|/root/.cache/torch/hub/checkpoints/esm1b_t33_650M_UR50S-contact-regression.pt

Status Legend:
(OK):download completed.

04/12 20:46:45 [NOTICE] Downloading 1 item(s)
04/12 20:47:00 [NOTICE] Download complete: /root/.cache/torch/hub/checkpoints/esm_msa1b_t12_100M_UR50S.pt

Download Results:
gid   |stat|avg speed  |path/URI
======+====+===========+=======================================================
d9b729|OK  |    88MiB/s|/root/.cache/torch/hub/checkpoints/esm_msa1b_t12_100M_UR50S.pt

Status Legend:
(OK):download completed.

04/12 20:47:01 [NOTICE] Downloading 1 item(s)

04/12 20:47:01 [NOTICE] Download complete: /root/.cache/torch/hub/checkpoints/esm_msa1b_t12_100M_UR50S-contact-regression.pt

Download Results:
gid   |stat|avg speed  |path/URI
======+====+===========+=======================================================
fb2e95|OK  |        n/a|/root/.cache/torch/hub/checkpoints/esm_msa1b_t12_100M_UR50S-contact-regression.pt

Status Legend:
(OK):download completed.
```

--------------------------------

### Load and Initialize MSA Transformer Model

Source: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb

This code snippet loads and initializes the MSA Transformer model and its corresponding alphabet. It sets the model to evaluation mode and moves it to the GPU (CUDA) for efficient computation, and creates a batch converter.

```python
msa_transformer, msa_transformer_alphabet = esm.pretrained.esm_msa1b_t12_100M_UR50S()
msa_transformer = msa_transformer.eval().cuda()
msa_transformer_batch_converter = msa_transformer_alphabet.get_batch_converter()
```

--------------------------------

### Get Encoder Output for Single Chain

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/README.md

Extracts the encoder output from the ESM model, which serves as a structural representation. For an input of L amino acids, the output shape will be L x 512.

```python
rep = esm.inverse_folding.util.get_encoder_output(model, alphabet, coords)
```

--------------------------------

### Load ESM-2 15B Model with Fairscale FSDP CPU Offloading

Source: https://github.com/facebookresearch/esm/blob/main/examples/README.md

This Python script demonstrates an advanced method for loading the large ESM-2 15B model. It leverages Fairscale's Fully Sharded Data Parallel (FSDP) with CPU offloading capabilities, which is particularly useful for managing memory constraints when running large models on a single GPU.

```python
import torch
from fairscale.nn.model_parallel.initialize import initialize_model_parallel
from fairscale.optim.oss import OSS
from fairscale.nn.fsdp import FullyShardedDataParallel as FSDP

# Placeholder for model loading and FSDP configuration
# This is a conceptual representation of the script's functionality.
# Actual implementation details would be within the esm2_infer_fairscale_fsdp_cpu_offloading.py file.

def load_esm2_model_with_fsdp_cpu_offloading():
    # Initialize model parallel
    initialize_model_parallel(1, 0)

    # Load the ESM-2 15B model (details omitted for brevity)
    # model = ESM2_15B_Model()
    model = torch.nn.Module() # Dummy model for illustration

    # Configure FSDP with CPU offloading
    # fsdp_config = {
    #     "wrapper_cls": FSDP,
    #     "cpu_offload": True,
    #     "device": torch.device("cuda"),
    #     # Other FSDP parameters...
    # }
    # model = FSDP(model, **fsdp_config)

    print("ESM-2 15B model loaded with Fairscale FSDP and CPU offloading.")

if __name__ == "__main__":
    load_esm2_model_with_fsdp_cpu_offloading()

```

--------------------------------

### Compile Program with Constraints for Protein Design

Source: https://github.com/facebookresearch/esm/blob/main/examples/protein-programming-language/tutorial.ipynb

Assembles the complete protein design program by defining constraints (MaximizePTM, MaximizePLDDT, SymmetryRing) and their weights. It then compiles the program and prints a summary of the resulting energy function.

```python
from language import MaximizePTM, MaximizePLDDT, SymmetryRing

# Define the program.
program = ProgramNode(
    energy_function_terms=[MaximizePTM(), MaximizePLDDT(), SymmetryRing()],
    energy_function_weights=[1.0, 1.0, 1.0],
    children=[make_protomer_node() for _ in range(N)],
)

# Set up the program.
sequence, residue_indices = program.get_sequence_and_set_residue_index_ranges()
```

--------------------------------

### Get Encoder Output for Multi-chain Complex

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/README.md

Extracts the encoder output as a structural representation for a specific target chain within a multi-chain complex. This output can be used for downstream tasks requiring structural embeddings.

```python
rep = esm.inverse_folding.multichain_util.get_encoder_output_for_complex(
    model, alphabet, coords, target_chain_id
)
```

--------------------------------

### Run Simulated Annealing for Protein Design Optimization

Source: https://github.com/facebookresearch/esm/blob/main/examples/protein-programming-language/tutorial.ipynb

This Python code sets up and runs a simulated annealing optimization algorithm to design a protein sequence and structure. It utilizes the `run_simulated_annealing` function with parameters like initial temperature, annealing rate, and total steps. The final optimized sequence is then printed.

```python
from language import run_simulated_annealing

optimized_program = run_simulated_annealing(
    program=program,
    initial_temperature=1.0,
    annealing_rate=0.97,
    total_num_steps=10_000,
    folding_callback=folding_callback,
    display_progress=True,
)
print("Final sequence = {}".format(optimized_program.get_sequence_and_set_residue_index_ranges()[0]))
```

--------------------------------

### Load ESM Models via PyTorch Hub

Source: https://context7.com/facebookresearch/esm/llms.txt

Demonstrates how to load ESM-2, ESM-1b, and ESM-1v models directly from PyTorch Hub without cloning the repository. It also shows how to process sequences and extract embeddings.

```python
import torch

# Load ESM-2 model
model, alphabet = torch.hub.load("facebookresearch/esm:main", "esm2_t33_650M_UR50D")

# Load ESM-1b model
model, alphabet = torch.hub.load("facebookresearch/esm:main", "esm1b_t33_650M_UR50S")

# Load ESM-1v variant prediction model
model, alphabet = torch.hub.load("facebookresearch/esm:main", "esm1v_t33_650M_UR90S_1")

# Process sequences
batch_converter = alphabet.get_batch_converter()
model.eval()

data = [("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG")]
batch_labels, batch_strs, batch_tokens = batch_converter(data)

with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33])
    embeddings = results["representations"][33]
```

--------------------------------

### ESMFold CLI: CPU Offloading for Long Sequences (Python)

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Illustrates the use of the `--cpu-offload` flag in the ESMFold CLI. This option helps manage memory usage for predicting structures of long protein sequences by offloading some model parameters to CPU RAM.

```bash
esm-fold -i long_sequence.fasta -o output_directory --cpu-offload
```

--------------------------------

### Visualize Protein Chain using py3Dmol (Python)

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/notebook_multichain.ipynb

Defines a function `view_pdb` to visualize a specified chain from a PDB or CIF file using py3Dmol. It sets the cartoon representation with a specific color for the selected chain. Includes a fallback message if py3Dmol is not installed.

```python
try:
    import py3Dmol

    def view_pdb(fpath, chain_ids, colors=['red','blue','green']):
        with open(fpath) as ifile:
            system = "".join([x for x in ifile])

        view = py3Dmol.view(width=600, height=400)
        view.addModelsAsFrames(system)
        for chain_id, color in zip(chain_ids, colors):
            view.setStyle({'model': -1, 'chain': chain_id}, {"cartoon": {'color': color}})
        view.zoomTo()
        view.show()

except ImportError:
    def view_pdb(fpath, chain_id):
        print("Install py3Dmol to visualize, or use pymol")
```

--------------------------------

### Visualize Protein Structure using py3Dmol

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/notebook.ipynb

This Python function, `view_pdb`, visualizes a protein structure from a given file path and chain ID using the py3Dmol library. It displays the structure with a spectrum-colored cartoon representation. If py3Dmol is not installed, it prints a message suggesting alternatives.

```python
try:
    import py3Dmol

    def view_pdb(fpath, chain_id):
        with open(fpath) as ifile:
            system = "".join([x for x in ifile])

        view = py3Dmol.view(width=600, height=400)
        view.addModelsAsFrames(system)
        view.setStyle({'model': -1, 'chain': chain_id}, {"cartoon": {'color': 'spectrum'}})
        view.zoomTo()
        view.show()

except ImportError:
    def view_pdb(fpath, chain_id):
        print("Install py3Dmol to visualize, or use pymol")
```

--------------------------------

### ESMFold Command Line Interface Usage

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Demonstrates the command-line interface for performing bulk protein structure prediction using ESMFold. It takes a FASTA file as input and outputs PDB files. Key arguments control the number of recycles, batching, chunk size, and CPU usage.

```bash
esm-fold -i input.fasta -o output_directory --num-recycles 4 --max-tokens-per-batch 512 --chunk-size 128 --cpu-offload
```

--------------------------------

### Download PDB/CIF File using wget

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/notebook_multichain.ipynb

This snippet downloads a PDB or CIF file from a given URL and saves it to a specified directory. It's useful for fetching protein structure data for local analysis.

```shell
!wget https://files.rcsb.org/download/5YH2.cif -P data/    # save this to the data folder in colab
```

--------------------------------

### Sample Sequence Designs with ESM-IF1 (Bash)

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Samples sequence designs for a given protein structure in PDB or mmCIF format using the ESM-IF1 model. The script outputs sequences in FASTA format. The temperature parameter controls sequence diversity versus native recovery.

```bash
python examples/inverse_folding/sample_sequences.py examples/inverse_folding/data/5YH2.pdb \
  --chain C --temperature 1 --num-samples 3 --outpath examples/inverse_folding/output/sampled_sequences.fasta
```

--------------------------------

### CPU Offloading for Large Models (Python)

Source: https://context7.com/facebookresearch/esm/llms.txt

Illustrates how to load and run inference with very large ESM models (e.g., 15B parameters) using PyTorch's FullyShardedDataParallel (FSDP) with CPU offloading to manage memory. Requires PyTorch and ESM.

```python
import torch
import esm
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, CPUOffload
from torch.distributed.fsdp.wrap import enable_wrap, wrap

# Initialize distributed process (required for FSDP)
torch.distributed.init_process_group(
    backend="nccl",
    init_method="tcp://localhost:9999",
    world_size=1,
    rank=0
)

# Load large model
model, alphabet = esm.pretrained.esm2_t48_15B_UR50D()
model.eval()

# Wrap model with FSDP and CPU offloading
wrapper_kwargs = dict(cpu_offload=CPUOffload(offload_params=True))
with enable_wrap(wrapper_cls=FSDP, **wrapper_kwargs):
    for layer_name, layer in model.layers.named_children():
        wrapped_layer = wrap(layer)
        setattr(model.layers, layer_name, wrapped_layer)
    model = wrap(model)

model = model.cuda()

# Run inference
batch_converter = alphabet.get_batch_converter()
data = [("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG")]
batch_labels, batch_strs, batch_tokens = batch_converter(data)

with torch.no_grad():
    batch_tokens = batch_tokens.cuda()
    results = model(batch_tokens, repr_layers=[48])
    embeddings = results["representations"][48]
    print(f"Successfully computed embeddings with shape: {embeddings.shape}")
```

--------------------------------

### Load and Initialize ESM-2 Model

Source: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb

This code snippet loads and initializes the ESM-2 model and its associated alphabet. It sets the model to evaluation mode and moves it to the GPU (CUDA) for faster processing. It then creates a batch converter for preparing input sequences.

```python
esm2, esm2_alphabet = esm.pretrained.esm2_t33_650M_UR50D()
esm2 = esm2.eval().cuda()
esm2_batch_converter = esm2_alphabet.get_batch_converter()
```

--------------------------------

### Load ESM-2 Model via PyTorch Hub

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Loads a pretrained ESM-2 model and its corresponding alphabet directly from PyTorch Hub without needing to clone the repository. This is a convenient way to access state-of-the-art models.

```python
import torch
model, alphabet = torch.hub.load("facebookresearch/esm:main", "esm2_t33_650M_UR50D")
```

--------------------------------

### Initialize and Configure Designer

Source: https://github.com/facebookresearch/esm/blob/main/examples/lm-design/free_generation.ipynb

Imports required libraries and sets up the configuration for the Free Generation task using Hydra. This includes defining the task type, random seed, and desired sequence length.

```python
# Imports
import os
import time
import hydra
import py3Dmol
from lm_design import Designer

# Params
seed = 0  # Use different seeds to get different sequence results
free_generation_length = 100
TASK = "free_generation"

# Load hydra config from config.yaml
with hydra.initialize_config_module(config_module="conf"):
    cfg = hydra.compose(
        config_name="config", 
        overrides=[
            f"task={TASK}", 
            f"seed={seed}", 
            f"free_generation_length={free_generation_length}", 
            # 'tasks.free_generation.num_iter=100'  # DEBUG - use a smaller number of iterations
        ])

# Create a designer from configuration
des = Designer(cfg)
```

--------------------------------

### Import ESM Project Dependencies

Source: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb

This snippet imports essential libraries for the ESM project, including those for typing, itertools, file system operations, string manipulation, numerical computation (numpy), PyTorch, scientific computing (scipy), plotting (matplotlib), bioinformatics (Bio, biotite), progress bars (tqdm), and data manipulation (pandas). It also sets PyTorch's gradient calculation to be disabled by default.

```python
from typing import List, Tuple, Optional, Dict, NamedTuple, Union, Callable
import itertools
import os
import string
from pathlib import Path

import numpy as np
import torch
from scipy.spatial.distance import squareform, pdist, cdist
import matplotlib.pyplot as plt
import matplotlib as mpl
from Bio import SeqIO
import biotite.structure as bs
from biotite.structure.io.pdbx import PDBxFile, get_structure
from biotite.database import rcsb
from tqdm import tqdm
import pandas as pd

import esm

torch.set_grad_enabled(False)
```

--------------------------------

### Download PDB/CIF File using wget

Source: https://github.com/facebookresearch/esm/blob/main/examples/inverse_folding/notebook.ipynb

This snippet downloads a PDB or CIF file from a given URL and saves it to a specified directory. It's useful for obtaining protein structure data for further analysis.

```python
#save this to the data folder in colab
!wget https://files.rcsb.org/download/5YH2.cif -P data/
```

--------------------------------

### Batch Processing with FastaBatchedDataset

Source: https://context7.com/facebookresearch/esm/llms.txt

Explains how to efficiently process large FASTA files using `FastaBatchedDataset` and `DataLoader` with batching based on token counts. This method is suitable for handling large datasets and utilizes GPU acceleration if available.

```python
import torch
from esm import Alphabet, FastaBatchedDataset, pretrained

# Load model
model, alphabet = pretrained.esm2_t33_650M_UR50D()
model.eval()
if torch.cuda.is_available():
    model = model.cuda()

# Load FASTA file
dataset = FastaBatchedDataset.from_file("sequences.fasta")
batches = dataset.get_batch_indices(toks_per_batch=4096, extra_toks_per_seq=1)
data_loader = torch.utils.data.DataLoader(
    dataset,
    collate_fn=alphabet.get_batch_converter(),
    batch_sampler=batches
)

print(f"Processing {len(dataset)} sequences in {len(batches)} batches")

# Process batches
all_representations = {}
with torch.no_grad():
    for batch_idx, (labels, strs, toks) in enumerate(data_loader):
        if torch.cuda.is_available():
            toks = toks.cuda()

        out = model(toks, repr_layers=[33])
        representations = out["representations"][33].cpu()

        for i, label in enumerate(labels):
            seq_len = len(strs[i])
            all_representations[label] = representations[i, 1:seq_len+1].mean(0)

        print(f"Processed batch {batch_idx + 1}/{len(batches)}")
```

--------------------------------

### Initialize LM Design Parameters and Imports

Source: https://github.com/facebookresearch/esm/blob/main/examples/lm-design/fixed_backbone.ipynb

Imports essential libraries and defines initial parameters for the protein design task. This includes setting the PDB file path, random seed, and the specific design task type ('fixedbb'). These parameters configure the LM for the desired design objective.

```python
# Imports
import os
import time
import hydra
import py3Dmol
from lm_design import Designer

# Params
pdb_fn = os.getcwd() + '/2N2U.pdb'
seed = 0  # Use different seeds to get different sequence designs for the same structure
TASK = "fixedbb"

```

--------------------------------

### Prepare Model and Parameter Lists for Grid Search in Python

Source: https://github.com/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb

This code prepares lists of regression model classes and their corresponding parameter grids. These lists are used in conjunction with the grid search process to efficiently iterate through different models and their hyperparameters.

```python
cls_list = [KNeighborsRegressor, SVR, RandomForestRegressor]
param_grid_list = [knn_grid, svm_grid, rfr_grid]
```

--------------------------------

### Unsupervised Contact Prediction with ESM

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Python code demonstrating how to perform unsupervised contact prediction using ESM models. This involves calling the `predict_contacts` method or using the model with `return_contacts=True`.

```python
# Call model.predict_contacts(tokens)
# or
# model(tokens, return_contacts=True)
```

--------------------------------

### Load Configuration for Fixed Backbone Design

Source: https://github.com/facebookresearch/esm/blob/main/examples/lm-design/fixed_backbone.ipynb

Loads the Hydra configuration, specifying the 'fixedbb' task, a given seed, and the input PDB file. This step allows for flexible and organized management of design parameters, enabling reproducibility and experimentation by overriding default settings.

```python
# Load hydra config from config.yaml
with hydra.initialize_config_module(config_module="conf"):
    cfg = hydra.compose(
        config_name="config", 
        overrides=[
            f"task={TASK}", 
            f"seed={seed}", 
            f"pdb_fn={pdb_fn}", 
            # 'tasks.fixedbb.num_iter=100'  # DEBUG - use a smaller number of iterations
        ])
```

--------------------------------

### Initialize Grid Search Parameters for Regression Models in Python

Source: https://github.com/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb

This snippet initializes parameter grids for three different regression models: KNeighborsRegressor, SVR, and RandomForestRegressor. These grids define the hyperparameter values to be explored during the grid search process. PCA-projected features are used for faster training.

```python
knn_grid = [
    {
        'model': [KNeighborsRegressor()],
        'model__n_neighbors': [5, 10],
        'model__weights': ['uniform', 'distance'],
        'model__algorithm': ['ball_tree', 'kd_tree', 'brute'],
        'model__leaf_size' : [15, 30],
        'model__p' : [1, 2],
    }
    ]

svm_grid = [
    {
        'model': [SVR()],
        'model__C' : [0.1, 1.0, 10.0],
        'model__kernel' : ['linear', 'poly', 'rbf', 'sigmoid'],
        'model__degree' : [3],
        'model__gamma': ['scale'],
    }
]

rfr_grid = [
    {
        'model': [RandomForestRegressor()],
        'model__n_estimators' : [20],
        'model__criterion' : ['squared_error', 'absolute_error'],
        'model__max_features': ['sqrt', 'log2'],
        'model__min_samples_split' : [5, 10],
        'model__min_samples_leaf': [1, 4]
    }
]
```

--------------------------------

### Visualize Wild-Type Protein Structure

Source: https://github.com/facebookresearch/esm/blob/main/examples/lm-design/fixed_backbone.ipynb

Visualizes the 3D structure of the original wild-type protein using the 'py3Dmol' library. The PDB file specified during initialization is loaded, styled similarly to the designed structure (spectrum cartoon), and zoomed. This provides a reference for comparison with the designed protein.

```python
# Visualize wild type structure
wt_struct_file = pdb_fn
view = py3Dmol.view(width=800, height=800)
view.addModel(open(wt_struct_file).read(), 'pdb')
view.setStyle({'cartoon': {'color': 'spectrum'}})
view.zoomTo()
view.show()
```

--------------------------------

### Load and Use ESM-2 Model for Protein Representation

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Loads an ESM-2 model and prepares it for inference. It demonstrates data preparation, including handling masked tokens, and extracting per-residue and per-sequence representations. It also shows how to visualize contact predictions.

```python
import torch
import esm

# Load ESM-2 model
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()
model.eval()  # disables dropout for deterministic results

# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
data = [
    ("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
    ("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
    ("protein2 with mask","KALTARQQEVFDLIRD<mask>ISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
    ("protein3",  "K A <mask> I S Q")
]
batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)

# Extract per-residue representations (on CPU)
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
token_representations = results["representations"][33]

# Generate per-sequence representations via averaging
# NOTE: token 0 is always a beginning-of-sequence token, so the first residue is token 1.
sequence_representations = []
for i, tokens_len in enumerate(batch_lens):
    sequence_representations.append(token_representations[i, 1 : tokens_len - 1].mean(0))

# Look at the unsupervised self-attention map contact predictions
import matplotlib.pyplot as plt
for (_, seq), tokens_len, attention_contacts in zip(data, batch_lens, results["contacts"]):
    plt.matshow(attention_contacts[: tokens_len, : tokens_len])
    plt.title(seq)
    plt.show()
```

--------------------------------

### Predict Protein Structures in Bulk via Command Line (Bash)

Source: https://context7.com/facebookresearch/esm/llms.txt

Utilizes the `esm-fold` command-line interface to predict protein structures for multiple sequences provided in a FASTA file. Offers options for custom settings like number of recycles, batch size, chunk size, CPU-only execution, and CPU offloading for memory management.

```bash
# Basic structure prediction
esm-fold -i input_sequences.fasta -o output_pdb_dir/

# With custom settings for memory optimization
esm-fold -i sequences.fasta -o structures/ \
  --num-recycles 4 \
  --max-tokens-per-batch 1024 \
  --chunk-size 128

# CPU-only mode (for machines without GPU)
esm-fold -i sequences.fasta -o structures/ --cpu-only

# CPU offloading for very large models or long sequences
esm-fold -i sequences.fasta -o structures/ --cpu-offload
```

--------------------------------

### Load ESM-2 Language Model and Extract Representations (Python)

Source: https://context7.com/facebookresearch/esm/llms.txt

Demonstrates loading the ESM-2 language model and extracting per-residue and per-sequence representations for protein sequences. Requires the `torch` and `esm` libraries. Input is a list of tuples containing sequence names and amino acid sequences. Output includes token representations and averaged sequence representations.

```python
import torch
import esm

# Load ESM-2 model (650M parameters)
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()
model.eval()  # disable dropout for deterministic results

# Prepare protein sequences
data = [
    ("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
    ("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
    ("protein_with_mask", "KALTARQQEVFDLIRD<mask>ISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
]

# Convert to tokens
batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)

# Extract per-residue representations
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)

token_representations = results["representations"][33]

# Generate per-sequence representations via averaging
sequence_representations = []
for i, tokens_len in enumerate(batch_lens):
    sequence_representations.append(
        token_representations[i, 1 : tokens_len - 1].mean(0)
    )

print(f"Sequence representation shape: {sequence_representations[0].shape}")
# Output: Sequence representation shape: torch.Size([1280])
```

--------------------------------

### Import Necessary Libraries

Source: https://github.com/facebookresearch/esm/blob/main/examples/esm_structural_dataset.ipynb

Imports essential Python libraries for data manipulation, visualization, and the ESM protein analysis tools. This includes `Path`, `seaborn`, `matplotlib`, `os`, `torch`, and components from the `esm` library.

```python
from pathlib import Path
import seaborn as sns
import matplotlib.pyplot as plt
import os
import torch

import esm
from esm.data import ESMStructuralSplitDataset
```

--------------------------------

### Examine a Single Protein Sample

Source: https://github.com/facebookresearch/esm/blob/main/examples/esm_structural_dataset.ipynb

Loads a specific sample from the training dataset and prints its available keys and content. Each sample is a dictionary containing 'seq' (sequence), 'ssp' (secondary structure prediction), 'dist' (distance map), and 'coords' (3D coordinates).

```python
esm_structural_train = ESMStructuralSplitDataset(
    split_level='superfamily', 
    cv_partition='4', 
    split='train', 
    root_path = os.path.expanduser('~/.cache/torch/data/esm'),
    download=True
)
ele = esm_structural_train[0]
print(ele.keys()) 
print()
print('sequence', ele['seq'])
print()
print('SSP labels', ele['ssp'])
```

--------------------------------

### Load Metadata with Pandas

Source: https://github.com/facebookresearch/esm/blob/main/scripts/atlas/README.md

Demonstrates how to load the Atlas metadata into a pandas DataFrame using the parquet file. This allows for easy manipulation and querying of the metadata, which contains information about each protein structure.

```python
import pandas as pd

df = pd.read_parquet('metadata.parquet')
```

--------------------------------

### Run Fixed Backbone Design Process

Source: https://github.com/facebookresearch/esm/blob/main/examples/lm-design/fixed_backbone.ipynb

Executes the protein sequence design process using the initialized Designer. It measures the execution time and prints the total duration in hours. This step iteratively refines the sequence to match the target structure.

```python
# Run the designer
start_time = time.time()
des.run_from_cfg()
print("finished after %s hours", (time.time() - start_time) / 3600)
```

--------------------------------

### Execute Free Generation Design

Source: https://github.com/facebookresearch/esm/blob/main/examples/lm-design/free_generation.ipynb

Runs the Free Generation design process using the configured Designer object. It measures and prints the execution time in hours.

```python
# Run the designer
start_time = time.time()
des.run_from_cfg()    
print("finished after %s hours", (time.time() - start_time) / 3600)
```

--------------------------------

### ESMFold CLI: Disabling Batching (Python)

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Shows how to disable the default batching behavior in the ESMFold CLI. Setting `--max-tokens-per-batch=0` processes sequences individually, which can be useful for debugging or specific use cases, though it may reduce speed.

```bash
esm-fold -i input.fasta -o output_directory --max-tokens-per-batch=0
```

--------------------------------

### Generate MSA with HHblits

Source: https://github.com/facebookresearch/esm/blob/main/examples/README.md

This command demonstrates how to generate Multiple Sequence Alignments (MSAs) using the HHblits tool. It takes a FASTA file as input and outputs an MSA in the a3m format. The command specifies the number of iterations and the database to use for searching.

```bash
hhblits -i UniRef50_$id.fas -oa3m UniRef50_$id.a3m -n 3 -d /uniclust30_2017_10/uniclust30_2017_10
```

--------------------------------

### ESMFold CLI: Handling Multimers in FASTA

Source: https://github.com/facebookresearch/esm/blob/main/README.md

Explains how to input multimeric protein structures into the FASTA file for prediction with ESMFold CLI. Chains within a single multimer should be represented as a single sequence, with chain identifiers separated by a colon character.

```plaintext
>MultimerSequence
SEQUENCE_CHAIN_A:SEQUENCE_CHAIN_B:SEQUENCE_CHAIN_C
```

--------------------------------

### ESM-2 Contact Prediction with PyTorch

Source: https://github.com/facebookresearch/esm/blob/main/examples/esm_structural_dataset.ipynb

This snippet shows how to load an ESM-2 model, prepare protein sequence data, and generate contact predictions using self-attention maps. It requires PyTorch and the ESM library. The output is a contact map visualization.

```python
import esm
import torch
import matplotlib.pyplot as plt

# Load the ESM-2 model and alphabet
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()

# Prepare dummy data (replace with your actual data)
data = [("protein_1", "ACDEFGHIKLMNPQRSTVWY"), ("protein_2", "MKLTVQYEDRFGHIJ")]

# Convert data to batch
batch_labels, batch_strs, batch_tokens = batch_converter(data)

# Calculate batch lengths
batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)

# Generate contact predictions
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)

# Visualize contact maps
for (_, seq), tokens_len, attention_contacts in zip(data, batch_lens, results["contacts"]):
    plt.matshow(attention_contacts[: tokens_len, : tokens_len])
    plt.title(f"Contact Map for {seq}")
    plt.show()
```

--------------------------------

### Cite Variant Prediction Benchmark Data (LaTeX)

Source: https://github.com/facebookresearch/esm/blob/main/README.md

This snippet provides a LaTeX command to cite all data used in the variant prediction benchmark from Meier et al. (2021). It is useful for researchers who want to include all relevant citations in their work.

```tex
\nocite{wrenbeck2017deep,klesmith2015comprehensive,haddox2018mapping,romero2015dissecting,firnberg2014comprehensive,deng2012deep,stiffler2015evolvability,jacquier2013capturing,findlay2018comprehensive,mclaughlin2012spatial,kitzman2015massively,doud2016accurate,pokusaeva2019experimental,mishra2016systematic,kelsic2016rna,melnikov2014comprehensive,brenan2016phenotypic,rockah2015systematic,wu2015functional,aakre2015evolving,qi2014quantitative,matreyek2018multiplex,bandaru2017deconstruction,roscoe2013analyses,roscoe2014systematic,mavor2016determination,chan2017correlation,melamed2013deep,starita2013activity,araya2012fundamental}
```

--------------------------------

### Import Scikit-learn ML Modules (Python)

Source: https://github.com/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb

Imports various modules from the scikit-learn library, including tools for model selection (GridSearchCV, train_test_split), dimensionality reduction (PCA), and different machine learning models like classifiers and regressors.

```python
import scipy
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.decomposition import PCA
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.svm import SVC, SVR
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.naive_bayes import GaussianNB
from sklearn.linear_model import LogisticRegression, SGDRegressor
from sklearn.pipeline import Pipeline
```

--------------------------------

### Load Protein Data

Source: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb

This code snippet defines the protein structures to be analyzed, retrieves the data using PDB IDs, and extracts contact information and multiple sequence alignments (MSAs) from the retrieved structures. It prepares the data for subsequent processing by the ESM models.

```python
PDB_IDS = ["1a3a", "5ahw", "1xcr"]

structures = {
    name.lower(): get_structure(PDBxFile.read(rcsb.fetch(name, "cif")))[0]
    for name in PDB_IDS
}

contacts = {
    name: contacts_from_pdb(structure, chain="A") 
    for name, structure in structures.items()
}

msas = {
    name: read_msa(f"data/{name.lower()}_1_A.a3m")
    for name in PDB_IDS
}

sequences = {
    name: msa[0] for name, msa in msas.items()
}
```

--------------------------------

### Run MSA Transformer Contact Prediction and Evaluate

Source: https://github.com/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb

This code snippet processes the loaded MSAs, converts them, and uses the MSA Transformer to predict contacts. It then evaluates these predictions against the known contacts, computes metrics, and organizes the results in a Pandas DataFrame for display.

```python
msa_transformer_predictions = {}
msa_transformer_results = []
for name, inputs in msas.items():
    inputs = greedy_select(inputs, num_seqs=128) # can change this to pass more/fewer sequences
    msa_transformer_batch_labels, msa_transformer_batch_strs, msa_transformer_batch_tokens = msa_transformer_batch_converter([inputs])
    msa_transformer_batch_tokens = msa_transformer_batch_tokens.to(next(msa_transformer.parameters()).device)
    msa_transformer_predictions[name] = msa_transformer.predict_contacts(msa_transformer_batch_tokens)[0].cpu()
    metrics = {"id": name, "model": "MSA Transformer (Unsupervised)"}
    metrics.update(evaluate_prediction(msa_transformer_predictions[name], contacts[name]))
    msa_transformer_results.append(metrics)
msa_transformer_results = pd.DataFrame(msa_transformer_results)
display(msa_transformer_results)
```