---
license: apache-2.0
---
# InterProt ESM2 SAE Models

A set of SAE models trained on [ESM2-650](https://huggingface.co/facebook/esm2_t33_650M_UR50D) activations using 1M protein sequences from [UniProt](https://www.uniprot.org/). The SAE implementation mostly followed [Gao et al.](https://arxiv.org/abs/2406.04093) with Top-K activation function, though with much fewer latent dimensions.

Check out [https://interprot.com](https://interprot.com) for an interactive visualizer of the 4096-dimensional SAE on ESM layer 24.

## Installation

```bash
pip install git+https://github.com/etowahadams/interprot.git
```

## Usage

Load the SAE
```python
from safetensors.torch import load_file
from interprot.sae_model import SparseAutoencoder

sae_model = SparseAutoencoder(1280, 4096)
checkpoint_path = 'esm2_plm1280_l24_sae4096.safetensors'
sae_model.load_state_dict(load_file(checkpoint_path))
```

Load ESM and run ESM inference -> SAE inference
```
import torch
from transformers import AutoTokenizer, EsmModel

# Load ESM model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
esm_model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")

# Run ESM inference with some sequence and take layer 24 activations
seq = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVVAAIVQDIAYLRSLGYNIVATPRGYVLAGG"
esm_layer = 24

inputs = tokenizer([seq], padding=True, return_tensors="pt")
with torch.no_grad():
    outputs = esm_model(**inputs, output_hidden_states=True)
esm_layer_acts = outputs.hidden_states[esm_layer] # (1, sequence length + 2, 1280)

# Run SAE inference with ESM activations as input
sae_acts = sae_model.get_acts(esm_layer_acts)
sae_acts # (1, sequence length + 2, 4096)
```