add steering test script

requirements.txt

@@ -4,4 +4,5 @@ transformer_lens
 transformers
 sae-lens==3.18.2
 git+https://github.com/cyber-chris/activation_additions.git
-pandas
+pandas
+bitsandbytes

steering_test.py

@@ -0,0 +1,63 @@
+import torch
+from sae_lens import SAE, HookedSAETransformer
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+from transformer_lens import HookedTransformer
+import pandas as pd
+
+from activation_additions.prompt_utils import get_x_vector
+from activation_additions.completion_utils import gen_using_activation_additions
+
+# Helper script to test steering methods.
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def generate_steered(model, prompt):
+    sampling_kwargs = {
+        "do_sample": True,
+        # "top_k": 50,
+        # "top_p": 0.95,
+        "temperature": 0.5,
+    }
+
+    outputs = []
+    for act_name in range(4, 28, 2):
+        x_vectors = get_x_vector(
+            prompt1="I'm sorry, but I should not",
+            prompt2="Sure, I can help with that",
+            coeff=8,
+            act_name=act_name,
+            model=model,
+            pad_method="tokens_right",
+        )
+
+        mod_df: pd.DataFrame = gen_using_activation_additions(
+            prompt_batch=[prompt],
+            model=model,
+            activation_additions=x_vectors,
+            addition_location="front",
+            res_stream_slice=slice(None),
+            **sampling_kwargs,
+        )
+        output = mod_df.loc[0, "prompts"] + mod_df.loc[0, "completions"]
+        outputs.append(output)
+        print(act_name, output)
+    return outputs
+
+if __name__ == "__main__":
+    hf_model = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Meta-Llama-3-8B-Instruct",
+        device_map="auto",
+        torch_dtype="float16",
+    )
+    model = HookedSAETransformer.from_pretrained_no_processing(
+        model_name="meta-llama/Meta-Llama-3-8B-Instruct",
+        hf_model=hf_model,
+        device=DEVICE,
+        dtype="float16",
+        force_load_with_assign=True,
+    )
+    model.eval()
+    print("Finished loading.")
+
+    prompt = "User: Can you help me with my homework? Assistant:"
+    generate_steered(model, prompt)