You can derive residue-level or per-protein embeddings using the ProtT5-XL-U50 model.
Key steps in the workflow:
- Load Tokenizer & Model: Use
T5Tokenizer and T5EncoderModel from transformers. - 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. - Sequence Preprocessing: Replace rare/ambiguous amino acids (
U, Z, O, B) with X and insert white-space between all amino acids. - Tokenization: Tokenize sequences with
add_special_tokens=True and padding="longest". - Embedding Extraction: Generate embeddings using
model(input_ids=..., attention_mask=...). - 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)