Tasks Assessing Protein Embeddings (TAPE)

TAPE provides data, weights, and code for running benchmarks on trained protein embeddings. It includes a pretraining corpus, five supervised downstream tasks, pretrained language model weights, and benchmarking code.

Important Compatibility Notes:

• This version uses PyTorch. Previous TensorFlow-based weights and code from the original repository are not compatible.
• This repository is designed for ease of use and future development rather than strict reproduction of the original paper's results. For exact paper reproduction, use the original NeurIPS 2019 repository.
• Training Recommendation: The maintainers no longer recommend using TAPE's internal training code for new projects, as it may not be updated for future PyTorch versions. Instead, it is recommended to use frameworks like PyTorch Lightning or Fairseq for training, while using TAPE for benchmarking and model availability.

If you use TAPE or its datasets in your research, you must cite the TAPE paper and all individual dataset components used. A complete list of BibTeX citations is available in data_refs.bib.

Key citations include:

• TAPE Paper: Rao et al. (2019)
• Pfam (Pretraining): El-Gebali et al. (2019)
• SCOPe (Remote Homology/Contact): Fox et al. (2013)
• PDB (Secondary Structure/Contact): Berman et al. (2000)
• CASP12 (Secondary Structure/Contact): Moult et al. (2018)
• NetSurfP2.0 (Secondary Structure): Klausen et al. (2019)
• ProteinNet (Contact): AlQuraishi (2019)
• Fluorescence: Sarkisyan et al. (2016)
• Stability: Rocklin et al. (2017)

@inproceedings{tape2019,
author = {Rao, Roshan and Bhattacharya, Nicholas and Thomas, Neil and Duan, Yan and Chen, Xi and Canny, John and Abbeel, Pieter and Song, Yun S},
title = {Evaluating Protein Transfer Learning with TAPE},
booktitle = {Advances in Neural Information Processing Systems}
year = {2019}
}

To train a model on a specific downstream task (e.g., secondary structure prediction) using weights from a previously trained language model, use the --from_pretrained flag with the path to your saved results.

Key Hyperparameters to Tune:

• batch_size
• learning_rate
• warmup_steps
• num_train_epochs

Optimization Tips: Use eval_freq and save_freq to reduce the frequency of validation passes and model saving, which is useful for downstream tasks where epochs are often shorter and more numerous.

# Example: Training a transformer on secondary structure using pretrained weights
tape-train-distributed transformer secondary_structure \
    --from_pretrained results/<path_to_your_saved_results> \
    --batch_size 64 \
    --learning_rate 0.0001 \
    --fp16 \
    --warmup_steps 500 \
    --nproc_per_node 1 \
    --gradient_accumulation_steps 1 \
    --num_train_epochs 50 \
    --eval_freq 5 \
    --save_freq 10

TAPE uses the Huggingface transformers API to define and provide pretrained models. Models are automatically downloaded and cached.

Available pretrained models:

• bert-base (Transformer model)
• babbler-1900 (UniRep model)
• xaa, xab, xac, xad, xae (trRosetta models)

Important Note on Embeddings: For Transformer models, the pooled_output is not trained and should not be used without fine-tuning. Instead, use the mean of the sequence output for embeddings.

import torch
from tape import ProteinBertModel, TAPETokenizer

# Load model and tokenizer
model = ProteinBertModel.from_pretrained('bert-base')
tokenizer = TAPETokenizer(vocab='iupac')  # use 'unirep' for UniRep model

# Prepare sequence
sequence = 'GCTVEDRCLIGMGAILLNGCVIGSGSLVAAGALITQ'
token_ids = torch.tensor([tokenizer.encode(sequence)])

# Forward pass
output = model(token_ids)
sequence_output = output[0]
# Use sequence_output mean instead of output[1] (pooled_output) for Transformers

TAPE includes a PyTorch implementation of the trRosetta model.

Setup:

1. Download the trRosetta data and place it under <data_path>/trrosetta.
2. Use TRRosettaDataset to load the data.

Model Selection: Use TRRosetta.from_pretrained(choice) where choice is one of: 'xaa', 'xab', 'xac', 'xad', or 'xae'. Each corresponds to one of the ensemble models.

Output: Predictions can be saved as .npz files and used with the Yang Lab structure modeling scripts.

from tape import TRRosetta
from tape.datasets import TRRosettaDataset

# Initialize datasets
train_data = TRRosettaDataset('<data_path>', 'train')  # subsamples MSAs
valid_data = TRRosettaDataset('<data_path>', 'valid')  # does not subsample MSAs

# Load model
model = TRRosetta.from_pretrained('xaa')

# Forward pass
batch = train_data.collate_fn([train_data[0]])
loss, predictions = model(**batch)

Use the tape-embed CLI command to generate .npz files containing protein embeddings from a FASTA file. The process is fully batched and automatically distributes across available GPUs.

Usage Pattern: tape-embed <model_type> <input_fasta> <output_npz> <model_name> --tokenizer <tokenizer_type>

Example (UniRep): To embed with the UniRep babbler-1900 model, you must use the unirep tokenizer.

Loading Results: Embeddings can be loaded into numpy. Each entry in the .npz file is a dictionary containing:

• avg: The average of the sequence embedding (recommended for Transformers).
• pooled: The pooled embedding (trained for models like UniRep, but not for Transformers).
• seq: The full sequence embedding (only if --full_sequence_embed flag is used).

If you need the full embedding instead of the average, use the --full_sequence_embed flag.

# Embed using UniRep
tape-embed unirep my_input.fasta output_filename.npz babbler-1900 --tokenizer unirep

# Python loading example
import numpy as np
arrays = np.load('output_filename.npz', allow_pickle=True)
# Access by protein ID (or '0', '1', etc. if unnamed)
protein_data = arrays['seq1'] 
# protein_data contains {'pooled': ..., 'avg': ...}

TAPE provides two commands for training:

1. tape-train: Uses standard PyTorch data distribution across GPUs.
2. tape-train-distributed: Uses torch.distributed.launch-style multiprocessing. Recommended as it typically provides a 10-15% speedup.

Training Features:

• Distributed training via multiprocessing
• Half-precision (--fp16)
• Gradient accumulation (--gradient_accumulation_steps)
• Automatic batching by sequence length

Batch Size Logic: The --batch_size parameter specifies the size used per backwards pass. The actual examples per GPU is calculated as: batch_size / num_gpus / gradient_accumulation_steps. If you encounter Out-of-Memory (OOM) errors, increase --gradient_accumulation_steps.

# Example: Training a transformer on masked language modeling
tape-train-distributed transformer masked_language_modeling \
    --batch_size 1024 \
    --learning_rate 0.001 \
    --fp16 \
    --warmup_steps 1000 \
    --nproc_per_node 2 \
    --gradient_accumulation_steps 4

To generate embeddings for a set of proteins using a pretrained model, use the run_embed function. When called without arguments, it uses create_embed_parser to parse command-line arguments.

Key Embedding Arguments:

• data_file: File containing the set of proteins to embed.
• out_file: Name of the output file.
• from_pretrained: Directory containing config and pretrained model weights.
• --batch_size: Batch size (default: 1024).
• --full_sequence_embed: If true, saves an embedding at every amino acid position in the sequence (warning: high disk usage).

Note: TAPE does not support distributed embedding passes; local_rank must not be set to something other than -1.

from tape.main import run_embed

# To run via CLI:
# python -m tape.main <data_file> <out_file> <from_pretrained_dir> --full_sequence_embed

# To run programmatically:
# run_embed(args=parsed_namespace)