OneScience/ProteinMPNN

ProteinMPNN

Model Introduction

ProteinMPNN is a protein sequence design model based on a Message Passing Neural Network. Given a protein backbone structure, it can efficiently generate highly expressible and foldable amino acid sequences.

Model Description

ProteinMPNN uses an encoder-decoder architecture. The encoder extracts geometric and topological features of the backbone structure through a graph neural network, while the decoder generates the amino acid sequence position by position in an autoregressive manner.

Usage

1. Using OneCode

You can try intelligent one-click AI4S programming through the OneCode online environment:

Try intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• Running on a GPU or DCU is recommended.
• A CPU can be used for import checks and small-configuration connectivity validation, but full training and inference will be slow.
• DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster.

3. Quick Start

Download the Model Package

modelscope download --model OneScience/ProteinMPNN --local_dir ./ProteinMPNN 
cd ProteinMPNN 

Install the Runtime Environment

DCU Environment

# Activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is supported
pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

If the runtime environment explicitly needs to point to the OneScience root directory, set:

export ONESCIENCE_ROOT=/path/to/onescience

Quick Verification

export PYTHONPATH=$(pwd)/model:${ONESCIENCE_ROOT}/src:${PYTHONPATH:-}
python -c "from proteinmpnn.protein_mpnn_utils import ProteinMPNN; print('proteinmpnn wrapper ok')"
python scripts/inference.py --help
python scripts/training.py --help

Inference

The current weights have been placed in subdirectories under weight/. The standard ProteinMPNN uses weight/vanilla_model_weights/; if --path_to_model_weights is not explicitly passed, scripts/inference.py uses this directory by default.

Run Minimal Inference with the Script

cd /path/to/proteinmpnn
bash scripts/test_inference.sh

This script uses the following by default:

Input PDB: data/inputs/PDB_monomers/pdbs/5L33.pdb
Designed chain: A
Model weights: weight/vanilla_model_weights
Output directory: outputs/test_inference/

View the generated sequences:

ls outputs/test_inference/seqs

Equivalent command:

python scripts/inference.py \
  --pdb_path ./data/inputs/PDB_monomers/pdbs/5L33.pdb \
  --pdb_path_chains "A" \
  --out_folder ./outputs/test_inference \
  --path_to_model_weights ./weight/vanilla_model_weights \
  --model_name v_48_020 \
  --num_seq_per_target 2 \
  --sampling_temp "0.1" \
  --seed 37 \
  --batch_size 1

Available weights:

• Standard model: --path_to_model_weights ./weight/vanilla_model_weights
• Soluble protein model: --path_to_model_weights ./weight/soluble_model_weights or add --use_soluble_model
• CA-only model: --path_to_model_weights ./weight/ca_model_weights or add --ca_only

Inference Example Scripts

The scripts/infer_examples/ directory contains examples for 12 inference scenarios, all adapted to the current directory structure:

Script	Scenario
submit_example_1.sh	Inference on multiple single-chain PDBs.
submit_example_2.sh	Multi-chain complex; design only the specified chains.
submit_example_3.sh	Inference on a single PDB complex.
submit_example_3_score_only.sh	Score existing structures/sequences without generating new sequences.
submit_example_3_score_only_from_fasta.sh	Score a structure using FASTA sequences.
submit_example_4.sh	Fix certain residue positions and exclude them from design.
submit_example_4_non_fixed.sh	Design only the specified positions.
submit_example_5.sh	Tied positions, with multi-position tied design.
submit_example_6.sh	Homooligomer-constrained design.
submit_example_7.sh	Output unconditional probabilities.
submit_example_8.sh	Add a global amino acid bias.
submit_example_pssm.sh	Add PSSM constraints to assist design.

Run a single example:

bash scripts/infer_examples/submit_example_3.sh

Note: submit_example_3_score_only_from_fasta.sh depends on submit_example_3.sh first generating outputs/example_3_outputs/seqs/3HTN.fa.

Training

The current example training data is placed in data/pdb_2021aug02_sample/. This directory should contain:

list.csv
valid_clusters.txt
test_clusters.txt
pdb/<2nd-3rd characters of pdbid>/<pdbid>.pt
pdb/<2nd-3rd characters of pdbid>/<pdbid>_<chain>.pt

Run the training example script directly:

cd /path/to/proteinmpnn
bash scripts/test_train.sh

This script uses the following by default:

Training data: data/pdb_2021aug02_sample
Output directory: outputs/train/exp_020/
Number of samples per epoch: 1000
Save a checkpoint every 50 epochs

View the training log and weights:

cat outputs/train/exp_020/log.txt
ls outputs/train/exp_020/model_weights

Equivalent command:

python scripts/training.py \
  --path_for_training_data ./data/pdb_2021aug02_sample \
  --path_for_outputs ./outputs/train/exp_020 \
  --num_examples_per_epoch 1000 \
  --save_model_every_n_epochs 50

To resume training, pass:

python scripts/training.py \
  --path_for_training_data ./data/pdb_2021aug02_sample \
  --path_for_outputs ./outputs/train/exp_020 \
  --previous_checkpoint ./outputs/train/exp_020/model_weights/epoch_last.pt

Common Parameters

Inference Parameters

Parameter	Description	Default/Example
--pdb_path	Input path for a single PDB	./data/inputs/PDB_monomers/pdbs/5L33.pdb
--jsonl_path	Parsed PDB JSONL input path	Generated by parse_multiple_chains.py
--pdb_path_chains	Chains to design in single-PDB mode	"A" or "A B"
--out_folder	Inference output directory	./outputs/test_inference
--path_to_model_weights	Weight directory	./weight/vanilla_model_weights
--model_name	Weight file name without .pt	v_48_020
--num_seq_per_target	Number of sequences to generate for each target	2
--sampling_temp	Sampling temperature	"0.1"
--score_only	Score only, without generating new sequences	0 or 1
--save_score	Save score files	0 or 1
--save_probs	Save probability files	0 or 1
--ca_only	Use the CA-only model	Disabled by default
--use_soluble_model	Use the soluble protein model	Disabled by default

Training Parameters

Parameter	Description	Default/Example
--path_for_training_data	Preprocessed training data directory	./data/pdb_2021aug02_sample
--path_for_outputs	Training output directory	./outputs/train/exp_020
--previous_checkpoint	Checkpoint for resuming training	epoch_last.pt
--num_epochs	Number of training epochs	Default 200
--num_examples_per_epoch	Number of samples loaded per epoch	Example 1000
--batch_size	Token batch size	Default 10000
--save_model_every_n_epochs	Save a checkpoint every N epochs	50 in the example script
--mixed_precision	Whether to use mixed precision	Default True

Official OneScience Information

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

Citations and License

• Original ProteinMPNN paper: Robust deep learning-based protein sequence design using ProteinMPNN.

• ProteinMPNN-related source code uses the MIT License. See LICENSE in the repository root for details. The specific terms of use for model weights and data should follow the instructions provided by the corresponding publishers.

• If you use ProteinMPNN in research, it is recommended to cite the corresponding original ProteinMPNN paper and relevant OneScience project information, and to add citations for downstream analysis tools or datasets according to the actual task.