ProtTrans

repository·master·Indexed 23 days ago

https://github.com/agemagician/prottrans

A collection of pre-trained transformer models for proteins designed for bioinformatics research. It supports feature extraction, fine-tuning, protein sequence generation, and prediction tasks. The library includes various model series such as ProtT5 (including ProtT5-XL-U50), ProtBert, ProtAlbert, ProtXLNet, and ProtElectra, available via Hugging Face, Zenodo, or Google Colab.

Tokens
24.1K
Snippets
64
Records
100
Agent score
79%

What's inside ProtTrans

  1. Explore ProtTrans usage modes

    master

    ProtTrans supports several primary workflows for protein analysis:

  2. Compare ProtTrans performance with other pLMs

    master

    ProtTrans models (specifically ProtT5-XL-UniRef50, or ProtT5-XL-U50) have been benchmarked against other protein language models (pLMs) like ESM.

    Key Performance Notes:

    • Half-Precision: For ProtT5-XL-U50, performance measurements were taken using embeddings extracted from the encoder-side of the T5 model in half-precision mode (model.half()). No performance degradation was observed when using half-precision to speed up embedding generation.
    • Top Performing Tasks: ProtT5-XL-U50 shows competitive or superior accuracy in tasks such as Subcellular localization, Conservation (ConSurf-DB), Variant effect (DMS-data), CATH superfamily classification, and Binding residues prediction.
  3. Set up a conda environment for fine-tuning

    master

    To run the fine-tuning notebooks, you can use the provided finetuning.yml file to create a consistent conda environment.

    Follow the official Conda documentation for creating an environment from a .yml file to ensure all dependencies are correctly installed.

  4. Install ProtTrans dependencies

    master

    ProtTrans models are available via the Hugging Face transformers library. You need to install torch, transformers, and sentencepiece.

    Additionally, if you encounter an UnboundLocalError: cannot access local variable 'sentencepiece_model_pb2', you may need to install protobuf manually.

    Note: If using a recent version of transformers, you might see a warning regarding T5 tokenization. You can resolve this by explicitly setting legacy=True in the tokenizer or by ignoring it, as legacy=True is the default behavior.

    pip install torch
    pip install transformers
    pip install sentencepiece
    # If encountering sentencepiece_model_pb2 errors:
    pip install protobuf
  5. Derive protein embeddings with ProtT5 (Quick Start)

    master

    You can derive residue-level or per-protein embeddings using the ProtT5-XL-U50 model.

    Key steps in the workflow:

    1. Load Tokenizer & Model: Use T5Tokenizer and T5EncoderModel from transformers.
    2. Device Management: Use GPU for half-precision. If using CPU, you must explicitly convert the model to torch.float32 as half-precision is currently only supported on GPUs.
    3. Sequence Preprocessing: Replace rare/ambiguous amino acids (U, Z, O, B) with X and insert white-space between all amino acids.
    4. Tokenization: Tokenize sequences with add_special_tokens=True and padding="longest".
    5. Embedding Extraction: Generate embeddings using model(input_ids=..., attention_mask=...).
    6. Post-processing: Extract the last_hidden_state. For residue embeddings, remove padded and special tokens (e.g., [0,:7]). For a single per-protein representation, take the .mean(dim=0) of the residue embeddings.
    from transformers import T5Tokenizer, T5EncoderModel
    import torch
    import re
    
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    
    # Load the tokenizer
    tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)
    
    # Load the model
    model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc").to(device)
    
    # only GPUs support half-precision currently; if you want to run on CPU use full-precision
    if device == torch.device("cpu"):
        model.to(torch.float32)
    
    # prepare your protein sequences as a list
    sequence_examples = ["PRTEINO", "SEQWENCE"]
    
    # replace all rare/ambiguous amino acids by X and introduce white-space between all amino acids
    sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", sequence))) for sequence in sequence_examples]
    
    # tokenize sequences and pad up to the longest sequence in the batch
    ids = tokenizer(sequence_examples, add_special_tokens=True, padding="longest")
    
    input_ids = torch.tensor(ids['input_ids']).to(device)
    attention_mask = torch.tensor(ids['attention_mask']).to(device)
    
    # generate embeddings
    with torch.no_grad():
        embedding_repr = model(input_ids=input_ids, attention_mask=attention_mask)
    
    # extract residue embeddings for the first ([0,:]) sequence in the batch and remove padded & special tokens ([0,:7]) 
    emb_0 = embedding_repr.last_hidden_state[0,:7] # shape (7 x 1024)
    
    # if you want to derive a single representation (per-protein embedding) for the whole protein
    emb_0_per_protein = emb_0.mean(dim=0) # shape (1024)
  6. Install requirements for ProtTrans

    master

    To use ProtTrans for protein feature extraction or fine-tuning, you must install the following libraries:

    • pytorch
    • transformers (from Hugging Face)

    If you intend to perform model visualization, you also need to install:

    • BertViz
  7. Prepare protein sequences for embedding

    master

    Before tokenization, protein sequences should be pre-processed to ensure compatibility with the models:

    1. Remove whitespace: Ensure sequences are joined correctly.
    2. Handle ambiguous amino acids: Replace rare or ambiguous amino acids (like U, Z, O, B) with X.
    3. Tokenization format: When using batch_encode_plus, set is_split_into_words=True if your input is a list of individual amino acid characters.
  8. Configure LoRA for ProtT5

    master

    The implementation uses LoRAConfig to define how Low-Rank Adaptation is applied to the transformer layers. You can specify which modules and layers to target using regular expressions.

    Key configuration parameters:

    • lora_rank: The rank of the adaptation matrices.
    • lora_modules: Regex to match transformer modules (e.g., .*SelfAttention|.*EncDecAttention).
    • lora_layers: Regex to match specific layers within modules (e.g., q|k|v|o).
    • trainable_param_names: Regex to identify which parameters (like LayerNorm or LoRA weights) should be trainable.
    class LoRAConfig:
        def __init__(self):
            self.lora_rank = 4
            self.lora_init_scale = 0.01
            self.lora_modules = ".*SelfAttention|.*EncDecAttention"
            self.lora_layers = "q|k|v|o"
            self.trainable_param_names = ".*layer_norm.*|.*lora_[ab].*"
            self.lora_scaling_rank = 1
  9. How ProtT5-XL-UniRef50 encoder vs decoder works

    master

    ProtT5-XL-UniRef50 is a model containing both an encoder and a decoder.

    • For feature extraction: You should load only the encoder part using T5EncoderModel. This reduces inference speed and GPU memory requirements by approximately 50%.
    • For full model usage: If you require both the encoder and decoder (e.g., for generative tasks), you should use T5Model instead of T5EncoderModel.