Update README.md

README.md

@@ -20,8 +20,15 @@ It is a single dictionary of subspaces for 16K concepts and serves as a drop-in
 
 ```python
 from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import pyvene as pv
 
+# Load model and tokenizer
+model_name = "google/gemma-2-2b-it"
+model = AutoModelForCausalLM.from_pretrained(model_name).cuda()
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
 # Create an intervention.
 class Encoder(pv.CollectIntervention):
     """An intervention that reads concept latent from streams"""
@@ -33,7 +40,8 @@ class Encoder(pv.CollectIntervention):
         return torch.relu(self.proj(base))
 
 # Loading weights
-path_to_params = hf_hub_download(repo_id="pyvene/gemma-reft-2b-it-res", filename="l20/weight.pt")
+path_to_params = hf_hub_download(repo_id="pyvene/gemma-diffmean-9b-it-res", filename="l20/weight.pt")
+params = torch.load(path_to_params)
 encoder = Encoder(embed_dim=params.shape[0], latent_dim=params.shape[1])
 encoder.proj.weight.data = params.float()
 
@@ -48,6 +56,7 @@ input_ids = torch.tensor([tokenizer.apply_chat_template(
     [{"role": "user", "content": prompt}], tokenize=True, add_generation_prompt=True)]).cuda()
 acts = pv_model.forward(
     {"input_ids": input_ids}, return_dict=True).collected_activations[0]
+
 ```
 
 # 4. Point of Contact