Update README.md

README.md

@@ -21,8 +21,35 @@ It is a single dictionary of subspaces for 16K concepts and serves as a drop-in
 # 3. How can I use these dictionaries straight away?
 
 ```python
+from huggingface_hub import hf_hub_download
 import pyvene as pv
 
+# Create an intervention.
+class Encoder(pv.CollectIntervention):
+    """An intervention that reads concept latent from streams"""
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs, keep_last_dim=True)
+        self.proj = torch.nn.Linear(
+                self.embed_dim, kwargs["latent_dim"], bias=False)
+    def forward(self, base, source=None, subspaces=None):
+        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")
+encoder = Encoder(embed_dim=params.shape[0], latent_dim=params.shape[1])
+encoder.proj.weight.data = params.float()
+
+# Mount the loaded intervention.
+pv_model = pv.IntervenableModel({
+   "component": f"model.layers[20].output",
+   "intervention": encoder}, model=model)
+
+# use pv_model just as other torch model, and you can collect subspace latent.
+prompt = "Would you be able to travel through time using a wormhole?"
+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