File size: 2,844 Bytes
822e13f 1558ca7 822e13f b5d932c 822e13f b5d932c 822e13f 1b6c3e6 b5d932c 1558ca7 822e13f 1558ca7 822e13f 1558ca7 822e13f 1558ca7 822e13f b5d932c 4f950cb b5d932c 4f950cb 822e13f b5d932c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
---
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
|