README updates

README.md

@@ -24,8 +24,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"""
@@ -38,6 +45,7 @@ class Encoder(pv.CollectIntervention):
 
 # Loading weights
 path_to_params = hf_hub_download(repo_id="pyvene/gemma-reft-2b-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()