ProteinDT

KO | EN

gitlite — search

#javascript #python #hacktoberfest #ai #llm #go #golang #linux #macos #nodejs #awesome #cli #vue #mcp

(★ 111)

A Text-guided Protein Design Framework, Nat Mach Intell 2025 (https://www.nature.com/articles/s42256-025-01011-z)

File Explorer

Download Latest Version (.zip)

- - test.yml
- - prepare_01_peptide_editing_raw_and_processed_data.py
  - prepare_01_stability_raw_data_Pin1.py
  - prepare_01_stability_raw_data_Villin.py
  - prepare_02_processed_data.py
  - prepare_03_train_peptide_editing_evaluator.py
  - README.md
  - step_01_editing_Galactica.py
  - step_01_editing_latent_interpolation.py
  - step_01_editing_latent_optimization.py
  - step_01_evaluate_region.py
  - step_01_evaluate_stability.py
  - step_01_evaluate_structure.py
  - step_02_binding_editing_Galactica.py
  - step_02_binding_editing_latent_interpolation.py
  - step_02_binding_editing_latent_optimization.py
  - utils.py
  - utils_analysis.py
  - utils_peptide_editing.py
- - step_01_text_retrieval.py
  - step_01_text_retrieval.txt
  - step_02_inference_ChatGPT.py
  - step_02_inference_Galactica.py
  - step_02_inference_ProteinDT.py
  - utils.py
- downstream_TAPE.py
- pretrain_step_01_CLAP.py
- pretrain_step_02_empty_sequence.py
- pretrain_step_02_pairwise_representation.py
- pretrain_step_03_facilitator.py
- pretrain_step_04_decoder.py
- pretrain_step_05_AE.py
- pretrain_step_05_inference.py
- README.md
- final.gif
- pipeline.png
- pretraining_roadmap.png
- environment.yml
- - README.md
  - step_01_parse_XML.py
  - step_02_generate_protein_text_pair.py
- - __init__.py
  - dataset_Pin1.py
  - dataset_Protein.py
  - dataset_RepresentationPair.py
  - dataset_SecondaryStructure.py
  - dataset_Stability.py
  - dataset_SwissProtCLAP.py
  - dataset_Villin.py
- - __init__.py
  - model_BindingModel.py
  - model_CDConv.py
  - model_FoldingBindingInferenceModel.py
  - model_GaussianFacilitator.py
  - model_MultinomialDiffusionDecoder.py
  - model_ProteinText.py
  - model_RNNPrediction.py
  - model_SDE.py
  - model_TransformerDecoder.py
- - __init__.py
  - datasets.py
  - metrics.py
  - models.py
  - trainer.py
- - __init__.py
  - tokenize.py
- - submit_step_01_Galactica.sh
  - submit_step_01_latent_interpolation_Multinomial.sh
  - submit_step_01_latent_interpolation_T5.sh
  - submit_step_01_latent_optimization.sh
  - submit_step_02_Galactica.sh
  - submit_step_02_latent_interpolation_Multinomial.sh
  - submit_step_02_latent_interpolation_T5.sh
  - submit_step_02_latent_optimization.sh
- - submit_Multinomial_BertBase.sh
  - submit_Multinomial_RNN.sh
  - submit_T5_T5Base.sh
LICENSE
README.md
setup.py

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# ProteinDT: A Text-guided Protein Design Framework

Authors: Shengchao Liu, Yanjing Li, Zhuoxinran Li, Anthony Gitter, Yutao Zhu, Jiarui Lu, Zhao Xu, Weili Nie, Arvind Ramanathan, Chaowei Xiao\*, Jian Tang\*, Hongyu Guo\*, Anima Anandkumar\*

\* jointly supervised

[[Project Page](https://chao1224.github.io/ProteinDT)] [[ArXiv](https://arxiv.org/abs/2302.04611)]
[[Datasets on HuggingFace](https://huggingface.co/datasets/chao1224/ProteinDT/tree/main)] [[Checkpoints on HuggingFace](https://huggingface.co/chao1224/ProteinDT/tree/main)]

## 1 Environment

```
conda create -n ProteinDT python=3.7
conda activate ProteinDT

conda install -y numpy networkx scikit-learn

pip install torch==1.10.*

pip install transformers
pip install lxml

# for TAPE
pip install lmdb
pip install seqeval

# for baseline ChatGPT
pip install openai==0.28.0

# for baseline Galactica
pip install accelerate

# for visualization
pip install matplotlib

# for binding editing
pip install h5py
pip install torch_geometric==2.0 torch_scatter torch_sparse torch_cluster
pip install biopython

# for ESM folding
pip install "fair-esm[esmfold]"
pip install dm-tree omegaconf ml-collections einops
pip install fair-esm[esmfold]==2.0.0  --no-dependencies # Override deepspeed==0.5 
pip install 'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'

conda install -c conda-forge -yq mdtraj

# for ProteinDT
pip install .
```

## 2 Pretraining Datasets (SwissProtCLAP) Preparation

Please check folder `preprocess/SwissProtCLAP` for SwissProtCLAP construction from UniProt.

We also provide a copy of SwissProtCLAP at [this HuggingFace link](https://huggingface.co/datasets/chao1224/ProteinDT/tree/main). Or you can use the following script:
```
from huggingface_hub import HfApi, snapshot_download
api = HfApi()
snapshot_download(repo_id="chao1224/ProteinDT", repo_type="dataset", cache_dir='./')
```

Then move the data under `./data` folder. The data structure is
```
./data/
└── SwissProtCLAP
    ├── protein_sequence.txt
    └── text_sequence.txt
```

## 3 Pretraining

Go to folder `examples`, and do the pretraining in 5 steps. We summarize the logics of these 5 steps as below:

 <img src="figures/pretraining_roadmap.png" width="75%" />

The pretrained checkpoints can be found at [this HuggingFace link](https://huggingface.co/chao1224/ProteinDT/tree/main).
Before getting started, first we need to define our output home folder, e.g., `export OUTPUT_DIR=../output/ProteinDT/hyper_01`.

- Step 1. Conduct CLAP pretraining
    - On a single GPU card:
        ```
        python pretrain_step_01_CLAP.py \
        --protein_lr=1e-5 --protein_lr_scale=1 \
        --text_lr=1e-5 --text_lr_scale=1 \
        --protein_backbone_model=ProtBERT_BFD \
        --epochs=10 --batch_size=9 --num_workers=0 \
        --output_model_dir="$OUTPUT_DIR"
        ```

- We also support distribution learning with DDP. Example of using a server with 8 GPU cards:
        ```
        CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
        python pretrain_step_01_CLAP.py \
        --protein_lr=1e-5 --protein_lr_scale=1 \
        --text_lr=1e-5 --text_lr_scale=1 \
        --protein_backbone_model=ProtBERT_BFD \
        --epochs=10 --batch_size=9 --num_workers=0 \
        --output_model_dir="$OUTPUT_DIR"
        ```

- Step 2. Obtain frozen representation:
    ```
    python pretrain_step_02_empty_sequence.py \
    --protein_backbone_model=ProtBERT_BFD \
    --batch_size=16 --num_workers=0 \
    --pretrained_folder="$OUTPUT_DIR"

python pretrain_step_02_pairwise_representation.py \
    --protein_backbone_model=ProtBERT_BFD \
    --batch_size=16 --num_workers=0 \
    --pretrained_folder="$OUTPUT_DIR"
    ```

- Step 3. Learn the facilitator distribution:
    ```
    python pretrain_step_03_facilitator.py \
    --protein_lr=1e-5 --protein_lr_scale=1 \
    --text_lr=1e-5 --text_lr_scale=1 \
    --protein_backbone_model=ProtBERT_BFD \
    --epochs=10 --batch_size=9 --num_workers=0 \
    --pretrained_folder="$OUTPUT_DIR" \
    --output_model_folder="$OUTPUT_DIR"/step_03_Gaussian_10
    ```

- Step 4. Learn the decoder distribution. Notice that we have three types of decoder distribution models:
    - A Transformer-based auto-regressive decoder. Here we adopt the T5 architecture.
        ```
        python pretrain_step_04_decoder.py \
        --num_workers=0 --lr=1e-4 --epochs=50 \
        --decoder_distribution=T5Decoder \
        --score_network_type=T5Base \
        --hidden_dim=16 \
        --pretrained_folder="$OUTPUT_DIR" \
        --output_folder="$OUTPUT_DIR"/step_04_T5
        ```
    - A discrete denoising diffusion model (multinomial diffusion).
        - Using RNN as score network:
            ```
            python pretrain_step_04_decoder.py \
            --num_workers=0 --lr=1e-4 --epochs=50 \
            --decoder_distribution=MultinomialDiffusion \
            --score_network_type=RNN \
            --hidden_dim=16 \
            --pretrained_folder="$OUTPUT_DIR" \
            --output_folder="$OUTPUT_DIR"/step_04_MultiDiffusion_RNN
            ```

- Using BERT as score network:
            ```
            python pretrain_step_04_decoder.py \
            --num_workers=0 --lr=1e-4 --epochs=50 \
            --decoder_distribution=MultinomialDiffusion \
            --score_network_type=BertBase \
            --hidden_dim=16 \
            --pretrained_folder="$OUTPUT_DIR" \
            --output_folder="$OUTPUT_DIR"/step_04_MultiDiffusion_BERT
            ```

- Step 5. learn an auto-encoder that is specifically designed for text-guided editing task. You can also treat this as a downstream task.
    ```
    python pretrain_step_05_AE.py \
    --num_workers=0 --lr=1e-4 --epochs=50 \
    --pretrained_folder="$OUTPUT_DIR" \
    --output_folder="$OUTPUT_DIR"/step_05
    ```

## 4 Downstream Tasks

We include three types of downstream tasks, as will be introduced below. You can find the scripts for first two downstream tasks under folder `scripts`.

### 4.1 Text-to-Protein Generation

First let's go to the folder `examples/downstream_Text2Protein`.

Then we sample text sequences for text-to-protein generation:
```
python step_01_text_retrieval.py
```
We also provide the sampled text data in `step_01_text_retrieval.txt`. You can replace it with the text sequences you want to use.

Now we can do the text-to-sequence generation, e.g., if we use T5 as the decoder:
```
export OUTPUT_DIR=../../output/ProteinDT/hyper_01

python step_02_inference_ProteinDT.py \
--decoder_distribution=T5Decoder --score_network_type=T5Base \
--num_workers=0 --hidden_dim=16 --batch_size=8 \
--pretrained_folder="$OUTPUT_DIR" \
--step_04_folder="$OUTPUT_DIR"/step_04_T5 \
--num_repeat=16 --use_facilitator --AR_generation_mode=01 \
--output_text_file_path="$OUTPUT_DIR"/step_04_T5/downstream_Text2Protein/step_02_inference.txt
```

### 4.2 Zero-shot Text-guided Protein Editing

First let's go to the folder `examples/downstream_Editing`.

The dataset preparation can be found at `examples/downstream_Edting/README.md`. You can also find it on [this HuggingFace link](https://huggingface.co/datasets/chao1224/ProteinDT/tree/main/downstream_Editing/datasets_and_checkpoints). We include three types of editing tasks: stability, structure, and peptide binding. In terms of the methods, we have two types: latent optimization and latent interpolation. The demo scripts are explained below.

#### 4.2.1 Latent Optimization
- Structure / Stability: `editing_task: alpha, beta, Villin, Pin1`.
    ```
    export OUTPUT_DIR=../../output/ProteinDT/hyper_01

python step_01_editing_latent_optimization.py \
    --num_workers=0 --batch_size=8 \
    --lambda_value=0.9 --num_repeat=16 --oracle_mode=text --temperature=2 \
    --editing_task=alpha --text_prompt_id=101 \
    --pretrained_folder="$OUTPUT_DIR" \
    --step_05_folder="$OUTPUT_DIR"/step_05_AE \
    --output_folder="$OUTPUT_DIR"/step_05_AE/downstream_Editing_latent_optimization/alpha_prompt_101_lambda_0.9_num_repeat_16_oracle_text_T_2 \
    --output_text_file_path="$OUTPUT_DIR"/step_05_AE/downstream_Editing_latent_optimization/alpha_prompt_101_lambda_0.9_num_repeat_16_oracle_text_T_2/step_01_editing.txt

python step_01_evaluate_structure.py \
    --num_workers=0 --batch_size=8 --editing_task=alpha --text_prompt_id=101 \
    --output_folder="$OUTPUT_DIR"/step_05_AE/downstream_Editing_latent_optimization/alpha_prompt_101_lambda_0.9_num_repeat_16_oracle_text_T_2 \
    --output_text_file_path="$OUTPUT_DIR"/step_05_AE/downstream_Editing_latent_optimization/alpha_prompt_101_lambda_0.9_num_repeat_16_oracle_text_T_2/step_01_editing.txt
    ```
- Peptide binding
    ```
    export OUTPUT_DIR=../../output/ProteinDT/hyper_01

python step_01_editing_latent_optimization.py \
    --num_workers=0 --batch_size=4 \
    --lambda_value=0.9 --num_repeat=16 --oracle_mode=text --temperature=2 \
    --editing_task=peptide_binding --text_prompt_id=101 \
    --pretrained_folder="$OUTPUT_DIR" \
    --step_05_folder="$OUTPUT_DIR"/step_05_AE \
    --output_folder="$OUTPUT_DIR"/step_05_AE/downstream_Editing_latent_optimization/peptide_binding_prompt_101_lambda_0.9_num_repeat_16_oracle_text_T_2 \
    --output_text_file_path="$OUTPUT_DIR"/step_05_AE/downstream_Editing_latent_optimization/peptide_binding_prompt_101_lambda_0.9_num_repeat_16_oracle_text_T_2/step_02_editing.txt
    ```

#### 4.2.2 Latent Interpolation
Notice that for latent interpolation, we have three models: auto-regressive (T5), denoising diffusion model (RNN and BERT). We provide demos scripts using T5.
- Structure / Stability: `editing_task: alpha, beta, Villin, Pin1`.
    ```
    export OUTPUT_DIR=../../output/ProteinDT/hyper_01

python step_01_editing_latent_interpolation.py \
    --editing_task=alpha --text_prompt_id=101 \
    --decoder_distribution=T5Decoder --score_network_type=T5Base \
    --num_workers=0 --hidden_dim=16 --batch_size=2 \
    --theta=0.9 --num_repeat=16 --oracle_mode=text --AR_generation_mode=01 --AR_condition_mode=expanded \
    --pretrained_folder="$OUTPUT_DIR" --step_04_folder="$OUTPUT_DIR"/step_04_T5 \
    --output_folder="$OUTPUT_DIR"/step_04_T5/downstream_Editing_latent_interpolation_alpha/prompt_101_theta_0.9_num_repeat_16_oracle_text_inference_01_expanded \
    --output_text_file_path="$OUTPUT_DIR"/step_04_T5/downstream_Editing_latent_interpolation_alpha/prompt_101_theta_0.9_num_repeat_16_oracle_text_inference_01_expanded/step_01_editing.txt

python step_01_evaluate_structure.py \
    --num_workers=0 --batch_size=1 \
    --editing_task=alpha --text_prompt_id=101 \
    --output_folder="$OUTPUT_DIR"/step_04_T5/downstream_Editing_latent_interpolation_alpha/prompt_101_theta_0.9_num_repeat_16_oracle_text_inference_01_expanded \
    --output_text_file_path="$OUTPUT_DIR"/step_04_T5/downstream_Editing_latent_interpolation_alpha/prompt_101_theta_0.9_num_repeat_16_oracle_text_inference_01_expanded/step_01_editing.txt
    ```
- Peptide binding
    ```
    export OUTPUT_DIR=../../output/ProteinDT/hyper_01

python step_02_binding_editing_latent_interpolation.py \
    --editing_task=peptide_binding --text_prompt_id=101 \
    --decoder_distribution=T5Decoder --score_network_type=T5Base \
    --num_workers=0 --hidden_dim=16 --batch_size=1 \
    --theta=0.9 --num_repeat=16 --oracle_mode=text --AR_generation_mode=01 --AR_condition_mode=expanded \
    --pretrained_folder="$OUTPUT_DIR" --step_04_folder="$OUTPUT_DIR"/step_04_T5 \
    --output_folder="$OUTPUT_DIR"/step_04_T5/downstream_Editing_latent_interpolation_peptide_binding/prompt_101_theta_0.9_num_repeat_16_oracle_text_inference_01_expanded \
    --output_text_file_path="$OUTPUT_DIR"/step_04_T5/downstream_Editing_latent_interpolation_peptide_binding/prompt_101_theta_0.9_num_repeat_16_oracle_text_inference_01_expanded/step_02_editing.txt
    ```

### 4.3 Protein Property Prediction

First please download the TAPE data following instructions [here](https://github.com/songlab-cal/tape?tab=readme-ov-file#lmdb-data). We also provide it at [this HuggingFace link](https://huggingface.co/datasets/chao1224/ProteinDT/tree/main).

Under `examples`, and the script is `downstream_TAPE.py`. We follow the exactly same hyper-parameter as [OntoProtein](https://github.com/zjunlp/OntoProtein).

```
python downstream_TAPE.py \
--task_name=ss3 \
--seed=3 \
--learning_rate=3e-5 \
--num_train_epochs=5 \
--per_device_train_batch_size=2 \
--gradient_accumulation_steps=8 \
--warmup_ratio=0.08 \
--pretrained_model=ProteinDT \
--pretrained_folder="$OUTPUT_DIR" \
--output_dir="$OUTPUT_DIR"/downstream_TAPE
```

## Cite Us
Feel free to cite this work if you find it useful to you!
```
@article{liu2023text,
    title={A Text-guided Protein Design Framework},
    author={Shengchao Liu, Yanjing Li, Zhuoxinran Li, Anthony Gitter, Yutao Zhu, Jiarui Lu, Zhao Xu, Weili Nie, Arvind Ramanathan, Chaowei Xiao, Jian Tang, Hongyu Guo, Anima Anandkumar},
    journal={arXiv preprint arXiv:2302.04611},
    year={2023}
}
```

ProteinDT

File Explorer

Today's picks

Was this content helpful?

Recent Feedback

ProteinDT

File Explorer

Today's picks

Was this content helpful?

Recent Feedback

Download README