SaProt
# Model Introduction SaProt (Protein Language Modeling with Structure-aware Vocabulary) is a protein language model that jointly models protein amino acid sequences and structural information. Its central idea is to combine amino acids (AA) with the 3Di structural alphabet generated by Foldseek into structure-aware tokens, allowing the model to learn representations from both protein sequences and structural context. SaProt can be used for protein representation extraction, zero-shot mutation effect prediction, protein inverse folding, and downstream task fine-tuning. Paper: > **SaProt: Protein Language Modeling with Structure-aware Vocabulary** > ICLR 2024 Spotlight > Follow-up work was published in Nature Biotechnology (2025) # Model Description SaProt models proteins using a structure-aware vocabulary formed by combining amino acids (AA) with the Foldseek 3Di structural alphabet. For example, a structure-aware sequence can be represented as: ```text M#EvVpQpL#VyQdYaKv ``` Every two characters form a structure-aware token: the first character represents the amino acid, and the second represents the corresponding 3Di structural state. `#` can be used to mask low-confidence structural regions. The official release provides pretrained models at multiple scales: | Model | Parameter scale | Training data | | --- | ---: | --- | | `SaProt_35M_AF2` | 35M | 40M AF2 structures | | `SaProt_650M_PDB` | 650M | 40M AF2 structures + 60K PDB structures | | `SaProt_650M_AF2` | 650M | 40M AF2 structures | | `SaProt_1.3B_AF2` | 1.3B | 40M AF2 structures | | `SaProt_1.3B_AFDB_OMG_NCBI` | 1.3B | AFDB + OMG_prot50 + NCBI | For the 35M and 650M SaProt models, the official recommendation is to use SA-token inputs containing structural information for the best results. The 1.3B version can handle both structure-aware sequences and amino-acid-only sequences relatively well. # Use Cases | Use case | Description | | --- | --- | | Protein representation extraction | Extract residue-level or protein-level embeddings | | Zero-shot mutation effect prediction | Evaluate single or multiple mutations directly without task-specific fine-tuning | | Structure-aware protein modeling | Jointly use amino acid and 3Di structural tokens | | Protein inverse folding | Design sequences from structural information | | Downstream task fine-tuning | Apply to tasks such as EC, GO, stability, PPI, Contact, and DeepLoc | # Usage ## 1. Using OneCode Experience intelligent one-click AI4S programming in the OneCode online environment: [Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) ## 2. Manual Installation and Usage **Hardware Requirements** - SaProt supports inference on CPUs and GPUs/DCUs. - The 35M model can be used for lightweight testing; GPUs/DCUs are recommended for the 650M and 1.3B models. - Batch embedding, mutation scanning, pretraining, and fine-tuning substantially increase GPU memory and host memory requirements. ### Set Up the Runtime Environment #### DCU Environment ```bash # Activate DTK and CONDA first conda create -n onescience311 python=3.11 -y conda activate onescience311 pip install onescience[bio] \ -i http://mirrors.onescience.ai:3141/pypi/simple/ \ --trusted-host mirrors.onescience.ai ``` #### Environment Notes - If you encounter missing dependencies or version incompatibilities during execution, refer to the dependency versions specified in `requirements.txt` and install or adjust the relevant dependencies as needed. ### Prepare Models and Data #### 1) SaProt Model Weights The official models are primarily released on Hugging Face: ```text SaProt_35M_AF2 https://huggingface.co/westlake-repl/SaProt_35M_AF2 SaProt_650M_PDB https://huggingface.co/westlake-repl/SaProt_650M_PDB SaProt_650M_AF2 https://huggingface.co/westlake-repl/SaProt_650M_AF2 SaProt_1.3B_AF2 https://huggingface.co/westlake-repl/SaProt_1.3B_AF2 SaProt_1.3B_AFDB_OMG_NCBI https://huggingface.co/westlake-repl/SaProt_1.3B_AFDB_OMG_NCBI ``` For example, download `SaProt_650M_AF2` in advance for offline use: ```bash huggingface-cli download \ westlake-repl/SaProt_650M_AF2 \ --local-dir ./weight/PLMs/SaProt_650M_AF2 ``` It is recommended to download all model weights under `weight/PLMs/`. The current SaProt configuration reads from: ```text weight/PLMs/SaProt_650M_AF2 ``` To run the ESM2 comparison experiment, also prepare: ```bash huggingface-cli download \ facebook/esm2_t33_650M_UR50D \ --local-dir ./weight/PLMs/esm2_t33_650M_UR50D ``` The corresponding configuration reads from: ```text weight/PLMs/esm2_t33_650M_UR50D ``` #### 2) Foldseek SaProt structure-aware inputs require PDB/CIF structures to be encoded as Foldseek 3Di sequences first. The official README provides the following download link: ```text https://drive.google.com/file/d/1B_9t3n_nlj8Y3Kpc_mMjtMdY0OPYa7Re/view ``` You can also use the official Foldseek Linux prebuilt package or a Foldseek installation already available on the system or platform. For this adapted version, place Foldseek at: ```text SaProt/ └── scripts/ └── bin/ └── foldseek ``` Then grant it execute permission: ```bash chmod +x scripts/bin/foldseek ``` The Foldseek path used by the current configuration is: ```text scripts/bin/foldseek ``` #### 3) Downstream Task Datasets The official downstream task datasets are available at: ```text https://drive.google.com/drive/folders/11dNGqPYfLE3M-Mbh4U7IQpuHxJpuRr4g?usp=sharing ``` For this adapted version, extract the downstream task data to: ```text scripts/LMDB/ ``` Typical paths used by the configuration include: ```text scripts/LMDB/Thermostability/foldseek/train scripts/LMDB/Thermostability/foldseek/valid scripts/LMDB/Thermostability/foldseek/test scripts/LMDB/ProteinGym/substitutions scripts/LMDB/ClinVar ``` #### 4) Pretraining Dataset To pretrain SaProt from scratch or continue pretraining, prepare the official pretraining data: ```text westlake-repl/AF2_UniRef50 https://huggingface.co/datasets/westlake-repl/AF2_UniRef50 ``` The official pretraining configuration uses LMDB data directories such as: ```text scripts/LMDB/AF2_Uniref50/foldseek/train scripts/LMDB/AF2_Uniref50/foldseek/valid ``` The pretraining dataset is large and is only needed when pretraining from scratch or continuing pretraining. ## 3. Quick Start ### Download the Model Package ```bash hf download OneScience-Group/SaProt \ --local-dir ./SaProt cd SaProt ``` ### Quick Verification Check the dependencies: ```bash python - <<'PY' import torch import transformers import esm import pytorch_lightning as pl print("torch:", torch.__version__) print("transformers:", transformers.__version__) print("pytorch_lightning:", pl.__version__) print("SaProt dependencies OK") PY ``` Check Foldseek: ```bash ./scripts/bin/foldseek version ``` Test model loading: ```bash python - <<'PY' from transformers import EsmTokenizer, EsmForMaskedLM model_path = "./weight/PLMs/SaProt_650M_AF2" tokenizer = EsmTokenizer.from_pretrained(model_path) model = EsmForMaskedLM.from_pretrained(model_path) print("SaProt load OK") PY ``` # Example Data The official repository provides: ```text scripts/example/8ac8.cif ``` This can be used to demonstrate conversion from a protein structure to a structure-aware sequence. Your own structure input can be: ```text *.pdb *.cif ``` If you already have a Foldseek-encoded structure-aware sequence, you can pass it directly to SaProt without processing the structure file again. # Inference Examples ## Load SaProt for Forward Inference ```bash python - <<'PY' import torch from transformers import EsmTokenizer, EsmForMaskedLM model_path = "weight/PLMs/SaProt_650M_AF2" device = "cuda" if torch.cuda.is_available() else "cpu" tokenizer = EsmTokenizer.from_pretrained(model_path) model = EsmForMaskedLM.from_pretrained(model_path) model.to(device) model.eval() seq = "M#EvVpQpL#VyQdYaKv" tokens = tokenizer.tokenize(seq) print(tokens) inputs = tokenizer(seq, return_tensors="pt") inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) print(outputs.logits.shape) PY ``` ## Load SaProt with the ESM Interface If the model directory contains `SaProt_650M_AF2.pt`, you can use the ESM loading function provided by the project: ```bash python - <<'PY' from scripts.utils.esm_loader import load_esm_saprot model_path = "weight/PLMs/SaProt_650M_AF2/SaProt_650M_AF2.pt" model, alphabet = load_esm_saprot(model_path) print("ESM SaProt load OK") PY ``` ## Convert a Structure File to a Structure-Aware Sequence ```bash python - <<'PY' from scripts.utils.foldseek_util import get_struc_seq pdb_path = "scripts/example/8ac8.cif" parsed_seqs = get_struc_seq("scripts/bin/foldseek", pdb_path, ["A"], plddt_mask=False)["A"] seq, foldseek_seq, combined_seq = parsed_seqs print(f"seq: {seq}") print(f"foldseek_seq: {foldseek_seq}") print(f"combined_seq: {combined_seq}") PY ``` The `A` chain selection extracts only chain A from the structure file. The `combined_seq` in the returned result is a structure-aware sequence that can be used directly by SaProt. ## Mutation Effect Prediction ```bash python - <<'PY' import torch from model.saprot.saprot_foldseek_mutation_model import SaprotFoldseekMutationModel config = { "foldseek_path": None, "config_path": "weight/PLMs/SaProt_650M_AF2", "load_pretrained": True, } model = SaprotFoldseekMutationModel(**config) tokenizer = model.tokenizer device = "cuda" if torch.cuda.is_available() else "cpu" model.eval() model.to(device) seq = "M#EvVpQpL#VyQdYaKv" mut_info = "V3A" mut_value = model.predict_mut(seq, mut_info) print(mut_value) mut_info = "V3A:Q4M" mut_value = model.predict_mut(seq, mut_info) print(mut_value) mut_pos = 3 mut_dict = model.predict_pos_mut(seq, mut_pos) print(mut_dict) mut_pos = 3 mut_dict = model.predict_pos_prob(seq, mut_pos) print(mut_dict) PY ``` ## Extract Protein Embeddings ```bash python - <<'PY' import torch from model.saprot.base import SaprotBaseModel from transformers import EsmTokenizer config = { "task": "base", "config_path": "weight/PLMs/SaProt_650M_AF2", "load_pretrained": True, } model = SaprotBaseModel(**config) tokenizer = EsmTokenizer.from_pretrained(config["config_path"]) device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) model.eval() seq = "M#EvVpQpL#VyQdYaKv" tokens = tokenizer.tokenize(seq) print(tokens) inputs = tokenizer(seq, return_tensors="pt") inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): embeddings = model.get_hidden_states(inputs, reduction="mean") print(embeddings[0].shape) PY ``` ## Protein Inverse Folding Inverse folding requires additional inverse folding model weights: ```text https://huggingface.co/westlake-repl/SaProt_650M_AF2_inverse_folding ``` After downloading, place them at: ```text weight/PLMs/SaProt_650M_AF2_inverse_folding ``` Example: ```bash python - <<'PY' import torch from model.saprot.saprot_if_model import SaProtIFModel config = { "config_path": "weight/PLMs/SaProt_650M_AF2_inverse_folding", "load_pretrained": True, } device = "cuda" if torch.cuda.is_available() else "cpu" model = SaProtIFModel(**config) model = model.to(device) aa_seq = "##########" struc_seq = "dddddddddd" pred_aa_seq = model.predict(aa_seq, struc_seq) print(pred_aa_seq) PY ``` # Training This repository uses a unified training entry point: ```bash python scripts/training.py -c