Update README.md

ajay-naidu-96 · Jan 21, 2024 · bc3bb3c · bc3bb3c
1 parent 89cde43
commit bc3bb3c
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -30,17 +30,11 @@ import pyvene as pv
 _, tokenizer, gpt2 = pv.create_gpt2()
 
 pv_gpt2 = pv.IntervenableModel({
-    "layer": 0,
-    "component": "mlp_output",
-    "source_representation": torch.zeros(
-        gpt2.config.n_embd)
+    "source_representation": torch.zeros(gpt2.config.n_embd)
 }, model=gpt2)
 
 orig_outputs, intervened_outputs = pv_gpt2(
-    base = tokenizer(
-        "The capital of Spain is", 
-        return_tensors="pt"
-    ), 
+    base = tokenizer("The capital of Spain is", return_tensors="pt"), 
     unit_locations={"base": 3}
 )
 print(intervened_outputs.last_hidden_state - orig_outputs.last_hidden_state)