1import gc
2import random
3
4import torch
5from tqdm import tqdm
6from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
7
8MODEL_ID = "T145/ZEUS-8B-V2"
9
10# More samples can help find the direction better.
11NUM_PROMPT_SAMPLES = 32
12
13# Used to skip the first and last layers for the modifications.
14SKIP_BEGIN_LAYERS = 1
15SKIP_END_LAYERS = 1
16
17# The layer we will use for the refusal_dir calculation will be floor(LAYER_FRACTION_TO_USE * model.layers).
18LAYER_FRACTION_TO_USE = 0.6
19
20# Use a negative scale_factor to "induce" and a positive scale_factor of < 1 to "ablate" less.
21SCALE_FACTOR = 1.0
22
23torch.inference_mode()
24torch.set_default_device("cpu")
25torch.set_grad_enabled(False)
26
27# Load the model on the GPU in quantized type if we can.
28model = AutoModelForCausalLM.from_pretrained(
29 MODEL_ID,
30 trust_remote_code=True,
31 torch_dtype=torch.float16,
32 quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16),
33 low_cpu_mem_usage=True,
34 device_map='auto'
35)
36model.requires_grad_(False)
37
38tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
39layer_idx = int(len(model.model.layers) * LAYER_FRACTION_TO_USE)
40
41print("Layer index for refusal direction: " + str(layer_idx))
42
43with open("harmful.txt", "r", encoding="utf-8") as f:
44 harmful = f.readlines()
45
46with open("harmless.txt", "r", encoding="utf-8") as f:
47 harmless = f.readlines()
48
49harmful_instructions = random.sample(harmful, min(NUM_PROMPT_SAMPLES, len(harmful)))
50harmless_instructions = random.sample(harmless, min(NUM_PROMPT_SAMPLES, len(harmless)))
51
52harmful_toks = [
53 tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}], add_generation_prompt=True, tokenize=False,
54 return_tensors="pt") for insn in harmful_instructions]
55harmless_toks = [
56 tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}], add_generation_prompt=True, tokenize=False,
57 return_tensors="pt") for insn in harmless_instructions]
58
59bar_generate = tqdm(total = len(harmful_instructions) + len(harmless_instructions), desc = "Generating samples")
60
61# Only return the final hidden state of the layer we care about, and use 'cpu' to save VRAM.
62def generate(toks):
63 inputs = tokenizer(toks, return_tensors="pt", padding=True)
64 inputs = inputs.to(model.device)
65 output = model.generate(
66 inputs['input_ids'],
67 use_cache=False,
68 max_new_tokens=1,
69 return_dict_in_generate=True,
70 output_hidden_states=True,
71 attention_mask=inputs["attention_mask"],
72 pad_token_id=tokenizer.eos_token_id
73 )
74 bar_generate.update(n=1)
75 return output.hidden_states[0][layer_idx][:, -1, :].to('cpu') # Final hidden state = -1.
76
77harmful_hidden = [generate(toks) for toks in harmful_toks]
78harmless_hidden = [generate(toks) for toks in harmless_toks]
79
80bar_generate.close()
81
82harmful_mean = torch.stack(harmful_hidden).mean(dim=0)
83harmless_mean = torch.stack(harmless_hidden).mean(dim=0)
84
85refusal_dir = harmful_mean - harmless_mean
86refusal_dir = refusal_dir.squeeze() / refusal_dir.norm()
87
88torch.save(refusal_dir, MODEL_ID.replace("/", "_") + "_refusal_dir.pt")
89
90# Free memory
91del model
92gc.collect()
93torch.cuda.empty_cache()
94
95# Reload the model in CPU memory with bfloat16 data type
96model = AutoModelForCausalLM.from_pretrained(
97 MODEL_ID,
98 trust_remote_code=True,
99 torch_dtype=torch.bfloat16,
100 low_cpu_mem_usage=True,
101 device_map='cpu'
102)
103model.requires_grad_(False)
104
105# Make sure it's on the 'cpu' device.
106if refusal_dir.device != model.device:
107 refusal_dir = refusal_dir.to(model.device)
108
109# Get the language model component and check it's as expected.
110lm_model = model.model
111assert hasattr(lm_model, 'layers'), "The model does not have the expected structure."
112
113# Check the ranges are valid.
114num_layers = len(lm_model.layers)
115assert SKIP_BEGIN_LAYERS >= 0, "SKIP_BEGIN_LAYERS must be >= 0."
116assert SKIP_END_LAYERS >= 0, "SKIP_END_LAYERS must be >= 0."
117assert SKIP_BEGIN_LAYERS + SKIP_END_LAYERS < num_layers, "SKIP_BEGIN_LAYERS + SKIP_END_LAYERS must be < num_layers."
118
119bar_layers = tqdm(total= (num_layers - (SKIP_BEGIN_LAYERS + SKIP_END_LAYERS)) * 2, desc = "Modifying tensors")
120
121# NOTE: Use a negative scale_factor to "induce" and a positive scale_factor of < 1 to "ablate" less.
122def modify_tensor(tensor_data, refusal_dir, scale_factor: float = 1.0):
123 assert scale_factor <= 1.0, "Using a scale_factor of > 1 doesn't make sense..."
124 tensor_float = tensor_data.to(torch.bfloat16)
125 refusal_dir_float = refusal_dir.to(torch.bfloat16)
126 tensor_float -= scale_factor * torch.matmul(torch.outer(refusal_dir_float, refusal_dir_float), tensor_float)
127 tensor_modified = tensor_float.to(torch.bfloat16)
128 bar_layers.update(1)
129 return torch.nn.Parameter(tensor_modified)
130
131# Modify the 'self_attn.o_proj.weight' and 'mlp.down_proj.weight' in each chosen layer.
132# NOTE: These tensors names are speific to "llama" and may need changing.
133# - See here for others: https://github.com/arcee-ai/mergekit/tree/main/mergekit/_data/architectures
134for layer_idx in range(SKIP_BEGIN_LAYERS, num_layers - SKIP_END_LAYERS):
135 lm_model.layers[layer_idx].self_attn.o_proj.weight = modify_tensor(
136 lm_model.layers[layer_idx].self_attn.o_proj.weight.data, refusal_dir, SCALE_FACTOR
137 )
138 lm_model.layers[layer_idx].mlp.down_proj.weight = modify_tensor(
139 lm_model.layers[layer_idx].mlp.down_proj.weight.data, refusal_dir, SCALE_FACTOR
140 )
141
142bar_layers.close()
143
144print("Saving modified model (with original tokenizer)...")
145
146FIXED_ID = f"{MODEL_ID}-abliterated"
147model.save_pretrained(FIXED_ID)
148tokenizer.save_pretrained(FIXED_ID)
