README.md

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

<p align="center">
  <strong>
    <span style="font-size: 30px;">SaProt</span>
  </strong>
</p>

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

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

This configuration reads from the following paths by default:

```text
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:

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

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

## Zero-Shot Mutation Effect Evaluation

ProteinGym evaluation:

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

The output file is saved by default to:

```text
output/ProteinGym/SaProt_650M_AF2.tsv
```

ClinVar evaluation:

```bash
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:

```text
output/ClinVar/SaProt_650M_AF2
```

Single-GPU environments can also use the adapted configuration:

```bash
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:

```bash
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.