Update to latest version of HF Transformers (Fixes Sumandora#4)

I hope this didn't break backwards compatibility in ways beyond what I
can test.

Anyways, lots of new stuff in HF Transformers, which make this quite a
lot harder than it was before. I'm not sure if I can continue supporting
this, if it this goes on. Perhaps they should think of a canonical way
of doing this...

Author: Sumandora

.gitignore

@@ -0,0 +1 @@
+*.pt

compute_refusal_dir.py

@@ -1,24 +1,27 @@
-import jaxtyping
-
 import random
 
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig
-
-import einops
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 from tqdm import tqdm
 
 torch.inference_mode()
 
-MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
-#MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
-#MODEL_ID = "Qwen/Qwen-1_8B-chat"
-#MODEL_ID = "google/gemma-1.1-2b-it"
-#MODEL_ID = "google/gemma-1.1-7b-it"
-#MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
-
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True, torch_dtype=torch.float16, quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16))
+MODEL_ID = "tiiuae/Falcon3-1B-Instruct"
+# MODEL_ID = "Qwen/Qwen3-1.7B"
+# MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
+# MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
+# MODEL_ID = "Qwen/Qwen-1_8B-chat"
+# MODEL_ID = "google/gemma-1.1-2b-it"
+# MODEL_ID = "google/gemma-1.1-7b-it"
+# MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID,
+                                             trust_remote_code=True,
+                                             dtype=torch.float16,
+                                             device_map="cuda",
+                                             quantization_config=BitsAndBytesConfig(load_in_4bit=True,
+                                                                                    bnb_4bit_compute_dtype=torch.float16))
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
 # settings:
@@ -39,18 +42,26 @@
 harmless_instructions = random.sample(harmless, instructions)
 
 harmful_toks = [
-    tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}], add_generation_prompt=True,
+    tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}],
+                                  add_generation_prompt=True,
                                   return_tensors="pt") for insn in harmful_instructions]
 harmless_toks = [
-    tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}], add_generation_prompt=True,
+    tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}],
+                                  add_generation_prompt=True,
                                   return_tensors="pt") for insn in harmless_instructions]
 
 max_its = instructions*2
 bar = tqdm(total=max_its)
 
+
 def generate(toks):
     bar.update(n=1)
-    return model.generate(toks.to(model.device), use_cache=False, max_new_tokens=1, return_dict_in_generate=True, output_hidden_states=True)
+    return model.generate(toks.to(model.device),
+                          use_cache=False,
+                          max_new_tokens=1,
+                          return_dict_in_generate=True,
+                          output_hidden_states=True)
+
 
 harmful_outputs = [generate(toks) for toks in harmful_toks]
 harmless_outputs = [generate(toks) for toks in harmless_toks]

inference.py

@@ -1,35 +1,50 @@
-from typing import Optional, Tuple
-
 import einops
 import jaxtyping
 import torch
 import torch.nn as nn
+from typing import Optional, Tuple
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig
+from inspect import signature
 
 torch.inference_mode()
 
-torch.set_default_device("cpu")
-
-MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
-#MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
-#MODEL_ID = "Qwen/Qwen-1_8B-chat"
-#MODEL_ID = "google/gemma-1.1-2b-it"
-#MODEL_ID = "google/gemma-1.1-7b-it"
-#MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
-
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True, device_map="cuda", quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16))
+MODEL_ID = "tiiuae/Falcon3-1B-Instruct"
+# MODEL_ID = "Qwen/Qwen3-1.7B"
+# MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
+# MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
+# MODEL_ID = "Qwen/Qwen-1_8B-chat"
+# MODEL_ID = "google/gemma-1.1-2b-it"
+# MODEL_ID = "google/gemma-1.1-7b-it"
+# MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID,
+                                             trust_remote_code=True,
+                                             dtype=torch.float16,
+                                             device_map="cuda",
+                                             quantization_config=BitsAndBytesConfig(load_in_4bit=True,
+                                                                                    bnb_4bit_compute_dtype=torch.float16))
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
 refusal_dir = torch.load(MODEL_ID.replace("/", "_") + "_refusal_dir.pt")
 
 
 def direction_ablation_hook(activation: jaxtyping.Float[torch.Tensor, "... d_act"],
                             direction: jaxtyping.Float[torch.Tensor, "d_act"]):
-    proj = einops.einsum(activation, direction.view(-1, 1), '... d_act, d_act single -> ... single') * direction
+    proj = einops.einsum(activation, direction.view(-1, 1),
+                         '... d_act, d_act single -> ... single') * direction
     return activation - proj
 
 
+# Some model developers thought it was stupid to pass a tuple of tuple of tuples around (rightfully so), but unfortunately now we have a divide
+sig = signature(model.model.layers[0].forward)
+simple = sig.return_annotation == torch.Tensor
+
+
 class AblationDecoderLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attention_type = "full_attention"
+
     def forward(
             self,
             hidden_states: torch.Tensor,
@@ -40,24 +55,32 @@ def forward(
             use_cache: Optional[bool] = False,
             cache_position: Optional[torch.LongTensor] = None,
             **kwargs,
-    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+    ):
         assert not output_attentions
 
-        ablated = direction_ablation_hook(hidden_states, refusal_dir.to(hidden_states.device)).to(hidden_states.device)
+        ablated = direction_ablation_hook(hidden_states, refusal_dir.to(
+            hidden_states.device)).to(hidden_states.device)
+
+        if simple:
+            return ablated
 
         outputs = (ablated,)
 
         if use_cache:
             outputs += (past_key_value,)
 
-        # noinspection PyTypeChecker
         return outputs
 
 
-for idx in reversed(range(len(model.model.layers))):  # for qwen 1 this needs to be changed to model.transformer.h
+# for qwen 1 this needs to be changed to model.transformer.h
+for idx in reversed(range(len(model.model.layers))):
     model.model.layers.insert(idx, AblationDecoderLayer())
 
-conversation=[]
+# bruh
+if hasattr(model, "config") and hasattr(model.config, "num_hidden_layers"):
+    model.config.num_hidden_layers *= 2
+
+conversation = []
 
 streamer = TextStreamer(tokenizer)
 
@@ -66,10 +89,9 @@ def forward(
     prompt = input()
     conversation.append({"role": "user", "content": prompt})
     toks = tokenizer.apply_chat_template(conversation=conversation,
-        add_generation_prompt=True, return_tensors="pt")
+                                         add_generation_prompt=True, return_tensors="pt")
 
     gen = model.generate(toks.to(model.device), streamer=streamer, max_new_tokens=1337)
 
     decoded = tokenizer.batch_decode(gen[0][len(toks[0]):], skip_special_tokens=True)
     conversation.append({"role": "assistant", "content": "".join(decoded)})
-