File size: 1,880 Bytes
f3915b2 9427608 f3915b2 2a56cd1 c5064e6 9427608 600090c 9427608 c5064e6 600090c c5064e6 600090c 9427608 |
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 |
---
tags:
- model_hub_mixin
- pytorch_model_hub_mixin
- model
paper: https://arxiv.org/abs/2502.09135
---
Vanilla Sparse AutoEncoder trained on embeddings from layer 3 of esm2_t6_8M_UR50D.
For more details check the [arxiv preprint](https://arxiv.org/abs/2502.09135) and
the [github repository](https://github.com/edithvillegas/plm-sae).
**To use:**
Download the class defining the sparse autoencoder from github.
```bash
git clone [email protected]:edithvillegas/plm-sae.git
cd plm-sae
```
Load the base ESM2 model and the sparse autoencoder from huggingface.
```python
from transformers import AutoTokenizer, AutoModelForMaskedLM
from sae.SAE_methods import AutoEncoder #import sparse autoencoder from local definition
#load ESM2 model
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
model = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
model = model.to("cuda")
#load SAE (GPU-only)
sparse_autoencoder = AutoEncoder.from_pretrained("evillegasgarcia/sae_esm2_6_l3")
```
Prepare auxiliary functions to extract embeddings from a specific point in the ESM2 model.
```python
#setup to extract ESM2 embeddings
layer_name = "esm.encoder.layer.3.output"
#define hook
intermediate_embs = dict()
def hook(module, input, output):
intermediate_embs[layer_name] = output
return hook
#attach hook
hook_handle = model.esm.encoder.layer[3].output.register_forward_hook(l3_hook)
```
Extract embeddings from the ESM2 model and then from the sparse autoencoder.
```python
#Inference
sequence = "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPL"
#PLM Inference
tokenized = tokenizer.encode(sequence, return_tensors="pt")
tokenized = tokenized.to("cuda")
outputs = model(tokenized)
embeddings = intermediate_embs[layer_name][0]
#SAE Inference
_, _, sae_embeddings, _, _ = sparse_autoencoder(embeddings)
```
|