SaeedLab/ProtEnrich-ProtBERT

ProtBERT

from transformers import AutoTokenizer, AutoModel
import torch

tokenizer = AutoTokenizer.from_pretrained("Rostlab/prot_bert")
encoder = AutoModel.from_pretrained("Rostlab/prot_bert")
protenrich = AutoModel.from_pretrained("SaeedLab/ProtEnrich-ProtBERT", trust_remote_code=True)

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

with torch.no_grad():
  outputs = encoder(**inputs)
  pooled = outputs.last_hidden_state[0, 1:-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)