Update README.md

README.md

@@ -41,19 +41,14 @@ All models share the same architecture and dictionary size (10,240). See [here](
 Extract interpretable features from protein sequences:
 
 ```python
-from huggingface_hub import hf_hub_download
-from interplm.sae.inference import load_model
+from interplm.sae.inference import load_sae_from_hf
 from interplm.esm.embed import embed_list_of_prot_seqs
 
 # Select ESM layer (must be one of 1-6)
-layer_num = 4
-
-# Download and load the model
-weights_path = hf_hub_download(
-    repo_id=f"Elana/InterPLM-esm2-8m",
-    filename=f"layer_{layer_num}/ae_normalized.pt"
-)
-sae = load_model(weights_path)
+layer = 4
+
+# Load specific layer SAE
+sae = load_sae_from_hf(plm_model="esm2-8m", plm_layer=layer)
 
 # Get ESM embeddings for protein
 protein_embeddings = embed_single_sequence(sequence="MRWQEMGYIFYPRKLR",