Update README.md

README.md

@@ -15,35 +15,44 @@ pip install git+https://github.com/etowahadams/interprot.git
 
 ## Usage
 
-Load the SAE
+Install InterProt, load ESM and SAE
 ```python
+import torch
+from transformers import AutoTokenizer, EsmModel
 from safetensors.torch import load_file
 from interprot.sae_model import SparseAutoencoder
+from huggingface_hub import hf_hub_download
 
-sae_model = SparseAutoencoder(1280, 4096)
-checkpoint_path = 'esm2_plm1280_l24_sae4096.safetensors'
-sae_model.load_state_dict(load_file(checkpoint_path))
-```
+ESM_DIM = 1280
+SAE_DIM = 4096
+LAYER = 24
 
-Load ESM and run ESM inference -> SAE inference
-```
-import torch
-from transformers import AutoTokenizer, EsmModel
-
-# Load ESM model and tokenizer
+# Load ESM model
 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
+# Load SAE model
+checkpoint_path = hf_hub_download(
+    repo_id="liambai/InterProt-ESM2-SAEs",
+    filename="esm2_plm1280_l24_sae4096.safetensors"
+)
+sae_model = SparseAutoencoder(ESM_DIM, SAE_DIM)
+sae_model.load_state_dict(load_file(checkpoint_path))
+```
+
+ESM -> SAE inference on an amino acid sequence of length `L`
+```
 seq = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVVAAIVQDIAYLRSLGYNIVATPRGYVLAGG"
-esm_layer = 24
 
-inputs = tokenizer([seq], padding=True, return_tensors="pt")
+# Tokenize sequence and run ESM inference
+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)
+# esm_layer_acts has shape (L+2, ESM_DIM), +2 for BoS and EoS tokens
+esm_layer_acts = outputs.hidden_states[LAYER][0]
+
+# Using ESM embeddings from LAYER, run SAE inference
+sae_acts = sae_model.get_acts(esm_layer_acts) # (L+2, SAE_DIM)
+sae_acts
 ```