Cerebra-Seq

Cerebra-Seq serves as the structure predictor for Cerebra-Epistasis. It uses features from ESM3 (1.4B) and ESMC (600M) to predict protein structures directly from amino acid sequences. The resulting node/edge embeddings, atom14 coordinates, and atom masks serve as inputs to the downstream components of Cerebra-Epistasis.

Requirements

With compatible PyTorch and Transformers installed, add einops for ESM3:

pip install einops

Example Usage

Single-sequence inference on GPU or CPU. Use checkpoint="model1" through "model5" to select a checkpoint. trust_remote_code=True loads the repositories' custom model code.

To reproduce the Cerebra-Epistasis results reported in the paper, use checkpoint="model1", as in the example below.

import torch
from transformers import AutoModel, AutoTokenizer

sequence = "MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"
device = "cuda"
L = len(sequence)


@torch.inference_mode()
def esm_features(repo, sequence, esm3=False):
    model = AutoModel.from_pretrained(
        repo, trust_remote_code=True,
        attn_implementation="sdpa",
    ).eval().to(device)

    if esm3:
        inputs = model.tokenize_sequences([sequence], device=device)
    else:
        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
        inputs = tokenizer(sequence, return_tensors="pt").to(device)

    return model(**inputs, return_dict=True).last_hidden_state[:, 1:-1].float().cpu()


esm3 = esm_features("Synthyra/ESM3_small", sequence, esm3=True)
esmc = esm_features("biohub/ESMC-600M-hf", sequence)

# Cerebra-Seq uses HHblits residue indices, not ESM tokenizer IDs.
aa_ids = dict(zip("ACDEFGHIKLMNPQRSTVWY", range(20)))
aa_ids.update(B=2, Z=3, U=1, X=20, J=20, O=20)
feats = {
    "target_feat": torch.tensor([[aa_ids.get(aa, 20) for aa in sequence]]),
    "residue_index": torch.arange(1, L + 1)[None],
    "seq_mask": torch.ones(1, L),
    "X1D_esm_c": esmc,
    "X1D_esm3": esm3,
}
feats = {key: value.to(device) for key, value in feats.items()}

# Select anchors according to sequence length; keep them on CPU.
n = next((count for limit, count in [(96, 18), (224, 24), (656, 32), (900, 24)] if L <= limit), 24)
anchors = torch.tensor([int(i * (L - 8) / n) + 5 for i in range(n)]).clamp(2, L - 2)
if n + 8 > L:
    anchors = torch.arange(2, L - 2)  # Keep anchors on CPU.

model = AutoModel.from_pretrained(
    "GongLab-THU/Cerebra-Seq",
    trust_remote_code=True, checkpoint="model1", device=device,
).eval().float()

prevs = [None, None, None]
with torch.inference_mode():
    for cycle in range(4):
        m, z, x, outputs = model(
            feats, prevs, anchors, _recycle=cycle < 3, return_aux=cycle == 3,
            return_dist=False, return_pae=False, reduce_plddt=True,
            conf_anchor_chunk_size=24, compress_recycle=True, keep_structure_all=False,
        )
        if cycle < 3:
            prevs = [m.detach(), z.detach(), x.detach()]

# Extract the four output tensors as detached CPU FP32.
features = model.extract_features(outputs, anchors, feats)

The extracted features are compatible with Cerebra-Epistasis and contain the following tensors:

node_embedding: (L, 256)
edge_embedding: (128, L, L)
atom14_coords:  (L, 14, 3)
atom14_masks:   (L, 14)

L is the sequence length. Atom ordering is residue-specific; atom14_masks is 1 for existing atoms and 0 for unused slots.

Optional: save a PDB

Append this to save the predicted structure:

from pathlib import Path

Path("prediction.pdb").write_text(
    model.to_pdb(features, feats, outputs=outputs), encoding="utf-8"
)

The unrelaxed PDB uses chain A, residue numbers 1..L, and pLDDT in the B-factor column. PDB export requires no additional dependencies.

License

The original Cerebra-Seq code and model weights are released under the MIT License. Third-party code retains its original licenses and copyright notices. The separately downloaded ESM3 and ESMC models remain subject to their respective license terms.