modernprost-base-deprecated

Model Card for modernprost-base

<!-- Provide a quick summary of what the model is/does. -->

modernprost-base is a lightweight 48M parameter specialist protein language model (pLM) which predicts a protein's Foldseek 3Di encoding from its amino acid sequence. It is designed to be a lightweight replacement for ProstT5 providing similar accurate in a more efficient and far smaller mode.

Model Details

Model Description

modernprost-base uses the ModernBERT architecture. It was trained using over 19M high quality protein structure predictions taken from AlphaFold Database, Big Fantastic Virus Database (BFVD), Phold's database and various metagenomic ColabFold database proteins.

• Developed by: George Bouras (GitHub @gbouras13) and Victor Mihaila (HuggingFace @Victor1306 )
• Model type: Encoder-only
• Language(s) (NLP): Protein sequence and structure
• License: MIT

How to Get Started with the Model

Feature extraction:

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

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained('gbouras13/modernprost-base',trust_remote_code=True)

# Load the model
model = AutoModel.from_pretrained('gbouras13/modernprost-base',trust_remote_code=True).to(device)

# prepare your protein sequences/structures as a list
sequence_examples = ["MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAED", "LQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN"]
# replace all rare/ambiguous amino acids by X 
for seq in sequence_examples:
    seq = seq.replace("U", "X").replace("Z", "X").replace("O", "X")

# tokenize seqs
tokenizer_kwargs = dict(
                        text=sequence_examples,
                        padding="longest",
                        truncation=False,
                        return_tensors="pt",
                        add_special_tokens=False
                    )

token_encoding = tokenizer(**tokenizer_kwargs).to(device)
attn_mask = token_encoding.attention_mask  # [B, L]
# generate embeddings
with torch.no_grad():
    outputs = model(
        token_encoding.input_ids,
        attention_mask=token_encoding.attention_mask,
    )
    

# if you want the embeddings for the first sequence length 60
emb_0 = outputs.hidden_states[-1][0,0:59]
# if you want the embeddings for the second sequence length 50
emb_1 = outputs.hidden_states[-1][0,0:49]

# 3Di


logits = outputs.logits
tokenized_preds = torch.argmax(logits, dim=-1)  # [B, L] 

ss_mapping = {
    0: "A", 1: "C", 2: "D", 3: "E", 4: "F",
    5: "G", 6: "H", 7: "I", 8: "K", 9: "L",
    10: "M", 11: "N", 12: "P", 13: "Q", 14: "R",
    15: "S", 16: "T", 17: "V", 18: "W", 19: "Y", 20: "X"
}

for preds, mask in zip(tokenized_preds, attn_mask):
    # keep only non-padding positions
    valid_preds = preds[mask.bool()]

    threedi = "".join(ss_mapping[int(tok.item())] for tok in valid_preds)
    print(threedi)
    

Training Details

Training Data

Pre-training data (3Di+AA sequences for over 19M proteins)

Training Procedure and Hyperparameters

Pre-training consists of predicting the Foldseek 3Di corresponding to the input amino acid. You can find all training scripts at https://github.com/gbouras13/distill_prostt5.

Speed

Predicting 3Di strings from AA input is extremely fast. On an AMD MI250x GPU, we were able to predict 3Di for 1 million proteins from Logan in just over 8 minutes. Pre-sorting proteins by length and batching is recommended for the best inference performance. Please see https://github.com/gbouras13/distill_prostt5 for a more feature-rich implementation.