| # ProtBERT |
|
|
| ```python |
| 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) |
| ``` |
