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:
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
1 hf download OneScience-Group/ProteinMPNN --local-dir ./ProteinMPNN
2 cd ProteinMPNN
Install the Runtime Environment
DCU Environment
1 # Activate DTK and CONDA first
2 conda create -n onescience311 python = 3.11 -y
3 conda activate onescience311
4 # uv installation is supported
5 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/onescienceQuick Verification
1 export PYTHONPATH = $( pwd ) /model: ${ONESCIENCE_ROOT} /src: ${PYTHONPATH :- }
2 python -c "from proteinmpnn.protein_mpnn_utils import ProteinMPNN; print('proteinmpnn wrapper ok')"
3 python scripts/inference.py --help
4 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
1 cd /path/to/proteinmpnn
2 bash scripts/test_inference.sh
This script uses the following by default:
1 Input PDB: data/inputs/PDB_monomers/pdbs/5L33.pdb
2 Designed chain: A
3 Model weights: weight/vanilla_model_weights
4 Output directory: outputs/test_inference/
View the generated sequences:
ls outputs/test_inference/seqsEquivalent command:
1 python scripts/inference.py \
2 --pdb_path ./data/inputs/PDB_monomers/pdbs/5L33.pdb \
3 --pdb_path_chains "A" \
4 --out_folder ./outputs/test_inference \
5 --path_to_model_weights ./weight/vanilla_model_weights \
6 --model_name v_48_020 \
7 --num_seq_per_target 2 \
8 --sampling_temp "0.1" \
9 --seed 37 \
10 --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.shInference on multiple single-chain PDBs. submit_example_2.shMulti-chain complex; design only the specified chains. submit_example_3.shInference on a single PDB complex. submit_example_3_score_only.shScore existing structures/sequences without generating new sequences. submit_example_3_score_only_from_fasta.shScore a structure using FASTA sequences. submit_example_4.shFix certain residue positions and exclude them from design. submit_example_4_non_fixed.shDesign only the specified positions. submit_example_5.shTied positions, with multi-position tied design. submit_example_6.shHomooligomer-constrained design. submit_example_7.shOutput unconditional probabilities. submit_example_8.shAdd a global amino acid bias. submit_example_pssm.shAdd PSSM constraints to assist design.
Run a single example:
bash scripts/infer_examples/submit_example_3.shNote: 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:
1 list.csv
2 valid_clusters.txt
3 test_clusters.txt
4 pdb/<2nd-3rd characters of pdbid>/<pdbid>.pt
5 pdb/<2nd-3rd characters of pdbid>/<pdbid>_<chain>.pt
Run the training example script directly:
1 cd /path/to/proteinmpnn
2 bash scripts/test_train.sh
This script uses the following by default:
1 Training data: data/pdb_2021aug02_sample
2 Output directory: outputs/train/exp_020/
3 Number of samples per epoch: 1000
4 Save a checkpoint every 50 epochs
View the training log and weights:
1 cat outputs/train/exp_020/log.txt
2 ls outputs/train/exp_020/model_weights
Equivalent command:
1 python scripts/training.py \
2 --path_for_training_data ./data/pdb_2021aug02_sample \
3 --path_for_outputs ./outputs/train/exp_020 \
4 --num_examples_per_epoch 1000 \
5 --save_model_every_n_epochs 50
To resume training, pass:
1 python scripts/training.py \
2 --path_for_training_data ./data/pdb_2021aug02_sample \
3 --path_for_outputs ./outputs/train/exp_020 \
4 --previous_checkpoint ./outputs/train/exp_020/model_weights/epoch_last.pt
Common Parameters
Inference Parameters
Parameter Description Default/Example --pdb_pathInput path for a single PDB ./data/inputs/PDB_monomers/pdbs/5L33.pdb--jsonl_pathParsed PDB JSONL input path Generated by parse_multiple_chains.py --pdb_path_chainsChains to design in single-PDB mode "A" or "A B"--out_folderInference output directory ./outputs/test_inference--path_to_model_weightsWeight directory ./weight/vanilla_model_weights--model_nameWeight file name without .pt v_48_020--num_seq_per_targetNumber of sequences to generate for each target 2--sampling_tempSampling temperature "0.1"--score_onlyScore only, without generating new sequences 0 or 1--save_scoreSave score files 0 or 1--save_probsSave probability files 0 or 1--ca_onlyUse the CA-only model Disabled by default --use_soluble_modelUse the soluble protein model Disabled by default
Training Parameters
Parameter Description Default/Example --path_for_training_dataPreprocessed training data directory ./data/pdb_2021aug02_sample--path_for_outputsTraining output directory ./outputs/train/exp_020--previous_checkpointCheckpoint for resuming training epoch_last.pt--num_epochsNumber of training epochs Default 200 --num_examples_per_epochNumber of samples loaded per epoch Example 1000 --batch_sizeToken batch size Default 10000 --save_model_every_n_epochsSave a checkpoint every N epochs 50 in the example script--mixed_precisionWhether to use mixed precision Default True
Official OneScience Information
Citations and License
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.
