SaeedLab/ProtEnrich-ProtT5

ProtT5

from transformers import AutoTokenizer, AutoModel, T5Tokenizer, T5EncoderModel
import torch

tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc')
encoder = T5EncoderModel.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc')
protenrich = AutoModel.from_pretrained("SaeedLab/ProtEnrich-ProtT5", trust_remote_code=True)

seqs = ["MKTFFVLLL"]
seqs = [" ".join(i) for i in seqs]
inputs = tokenizer(seqs, return_tensors="pt", add_special_tokens=True)

with torch.no_grad():
  outputs = encoder(**inputs)
  pooled = outputs.last_hidden_state[0, :-1].mean(axis=0)
  enriched = protenrich(pooled)

print('H enrich:', enriched.h_enrich)
print('H anchor:', enriched.h_anchor)
print('H algn:', enriched.h_algn)
print('Structure:', enriched.struct)
print('Dynamics:', enriched.dyn)