mirror of
https://github.com/Sumandora/remove-refusals-with-transformers.git
synced 2026-07-10 01:18:39 +00:00
7786b0a8c5
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...
97 lines
3.5 KiB
Python
97 lines
3.5 KiB
Python
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()
|
|
|
|
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
|
|
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,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
|
output_attentions: Optional[bool] = False,
|
|
use_cache: Optional[bool] = False,
|
|
cache_position: Optional[torch.LongTensor] = None,
|
|
**kwargs,
|
|
):
|
|
assert not output_attentions
|
|
|
|
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,)
|
|
|
|
return outputs
|
|
|
|
|
|
# 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())
|
|
|
|
# bruh
|
|
if hasattr(model, "config") and hasattr(model.config, "num_hidden_layers"):
|
|
model.config.num_hidden_layers *= 2
|
|
|
|
conversation = []
|
|
|
|
streamer = TextStreamer(tokenizer)
|
|
|
|
print(f"Chat with {MODEL_ID}:")
|
|
while True:
|
|
prompt = input()
|
|
conversation.append({"role": "user", "content": prompt})
|
|
toks = tokenizer.apply_chat_template(conversation=conversation,
|
|
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)})
|