OneScience-Group/ProteinMPNN

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

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

hf download OneScience-Group/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:

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

Training Parameters

Official OneScience Information

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.