SaProt / README.md

Upload folder using huggingface_hub

3b99abb verified about 18 hours ago

16.2 kB

license: mit
language:
  - en
  - zh
tags:
  - OneScience
  - life-science
  - protein
  - protein-language-model
  - structure-aware
  - mutation-effect
  - embedding
  - inverse-folding
  - SaProt
frameworks: PyTorch

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:

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

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

# 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:

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:

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:

weight/PLMs/SaProt_650M_AF2

To run the ESM2 comparison experiment, also prepare:

huggingface-cli download \
  facebook/esm2_t33_650M_UR50D \
  --local-dir ./weight/PLMs/esm2_t33_650M_UR50D

The corresponding configuration reads from:

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:

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:

SaProt/
└── scripts/
    └── bin/
        └── foldseek

Then grant it execute permission:

chmod +x scripts/bin/foldseek

The Foldseek path used by the current configuration is:

scripts/bin/foldseek

3) Downstream Task Datasets

The official downstream task datasets are available at:

https://drive.google.com/drive/folders/11dNGqPYfLE3M-Mbh4U7IQpuHxJpuRr4g?usp=sharing

For this adapted version, extract the downstream task data to:

scripts/LMDB/

Typical paths used by the configuration include:

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:

westlake-repl/AF2_UniRef50
https://huggingface.co/datasets/westlake-repl/AF2_UniRef50

The official pretraining configuration uses LMDB data directories such as:

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

hf download OneScience-Group/SaProt \
  --local-dir ./SaProt

cd SaProt

Quick Verification

Check the dependencies:

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:

./scripts/bin/foldseek version

Test model loading:

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:

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:

*.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

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:

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

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

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

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:

https://huggingface.co/westlake-repl/SaProt_650M_AF2_inverse_folding

After downloading, place them at:

weight/PLMs/SaProt_650M_AF2_inverse_folding

Example:

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:

python scripts/training.py -c <config_path>

Configuration files are located in conf/, model code is located in model/, and data processing and utility code is located in scripts/. The current configuration uses the pretrained model weights in weight/PLMs/SaProt_650M_AF2 by default.

Pretraining

To pretrain SaProt from scratch or continue pretraining, first prepare the pretraining LMDB dataset, then run:

python scripts/training.py -c conf/pretrain/saprot.yaml

This configuration reads from the following paths by default:

scripts/LMDB/AF2_Uniref50/foldseek/train
scripts/LMDB/AF2_Uniref50/foldseek/valid

Downstream Fine-Tuning

Use the following commands to fine-tune SaProt on downstream tasks:

# Thermostability
python scripts/training.py -c conf/Thermostability/saprot.yaml

# EC
python scripts/training.py -c conf/EC/saprot.yaml

# GO
python scripts/training.py -c conf/GO/MF/saprot.yaml
python scripts/training.py -c conf/GO/BP/saprot.yaml
python scripts/training.py -c conf/GO/CC/saprot.yaml

# Metal ion binding
python scripts/training.py -c conf/MetalIonBinding/saprot.yaml

# Human PPI
python scripts/training.py -c conf/HumanPPI/saprot.yaml

# Contact prediction
python scripts/training.py -c conf/Contact/saprot.yaml

# DeepLoc
python scripts/training.py -c conf/DeepLoc/cls2/saprot.yaml
python scripts/training.py -c conf/DeepLoc/cls10/saprot.yaml

For single-GPU or limited-memory environments, use conf/scnet/Thermostability_saprot_1gpu.yaml as a starting point:

python scripts/training.py -c conf/scnet/Thermostability_saprot_1gpu.yaml

Zero-Shot Mutation Effect Evaluation

ProteinGym evaluation:

python scripts/mutation_zeroshot.py -c conf/ProteinGym/saprot.yaml

The output file is saved by default to:

output/ProteinGym/SaProt_650M_AF2.tsv

ClinVar evaluation:

python scripts/mutation_zeroshot.py -c conf/ClinVar/saprot.yaml
python scripts/compute_clinvar_auc.py -c conf/ClinVar/saprot.yaml

ClinVar prediction results are saved by default to:

output/ClinVar/SaProt_650M_AF2

Single-GPU environments can also use the adapted configuration:

python scripts/mutation_zeroshot.py -c conf/scnet/ClinVar_saprot.yaml
python scripts/compute_clinvar_auc.py -c conf/scnet/ClinVar_saprot.yaml

ESM2 Comparison Experiment

To run the ESM2 baseline, additionally prepare the weights in weight/PLMs/esm2_t33_650M_UR50D and the corresponding normal LMDB data. Example commands:

python scripts/training.py -c conf/Thermostability/esm2.yaml
python scripts/mutation_zeroshot.py -c conf/ProteinGym/esm2.yaml

Output Description

Task	Main output
Structure encoding	AA sequence, 3Di sequence, and structure-aware sequence
Model forward pass	Token-level logits
Protein representation	Residue-level/protein-level embeddings
Mutation effect prediction	Mutation score
Zero-shot evaluation	ProteinGym Spearman results or ClinVar AUC results
Inverse folding	Protein sequences generated or evaluated under structural conditions
Downstream fine-tuning	Prediction results and model checkpoint for the corresponding task

Official OneScience Information

Platform	Main OneScience repository	Skills repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

The official SaProt source repository is released under the MIT License, which permits use, modification, distribution, sublicensing, and commercial use. When copying or distributing it, retain the original copyright notice and the MIT License text.
SaProt model weights are released independently through Hugging Face. For commercial use, redistribution, or other purposes, check and comply with the license on each corresponding model page. The relevant pretraining and downstream datasets are also subject to the licenses and terms of use on their respective dataset pages.
This repository is a DCU-adapted version of SaProt. Use of the repository code, model weights, and related data remains subject to the licenses and terms of use of their respective original projects.