|
|
--- |
|
|
library_name: transformers |
|
|
tags: [] |
|
|
--- |
|
|
# FastESM |
|
|
|
|
|
## A faster half-precision version of ESM2-650 that leverages FlashAttention2 |
|
|
|
|
|
FastESM is a fully Huggingface compatible version rewritten with a newer PyTorch Attention implementation which will run FlashAttention2 when possible. |
|
|
|
|
|
To produce the FastESM weights, we trained ESM2-650 50000 additional steps in fp16 mixed precision on [OMG50](tattabio/OMG_prot50) up to sequence length of **2048**. |
|
|
|
|
|
Outputting attentions and predicting contacts are not possible from SDPA. Various other optimizations also make the base implementation slightly different than the HF one. |
|
|
|
|
|
## Use with 🤗 transformers |
|
|
```python |
|
|
import torch |
|
|
from transformers import AutoModel, AutoTokenizer |
|
|
|
|
|
model_path = 'Synthyra/FastESM2_650' |
|
|
model = AutoModel.from_pretrained(model_path, torch_dtype=torch.float16, trust_remote_code=True).eval() |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) |
|
|
|
|
|
sequence = 'MSEQWENCE' |
|
|
tokenized = tokenizer(sequence, return_tensors='pt') |
|
|
with torch.no_grad(): |
|
|
embeddings = model(**tokenized).last_hidden_state |
|
|
|
|
|
print(embeddings.shape) # (1, 11, 1280) |
|
|
``` |
|
|
|
|
|
## Embed entire datasets with no new code |
|
|
To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the progress bar is usually much longer than the actual time. |
|
|
```python |
|
|
embeddings = model.embed_dataset( |
|
|
sequences=sequences, # list of protein strings |
|
|
batch_size=16, # embedding batch size |
|
|
max_len=2048, # truncate to max_len |
|
|
full_embeddings=True, # return residue-wise embeddings |
|
|
full_precision=False, # store as float32 |
|
|
pooling_type='mean', # use mean pooling if protein-wise embeddings |
|
|
num_workers=0, # data loading num workers |
|
|
sql=False, # return dictionary of sequences and embeddings |
|
|
) |
|
|
|
|
|
_ = model.embed_dataset( |
|
|
sequences=sequences, # list of protein strings |
|
|
batch_size=16, # embedding batch size |
|
|
max_len=2048, # truncate to max_len |
|
|
full_embeddings=True, # return residue-wise embeddings |
|
|
full_precision=False, # store as float32 |
|
|
pooling_type='mean', # use mean pooling if protein-wise embeddings |
|
|
num_workers=0, # data loading num workers |
|
|
sql=True, # store sequences in local SQL database |
|
|
sql_db_path='embeddings.db', # path to .db file of choice |
|
|
) |
|
|
``` |
|
|
|
|
|
## Comparison of half precisions |
|
|
Presumabely because we trained in mixed-precision fp16, fp16 has closer outputs to the fp32 weights then bf16. Therefore, we recommend loading in fp16. |
|
|
|
|
|
When summing the MSE of 1000 sequences vs. the fp32 weights: |
|
|
|
|
|
Average MSE for FP16: 0.00000140 |
|
|
|
|
|
Average MSE for BF16: 0.00004125 |
|
|
|
|
|
### FlashAttention2 |
|
|
Requires PyTorch 2.5+ for the most savings, see [SDPA](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html). |
|
|
|
|
|
### Citation |
|
|
|
