Implementation for the paper Evaluating Tuning Strategies for Sequence Generation with Protein Language Models.

Adapted (forked and modified/enhanced) from https://github.com/corolla-johnson/mkultra, which is a prompt tuning implementation for NLP, supporting GPT-2. The prompt tuning method was developed by Lester et al., see Lester et al. The power of scale for parameter-efficient prompt tuning. For setup, you can reconstruct our conda environment for prompt tuning by the

environment.yaml

file.

Models

This implementation currently supports ProtGPT2 (Ferruz et al. ProtGPT2 is a deep unsupervised language model for protein design.) and RITA (Hesslow et al. RITA: a Study on Scaling Up Generative Protein Sequence Models.). See

mkultra/tuning.py

for implementation details.

Training

Training scripts for RITA and ProtGPT2 are

RITA_prompt_tuning.py

and

ProtGPT2_prompt_tuning.py

, respectively. They can be configured with parameters specified in a JSON config, see

training_configs/

folder and the Trainer documentation in

mkultra/trainers.py

. You can enable memory tracking during training, but this makes the training slower. The training configurations in

training_configs/

are the ones that we used for our paper. This includes the configs for training and evaluation, as well as those for runtime measuring and memory tracking.

An example notebook for training a prompt for RITA is

RITA_prompt_tuning_example.ipynb

, open in Colab here.

Datasets

You can train the model with Fasta datasets if you use the

FastaDataset

class (

mkultra.sequence_loader.py

) as dataset input for a PyTorch

DataLoader

. Dataset preprocessing as in RITA can be done with the

prepare_dataset.ipynb

notebook in the

utils

folder. Also, have a look into the script

utils/clustering.sh

to see the configuration we used to cluster our datasets using MMseqs2. The current setup of the dataset notebook and clustering script are for clustering with 100% sequence similarity threshold, but you can adjust that.

We also provide the datasets that we used for our experiments in the

datasets/

folder. They contain sequences from the Pfam family PF03272. The dataset (

datasets/InterProUniprotPF03272.fasta

) was downloaded from InterPro on January 5, 2023 (Paysan-Lafosse et al. InterPro in 2022.). Then, we removed all sequences containing an X, which created the dataset

InterProUniprotPF03272_Xremoved.fasta

. We clustered our data with MMseqs2(Steinegger et al. MMseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. And: Steinegger et al. Clustering huge protein sequence sets in linear time.). The datasets clustered with a sequence similarity threshold of 100%, which were our main datasets, are

InterProUniprotPF03272prepared_<train|validation|test>.fasta

. Datasets clustered with different thresholds can be found as

InterProUniprotPF03272_<threshold>_<train|validation|test>.fasta

Evaluation

Perplexity

This implementation currently supports perplexity evaluation as perplexity per token. See

RITA_prompt_comparison_basemodel.py

and

RITA_prompt_comparison_trainvaltest.py

, which are the evaluation scripts that we used for our experiments. If you want to write an evaluation script yourself, have a look at

mkultra/evaluator.py

Protein family prediction for generated sequences

We further evaluated the amount of generated sequences that were classified by ProtCNN (Bileschi et al. Using deep learning to annotate the protein universe.) as belonging to our target family PF03272. The script we used for that is

protcnn.py

, taken and adapted from the official ProtCNN notebook. As ProtCNN uses Tensorflow and we wanted to keep it separate from our other scripts, we created an extra conda environment for it, specified in

environment-protcnn.yml

. For using the script, you have to download the ProtCNN model and vocabulary as described in the official notebook:

wget -qN https://storage.googleapis.com/brain-genomics-public/research/proteins/pfam/models/single_domain_per_sequence_zipped_models/seed_random_32.0/5356760.tar.gz
tar xzf 5356760.tar.gz
wget https://storage.googleapis.com/brain-genomics-public/research/proteins/pfam/models/single_domain_per_sequence_zipped_models/trained_model_pfam_32.0_vocab.json

You can set a fixed sliding window size and/or stride, or let the prediction be run on all possible windows of minimum size 50 (or the sequence length if shorter) for a sequence. Our code also supports running an ensemble of multiple ProtCNN models (as described in the ProtCNN paper), for this you have to modify the list of saved models in the script and download the respective additional models. Further, in single-model mode, you can set a probability threshold to discard predictions with a lower probability. In our final experiments, we used no fixed window size or stride, a single model, and a probability threshold of 0.5.

In addition to ProtCNN, we also evaluated protein family prediction with HMMER with profile-specific gathering thresholds, which you can run as follows:

hmmsearch --cut_ga --tblout <path/to/output/file> <path/to/HMM/file> <path/to/sequences/file>

The HMMER runs for all our sets of generated sequences are bundled in the script

hmmer_search.sh

. We use the Pfam profile HMM downloaded from InterPro for the family.

Activity filters for generated sequences

To reproduce our evaluations of protein activity (see Johnson et al. Computational Scoring and Experimental Evaluation of Enzymes Generated by Neural Networks.), run

activity_prediction_metrics/metrics.py

for all generated datasets. See also

activity_prediction_metrics/metrics.ipynb

for an interactive version. Because these evaluations need specific dependencies, use an extra conda environment that is created with the same packages as specifed in the

environment-activity.yml

file. Furthermore, the ESM likelihood computation requires the Github repo https://github.com/seanrjohnson/protein_gibbs_sampler. We advise to run the notebook

activity_prediction_metrics/metrics.ipynb

once initially, because it installs all required dependencies additional to those from the

environment-activity.yml

Then, use

activity_prediction_metrics/aggregate.ipynb

to aggregate the results. This counts the activity predictions for each generated dataset and writes everything into

activity_prediction_metrics/activity.csv

for further usage, e.g. plotting.

Sequence generation

For generating sequences, instantiate a prompt tuning model (see

mkultra/tuning.py

) and then load and add a prompt (

see mkultra/checkpoint_loader.py

) that was trained for that type of model, as for example in

RITA_prompt_sequence_generation.py

. In our experiments, we generated 193 sequences (size of our test set) in batches of 10.

A note about finetuning

In our paper, we compare the performance of prompt-tuned models to that of finetuned models. For finetuning, we use the run_clm.py script by Huggingface with the same batch sizes as for prompt tuning. You have to add

trust_remote_code=True

to the model loading call in line 376. To generate the datasets as txt, you can use the

prepare_dataset.ipynb

notebook in the

utils

folder.

Implementation for the paper Evaluating Tuning Strategies for Sequence Generation with Protein Language Models.

environment.yaml

file.

Models

mkultra/tuning.py

for implementation details.

Training

Training scripts for RITA and ProtGPT2 are

RITA_prompt_tuning.py

and

ProtGPT2_prompt_tuning.py

, respectively. They can be configured with parameters specified in a JSON config, see

training_configs/

folder and the Trainer documentation in

mkultra/trainers.py

. You can enable memory tracking during training, but this makes the training slower. The training configurations in

training_configs/

are the ones that we used for our paper. This includes the configs for training and evaluation, as well as those for runtime measuring and memory tracking.

An example notebook for training a prompt for RITA is

RITA_prompt_tuning_example.ipynb

, open in Colab here.

Datasets

You can train the model with Fasta datasets if you use the

FastaDataset

class (

mkultra.sequence_loader.py

) as dataset input for a PyTorch

DataLoader

. Dataset preprocessing as in RITA can be done with the

prepare_dataset.ipynb

notebook in the

utils

folder. Also, have a look into the script

utils/clustering.sh

We also provide the datasets that we used for our experiments in the

datasets/

folder. They contain sequences from the Pfam family PF03272. The dataset (

datasets/InterProUniprotPF03272.fasta

) was downloaded from InterPro on January 5, 2023 (Paysan-Lafosse et al. InterPro in 2022.). Then, we removed all sequences containing an X, which created the dataset

InterProUniprotPF03272_Xremoved.fasta

InterProUniprotPF03272prepared_<train|validation|test>.fasta

. Datasets clustered with different thresholds can be found as

InterProUniprotPF03272_<threshold>_<train|validation|test>.fasta

Evaluation

Perplexity

This implementation currently supports perplexity evaluation as perplexity per token. See

RITA_prompt_comparison_basemodel.py

and

RITA_prompt_comparison_trainvaltest.py

, which are the evaluation scripts that we used for our experiments. If you want to write an evaluation script yourself, have a look at

mkultra/evaluator.py

Protein family prediction for generated sequences

protcnn.py

, taken and adapted from the official ProtCNN notebook. As ProtCNN uses Tensorflow and we wanted to keep it separate from our other scripts, we created an extra conda environment for it, specified in

environment-protcnn.yml

. For using the script, you have to download the ProtCNN model and vocabulary as described in the official notebook:

wget -qN https://storage.googleapis.com/brain-genomics-public/research/proteins/pfam/models/single_domain_per_sequence_zipped_models/seed_random_32.0/5356760.tar.gz
tar xzf 5356760.tar.gz
wget https://storage.googleapis.com/brain-genomics-public/research/proteins/pfam/models/single_domain_per_sequence_zipped_models/trained_model_pfam_32.0_vocab.json

In addition to ProtCNN, we also evaluated protein family prediction with HMMER with profile-specific gathering thresholds, which you can run as follows:

hmmsearch --cut_ga --tblout <path/to/output/file> <path/to/HMM/file> <path/to/sequences/file>

The HMMER runs for all our sets of generated sequences are bundled in the script

hmmer_search.sh

. We use the Pfam profile HMM downloaded from InterPro for the family.

Activity filters for generated sequences

To reproduce our evaluations of protein activity (see Johnson et al. Computational Scoring and Experimental Evaluation of Enzymes Generated by Neural Networks.), run

activity_prediction_metrics/metrics.py

for all generated datasets. See also

activity_prediction_metrics/metrics.ipynb

for an interactive version. Because these evaluations need specific dependencies, use an extra conda environment that is created with the same packages as specifed in the

environment-activity.yml

file. Furthermore, the ESM likelihood computation requires the Github repo https://github.com/seanrjohnson/protein_gibbs_sampler. We advise to run the notebook

activity_prediction_metrics/metrics.ipynb

once initially, because it installs all required dependencies additional to those from the

environment-activity.yml

Then, use

activity_prediction_metrics/aggregate.ipynb

to aggregate the results. This counts the activity predictions for each generated dataset and writes everything into

activity_prediction_metrics/activity.csv

for further usage, e.g. plotting.

Sequence generation

For generating sequences, instantiate a prompt tuning model (see

mkultra/tuning.py

) and then load and add a prompt (

see mkultra/checkpoint_loader.py

) that was trained for that type of model, as for example in

RITA_prompt_sequence_generation.py

. In our experiments, we generated 193 sequences (size of our test set) in batches of 10.

A note about finetuning

trust_remote_code=True

to the model loading call in line 376. To generate the datasets as txt, you can use the

prepare_dataset.ipynb

notebook in the

utils

folder.

Implementation for the paper Evaluating Tuning Strategies for Sequence Generation with Protein Language Models.

Implementation for the paper Evaluating Tuning Strategies for Sequence Generation with Protein Language Models.

Models

Training

Datasets

Evaluation

Perplexity

Protein family prediction for generated sequences

Activity filters for generated sequences

Sequence generation

A note about finetuning

Related Skills

<h1 align="center">

2. Apply Deepthink Protocol (reason about dependencies

- Identify gaps

Implementation for the paper Evaluating Tuning Strategies for Sequence Generation with Protein Language Models.

Models

Training

Datasets

Evaluation

Perplexity

Protein family prediction for generated sequences

Activity filters for generated sequences

Sequence generation

A note about finetuning