evillegasgarcia/sae_esm2_6_l3

Vanilla Sparse AutoEncoder trained on embeddings from layer 3 of esm2_t6_8M_UR50D.

For more details check the arxiv preprint and the github repository.

To use:

Download the class defining the sparse autoencoder from github.

git clone [email protected]:edithvillegas/plm-sae.git
cd plm-sae

Load the base ESM2 model and the sparse autoencoder from huggingface.

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.

#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.

#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)