Update README.md

README.md

@@ -42,21 +42,18 @@ Extract interpretable features from protein sequences:
 
 ```python
 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 = 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",
-                                           model_name="esm2_t6_8M_UR50D",
-                                           layer=layer_num)
-
-# Extract features
-features = sae.encode(protein_embeddings)
+from interplm.esm.embed import embed_single_sequence
+
+# Get ESM embeddings for protein sequence
+embeddings = embed_single_sequence(
+   sequence="MRWQEMGYIFYPRKLR",
+   model_name="esm2_t6_8M_UR50D",
+   layer=4  # Choose ESM layer (1-6)
+)
+
+# Load SAE model and extract features 
+sae = load_sae_from_hf(plm_model="esm2-8m", plm_layer=4)
+features = sae.encode(embeddings)
 ```
 
 For detailed training and analysis examples, see the [GitHub README](https://github.com/ElanaPearl/InterPLM/blob/main/README.md).