Upload ortho_cookbook

ortho_cookbook.ipynb

@@ -0,0 +1,1129 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "82acAhWYGIPx"
+      },
+      "source": [
+        "# Demo of bypassing refusal"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "j7hOtw7UHXdD"
+      },
+      "source": [
+        "This notebook demonstrates our method for bypassing refusal, levaraging the insight that refusal is mediated by a 1-dimensional subspace. We recommend reading the [research post](https://www.lesswrong.com/posts/jGuXSZgv6qfdhMCuJ/refusal-in-llms-is-mediated-by-a-single-direction) for a more thorough explanation.\n",
+        "\n",
+        "Modified by FailSpy for easier usage"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fcxHyDZw6b86"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dLeei4-T6Wef"
+      },
+      "outputs": [],
+      "source": [
+        "%%capture\n",
+        "!pip install transformers transformers_stream_generator tiktoken transformer_lens einops jaxtyping colorama scikit-learn"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "_vhhwl-2-jPg"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "import functools\n",
+        "import einops\n",
+        "import requests\n",
+        "import pandas as pd\n",
+        "import io\n",
+        "import textwrap\n",
+        "import gc\n",
+        "\n",
+        "from datasets import load_dataset\n",
+        "from sklearn.model_selection import train_test_split\n",
+        "from tqdm import tqdm\n",
+        "from torch import Tensor\n",
+        "from typing import List, Callable\n",
+        "from transformer_lens import HookedTransformer, utils\n",
+        "from transformer_lens.hook_points import HookPoint\n",
+        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+        "from jaxtyping import Float, Int\n",
+        "from colorama import Fore\n",
+        "\n",
+        "# We turn automatic differentiation off, to save GPU memory, as this notebook focuses on model inference not model training. (credit: Undi95)\n",
+        "torch.set_grad_enabled(False)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Load harmful/harmless prompts datasets, and chat template\n",
+        "Please note you may need to change the `CHAT_TEMPLATE` for whichever model you're using"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Rth8yvLZJsXs"
+      },
+      "outputs": [],
+      "source": [
+        "def get_harmful_instructions():\n",
+        "    url = 'https://raw.githubusercontent.com/llm-attacks/llm-attacks/main/data/advbench/harmful_behaviors.csv'\n",
+        "    response = requests.get(url)\n",
+        "    dataset = pd.read_csv(io.StringIO(response.content.decode('utf-8')))\n",
+        "    instructions = dataset['goal'].tolist()\n",
+        "    train, test = train_test_split(instructions, test_size=0.2, random_state=42)\n",
+        "    return train, test\n",
+        "\n",
+        "def get_harmless_instructions():\n",
+        "    hf_path = 'tatsu-lab/alpaca'\n",
+        "    dataset = load_dataset(hf_path)\n",
+        "    # filter for instructions that do not have inputs\n",
+        "    instructions = []\n",
+        "    for i in range(len(dataset['train'])):\n",
+        "        if dataset['train'][i]['input'].strip() == '':\n",
+        "            instructions.append(dataset['train'][i]['instruction'])\n",
+        "    train, test = train_test_split(instructions, test_size=0.2, random_state=42)\n",
+        "    return train, test\n",
+        "\n",
+        "harmful_inst_train, harmful_inst_test = get_harmful_instructions()\n",
+        "harmless_inst_train, harmless_inst_test = get_harmless_instructions()\n",
+        "\n",
+        "CHAT_TEMPLATE = \"\"\"<|user|>\\n{instruction}<|end|>\\n<|assistant|>\"\"\" # phi-3 chat template\n",
+        "\n",
+        "CHAT_TEMPLATE = \"\"\"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\\n{instruction}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\"\"\" # llama-3 chat template\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6ZOoJagxD49V"
+      },
+      "source": [
+        "### Load model (multi-GPU support)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 191,
+          "referenced_widgets": [
+            "ad063e2c68a44f009bfab68c141c09be",
+            "89ee88168c474e9fbcf4a17f1483eff4",
+            "3877270cf4bc42a9b6142cce7a5d8c54",
+            "9a5611a341ed4673aaaf2f463f685d7c",
+            "a2de63dfbd6c485e841c6fcd1fefe451",
+            "c362d50107dd4a2db0d1a79da2af8d57",
+            "ffa85c694b694425999187b346c7ecfe",
+            "ec8f6f360a2243b0ac98d34e825ba378",
+            "f2ee188bfaa84e9680dbc296b1adbef6",
+            "e973493cd6d14381bb4ad2f82417e8a9",
+            "89797f6e82104058af92e3ceb094af66"
+          ]
+        },
+        "id": "Vnp65Vsg5x-5",
+        "outputId": "25fb5805-fe31-44b0-8f73-6fabc230d261"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_PATH = 'meta-llama/Meta-Llama-3-70B-Instruct'\n",
+        "\n",
+        "# little hack/tip:\n",
+        "# if you're dealing with a fine-tuned model of a \"supported\" model by transformer lens\n",
+        "# you can replicate the 'model path' of the supported model in your working directory\n",
+        "# e.g. rename the folder of 'dolphin-2.9-llama3-8b' to 'Meta-Llama-3-70B-Instruct', and put that into a folder called 'meta-llama'\n",
+        "# now transformers will accept 'meta-llama/Meta-Llama-3-70B-Instruct' as the model path for the model you're using, AND you don't have to add the model name to HookedTransformers\n",
+        "# make sure the model architecture and configs really do match though! \n",
+        "\n",
+        "model = HookedTransformer.from_pretrained_no_processing(\n",
+        "    MODEL_PATH,\n",
+        "    #local_files_only=True, # you can use local_files_only=True as a kwarg to from_pretrained_no_processing to enforce using the model from a local directory\n",
+        "    dtype=torch.bfloat16,   # you may want to try full precision if you can. bfloat16 is a good compromise though, but may not work in certain conditions or on certain hardware. DYOR\n",
+        "    default_padding_side='left'\n",
+        ")\n",
+        "\n",
+        "\n",
+        "model.tokenizer.padding_side = 'left'\n",
+        "model.tokenizer.pad_token = model.tokenizer.eos_token"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rF7e-u20EFTe"
+      },
+      "source": [
+        "#### Load model utility stuff"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "def tokenize_instructions_chat(\n",
+        "    tokenizer: AutoTokenizer,\n",
+        "    instructions: List[str]\n",
+        ") -> Int[Tensor, 'batch_size seq_len']:\n",
+        "    prompts = [CHAT_TEMPLATE.format(instruction=instruction) for instruction in instructions]\n",
+        "    return tokenizer(prompts, padding=True, truncation=False, return_tensors=\"pt\").input_ids\n",
+        "\n",
+        "tokenize_instructions_fn = functools.partial(tokenize_instructions_chat, tokenizer=model.tokenizer)\n",
+        "def _generate_with_hooks(\n",
+        "    model: HookedTransformer,\n",
+        "    toks: Int[Tensor, 'batch_size seq_len'],\n",
+        "    max_tokens_generated: int = 64,\n",
+        "    fwd_hooks = [],\n",
+        ") -> List[str]:\n",
+        "    all_toks = torch.zeros((toks.shape[0], toks.shape[1] + max_tokens_generated), dtype=torch.long, device=toks.device)\n",
+        "    all_toks[:, :toks.shape[1]] = toks\n",
+        "    for i in range(max_tokens_generated):\n",
+        "        with model.hooks(fwd_hooks=fwd_hooks):\n",
+        "            logits = model(all_toks[:, :-max_tokens_generated + i])\n",
+        "            next_tokens = logits[:, -1, :].argmax(dim=-1) # greedy sampling (temperature=0)\n",
+        "            all_toks[:,-max_tokens_generated+i] = next_tokens\n",
+        "    return model.tokenizer.batch_decode(all_toks[:, toks.shape[1]:], skip_special_tokens=True)\n",
+        "\n",
+        "def get_generations(\n",
+        "    model: HookedTransformer,\n",
+        "    instructions: List[str],\n",
+        "    tokenize_instructions_fn: Callable[[List[str]], Int[Tensor, 'batch_size seq_len']],\n",
+        "    fwd_hooks = [],\n",
+        "    max_tokens_generated: int = 64,\n",
+        "    batch_size: int = 4,\n",
+        ") -> List[str]:\n",
+        "    generations = []\n",
+        "    for i in tqdm(range(0, len(instructions), batch_size)):\n",
+        "        toks = tokenize_instructions_fn(instructions=instructions[i:i+batch_size])\n",
+        "        generation = _generate_with_hooks(\n",
+        "            model,\n",
+        "            toks,\n",
+        "            max_tokens_generated=max_tokens_generated,\n",
+        "            fwd_hooks=fwd_hooks,\n",
+        "        )\n",
+        "        generations.extend(generation)\n",
+        "    return generations\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "#### Useful cell for clearing RAM/VRAM"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# run this cell if you bump into memory issues\n",
+        "try:\n",
+        "    del harmless_logits\n",
+        "except Exception:\n",
+        "    pass\n",
+        "try:\n",
+        "    del harmful_logits\n",
+        "except Exception:\n",
+        "    pass\n",
+        "gc.collect(); torch.cuda.empty_cache()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "W9O8dm0_EQRk"
+      },
+      "source": [
+        "## Finding potential \"refusal directions\" (batched)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "vscode": {
+          "languageId": "julia"
+        }
+      },
+      "outputs": [],
+      "source": [
+        "harmful = {}\n",
+        "harmless = {}\n",
+        "\n",
+        "# may want to spare your RAM and cycles here. can use '32' here instead or something like the paper\n",
+        "N_INST_TRAIN = min(len(harmful_inst_train), len(harmless_inst_train)) \n",
+        "\n",
+        "# load the full training set here to align all the dimensions\n",
+        "toks = tokenize_instructions_fn(instructions=harmful_inst_train[:N_INST_TRAIN]+harmless_inst_train[:N_INST_TRAIN])\n",
+        "harmful_toks,harmless_toks = toks.split(N_INST_TRAIN)\n",
+        "\n",
+        "batch_size = 48 # adjust this based on available VRAM\n",
+        "\n",
+        "for i in tqdm(range(0, N_INST_TRAIN // batch_size + (N_INST_TRAIN % batch_size > 0))):\n",
+        "    id = i*batch_size\n",
+        "    e = min(N_INST_TRAIN,id+batch_size)\n",
+        "\n",
+        "    # run the models on harmful and harmless prompts, cache their activations separately.\n",
+        "    harmful_logits, harmful_cache = model.run_with_cache(harmful_toks[id:e], names_filter=lambda hook_name: 'resid' in hook_name, device='cpu', reset_hooks_end=True)\n",
+        "    harmless_logits, harmless_cache = model.run_with_cache(harmless_toks[id:e], names_filter=lambda hook_name: 'resid' in hook_name, device='cpu', reset_hooks_end=True)\n",
+        "    \n",
+        "    for key in harmful_cache:\n",
+        "        if key not in harmful:\n",
+        "            harmful[key] = [harmful_cache[key]]\n",
+        "            harmless[key] = [harmless_cache[key]]\n",
+        "        else:\n",
+        "            harmful[key].append(harmful_cache[key])\n",
+        "            harmless[key].append(harmless_cache[key])\n",
+        "\n",
+        "    # force Python & PyTorch to clear GPU and CPU RAM where possible\n",
+        "    del harmful_logits, harmless_logits, harmful_cache, harmless_cache\n",
+        "    gc.collect()\n",
+        "    torch.cuda.empty_cache()\n",
+        "\n",
+        "harmful = {k:torch.cat(v) for k,v in harmful.items()}\n",
+        "harmless = {k:torch.cat(v) for k,v in harmless.items()}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Compute activations into refusal directions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tqD5E8Vc_w5d"
+      },
+      "outputs": [],
+      "source": [
+        "# compute difference of means between harmful and harmless activations at intermediate layers\n",
+        "\n",
+        "def get_act_idx(cache_dict, act_name, layer):\n",
+        "    key = (act_name, layer,)\n",
+        "    return cache_dict[utils.get_act_name(*key)]\n",
+        "\n",
+        "activation_layers = ['resid_pre', 'resid_mid', 'resid_post']\n",
+        "\n",
+        "activation_refusals = {k:[] for k in activation_layers}\n",
+        "\n",
+        "for layer_num in range(1,model.cfg.n_layers):\n",
+        "    pos = -1\n",
+        "\n",
+        "    for layer in activation_layers:\n",
+        "        harmful_mean_act = get_act_idx(harmful, layer, layer_num)[:, pos, :].mean(dim=0)\n",
+        "        harmless_mean_act = get_act_idx(harmless, layer, layer_num)[:, pos, :].mean(dim=0)\n",
+        "    \n",
+        "        refusal_dir = harmful_mean_act - harmless_mean_act\n",
+        "        refusal_dir = refusal_dir / refusal_dir.norm()\n",
+        "        activation_refusals[layer].append(refusal_dir)\n",
+        "\n",
+        "# save to file so you don't have to re-build later\n",
+        "torch.save(activation_refusals, 'refusal_dirs.pth')\n",
+        "refusal_dirs = activation_refusals"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Alternatively, load a previous attempt's pth:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "activation_refusals = torch.load('refusal_dirs.pth')\n",
+        "refusal_dirs = activation_refusals"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2EoxY5i1CWe3"
+      },
+      "source": [
+        "## Ablate \"refusal direction\" via inference-time intervention\n",
+        "\n",
+        "Given a \"refusal direction\" $\\widehat{r} \\in \\mathbb{R}^{d_{\\text{model}}}$ with unit norm, we can ablate this direction from the model's activations $a_{l}$:\n",
+        "$${a}_{l}' \\leftarrow a_l - (a_l \\cdot \\widehat{r}) \\widehat{r}$$\n",
+        "\n",
+        "By performing this ablation on all intermediate activations, we enforce that the model can never express this direction (or \"feature\")."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### \"Score\" layer activation diffs\n",
+        "This is a place with room for improvement in methodology. For now, I'm just doing a rudimentary sort based on difference distance average to find dirs with the most \"change\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Get all calculated potential refusal dirs, sort them in Descending order (reverse) based on their mean()\n",
+        "activation_layers = ['resid_pre', 'resid_mid', 'resid_post'] # you can use a subset of these if you don't think certain activations are promising\n",
+        "\n",
+        "activation_layers = ['resid_pre'] # this is usually good enough, though if you've got the compute to spare...\n",
+        "activation_scored = sorted([activation_refusals[layer][l-1] for l in range(1,model.cfg.n_layers) for layer in activation_layers], key = lambda x: abs(x.mean()), reverse=True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "#### Model ablation testing/brute-forcing the best refusal dir"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "##### Inference-time intervention hook:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "26rf-yncB2PT"
+      },
+      "outputs": [],
+      "source": [
+        "def direction_ablation_hook(\n",
+        "    activation: Float[Tensor, \"... d_act\"],\n",
+        "    hook: HookPoint,\n",
+        "    direction: Float[Tensor, \"d_act\"]\n",
+        "):\n",
+        "    if activation.device != direction.device:\n",
+        "        direction = direction.to(activation.device)\n",
+        "    proj = einops.einsum(activation, direction.view(-1, 1), '... d_act, d_act single -> ... single') * direction\n",
+        "    return activation - proj"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "##### Testing baseline, can skip if you don't care (it will basically just be refusals with a regular model :P)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "N_INST_TEST = 12\n",
+        "baseline_generations = get_generations(model, harmful_inst_test[:N_INST_TEST], tokenize_instructions_fn, fwd_hooks=[])\n",
+        "for gen in baseline_generations:\n",
+        "    print(gen)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "##### Evaluating layers defined earlier (needs human evaluation to determine best layer for refusal inhibition)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "sR1G5bXoEDty",
+        "outputId": "2e610278-3d4c-473e-850f-a9b28d0b13f9"
+      },
+      "outputs": [],
+      "source": [
+        "if \"N_INST_TEST\" not in locals() or not N_INST_TEST:\n",
+        "    N_INST_TEST = 12 # you may want to evaluate more at the cost of additional compute time. by default, batches are size of 4, so I'd recommend making it a multiple of 4.\n",
+        "EVAL_N = 10 # Evaluate how many of the top N potential dirs\n",
+        "evals = []\n",
+        "for refusal_dir in tqdm(activation_scored[:EVAL_N]):\n",
+        "    intervention_layers = list(range(model.cfg.n_layers)) # all layers\n",
+        "\n",
+        "    hook_fn = functools.partial(direction_ablation_hook,direction=refusal_dir)\n",
+        "    fwd_hooks = [(utils.get_act_name(act_name, l), hook_fn) for l in intervention_layers for act_name in ['resid_pre', 'resid_mid', 'resid_post']]\n",
+        "\n",
+        "    intervention_generations = get_generations(model, harmful_inst_test[:N_INST_TEST], tokenize_instructions_fn, fwd_hooks=fwd_hooks)\n",
+        "    evals.append(intervention_generations)\n",
+        "    \n",
+        "    #print(intervention_generations) # if you want to watch it as it goes"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "#### Present evals to clever pre-trained non-refusing human"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "pxbJr4vCFCOL",
+        "outputId": "47810f58-813e-4014-fee8-cb1406e318c8"
+      },
+      "outputs": [],
+      "source": [
+        "for instruction in range(N_INST_TEST):\n",
+        "    if 'baseline_generations' in locals() and baseline_generations and len(baseline_generations) > instruction:\n",
+        "        print(f\"INSTRUCTION {instruction}: {repr(harmful_inst_test[instruction])}\")\n",
+        "        print(Fore.GREEN + f\"BASELINE COMPLETION:\")\n",
+        "        print(textwrap.fill(repr(baseline_generations[instruction]), width=100, initial_indent='\\t', subsequent_indent='\\t'))\n",
+        "    for layer_candidate in range(EVAL_N):\n",
+        "        if len(evals) > layer_candidate and len(evals[layer_candidate]) > instruction:\n",
+        "            print(Fore.RED + f\"LAYER CANDIDATE #{layer_candidate} INTERVENTION COMPLETION:\")\n",
+        "            print(textwrap.fill(repr(evals[layer_candidate][instruction]), width=100, initial_indent='\\t', subsequent_indent='\\t'))\n",
+        "    print(Fore.RESET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "t9KooaWaCDc_"
+      },
+      "source": [
+        "## Orthogonalize weights w.r.t. \"refusal direction\"\n",
+        "\n",
+        "We can implement the intervention equivalently by directly orthogonalizing the weight matrices that write to the residual stream with respect to the refusal direction $\\widehat{r}$:\n",
+        "$$W_{\\text{out}}' \\leftarrow W_{\\text{out}} - \\widehat{r}\\widehat{r}^{\\mathsf{T}} W_{\\text{out}}$$\n",
+        "\n",
+        "By orthogonalizing these weight matrices, we enforce that the model is unable to write direction $r$ to the residual stream at all!\n",
+        "\n",
+        "This is basically how you finalize your layers' weights to represent your orthogonalization for a saved model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Choose your fighter (favorite, ideally non-refusing layer)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "layer_candidate = 2 # e.g. you should choose based on the layer you think aligns to the behavior you like\n",
+        "refusal_dir = activation_scored[layer_candidate]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Write ortho'd weights into model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "def get_orthogonalized_matrix(matrix: Float[Tensor, '... d_model'], vec: Float[Tensor, 'd_model']) -> Float[Tensor, '... d_model']:\n",
+        "    proj = einops.einsum(matrix, vec.view(-1, 1), '... d_model, d_model single -> ... single') * vec\n",
+        "    return matrix - proj"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "GC7cpMXZCG64"
+      },
+      "outputs": [],
+      "source": [
+        "if refusal_dir.device != model.W_E.device:\n",
+        "    refusal_dir = refusal_dir.to(model.W_E.device)\n",
+        "model.W_E.data = get_orthogonalized_matrix(model.W_E, refusal_dir)\n",
+        "\n",
+        "for block in tqdm(model.blocks):\n",
+        "    if refusal_dir.device != block.attn.W_O.device:\n",
+        "        refusal_dir = refusal_dir.to(block.attn.W_O.device)\n",
+        "    block.attn.W_O.data = get_orthogonalized_matrix(block.attn.W_O, refusal_dir)\n",
+        "    block.mlp.W_out.data = get_orthogonalized_matrix(block.mlp.W_out, refusal_dir)\n",
+        "\n",
+        "# save your refusal_dir of choice separately to a file\n",
+        "torch.save(refusal_dir,\"ablation.pth\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Verify model weights are adjusted to match ablation (skippable)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "1Y-qtouNGf3t",
+        "outputId": "5f946460-4ed4-4cf1-e53d-809ebb880f6e"
+      },
+      "outputs": [],
+      "source": [
+        "orthogonalized_generations = get_generations(model, harmful_inst_test[:N_INST_TEST], tokenize_instructions_fn, fwd_hooks=[])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "r68O4_4DG3P7",
+        "outputId": "97eeb477-bfd1-4521-8c32-4657d99f3e0c"
+      },
+      "outputs": [],
+      "source": [
+        "for i in range(N_INST_TEST):\n",
+        "    if 'baseline_generations' in locals() and baseline_generations and len(baseline_generations) > i:\n",
+        "        print(f\"INSTRUCTION {i}: {repr(harmful_inst_test[i])}\")\n",
+        "        print(Fore.GREEN + f\"BASELINE COMPLETION:\")\n",
+        "        print(textwrap.fill(repr(baseline_generations[i]), width=100, initial_indent='\\t', subsequent_indent='\\t'))\n",
+        "    print(Fore.RED + f\"INTERVENTION COMPLETION:\")\n",
+        "    print(textwrap.fill(repr(evals[layer_candidate][i]), width=100, initial_indent='\\t', subsequent_indent='\\t'))\n",
+        "    print(Fore.MAGENTA + f\"ORTHOGONALIZED COMPLETION:\")\n",
+        "    print(textwrap.fill(repr(orthogonalized_generations[i]), width=100, initial_indent='\\t', subsequent_indent='\\t'))\n",
+        "    print(Fore.RESET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "vscode": {
+          "languageId": "julia"
+        }
+      },
+      "source": [
+        "## Save your unruly model\n",
+        "This is where you'll need to consult with the original structure of the model you're generating. Below is converting Phi-3 and Llama-3 examples, but you'll need to do differently for others. Or if you just want a \"no thinking\" save, you can use the pytorch save below. Be aware that the structure output by that is not directly convertable however."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Simple PyTorch save! (easiest, but least portable)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "torch.save(model, \"pytorch_model.bin\") # can name it whatever you want, and then reload it"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### Converting models back to HF safetensors (harder)\n",
+        "Do note that we only adjust a couple layers in get_orthogonalized_matrix(), so only need to copy 1 + (2*n_layers) over to the original trained model, not the whole lot.\n",
+        "\n",
+        "You can look to TransformerLens's source code `loading_from_pretrained.py` to see how the layers get converted in. e.g. https://github.com/neelnanda-io/TransformerLens/blob/main/transformer_lens/loading_from_pretrained.py#L1746-L1833 is `convert_llama_weights`, so you can just reverse the steps for the layers that you alter\n",
+        "\n",
+        "References to convert functions per model:\n",
+        "https://github.com/neelnanda-io/TransformerLens/blob/main/transformer_lens/loading_from_pretrained.py#L1475-L1504"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# this is probably useful for any conversion\n",
+        "cfg = model.cfg\n",
+        "\n",
+        "state_dict = model.state_dict()\n",
+        "\n",
+        "hf_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH,torch_dtype=torch.bfloat16) # load the original model as a regular unhooked Transformer -- don't need to load it into GPU as it's just for saving\n",
+        "lm_model = hf_model.model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "#### Llama-3 conversion"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "exUh3PEHRe9x"
+      },
+      "outputs": [],
+      "source": [
+        "lm_model.embed_tokens.weight = torch.nn.Parameter(state_dict[\"embed.W_E\"].cpu())\n",
+        "\n",
+        "for l in range(cfg.n_layers):\n",
+        "    lm_model.layers[l].self_attn.o_proj.weight = torch.nn.Parameter(einops.rearrange(state_dict[f\"blocks.{l}.attn.W_O\"], \"n h m->m (n h)\", n=cfg.n_heads).contiguous())\n",
+        "    lm_model.layers[l].mlp.down_proj.weight = torch.nn.Parameter(torch.transpose(state_dict[f\"blocks.{l}.mlp.W_out\"].contiguous(),0,1))\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "#### Phi-3 conversion"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "\n",
+        "lm_model.embed_tokens.weight = state_dict[\"embed.W_E\"]\n",
+        "\n",
+        "for l in range(cfg.n_layers):\n",
+        "\n",
+        "    W_O = einops.rearrange(\n",
+        "        state_dict[f\"blocks.{l}.attn.W_O\"], \"n_head d_head d_model -> d_model (n_head d_head)\", n_head=cfg.n_heads\n",
+        "    )\n",
+        "    lm_model.layers[l].self_attn.o_proj.weight = torch.nn.Parameter(W_O.contiguous())\n",
+        "\n",
+        "    lm_model.layers[l].mlp.down_proj.weight = torch.nn.Parameter(torch.transpose(state_dict[f\"blocks.{l}.mlp.W_out\"].cpu(), 0, 1).contiguous())\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "#### Save converted model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "hf_model.save(\"path/to/my_saved_based_model_dir/\")"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": [],
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "3877270cf4bc42a9b6142cce7a5d8c54": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_ec8f6f360a2243b0ac98d34e825ba378",
+            "max": 2,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_f2ee188bfaa84e9680dbc296b1adbef6",
+            "value": 2
+          }
+        },
+        "89797f6e82104058af92e3ceb094af66": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "89ee88168c474e9fbcf4a17f1483eff4": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_c362d50107dd4a2db0d1a79da2af8d57",
+            "placeholder": "​",
+            "style": "IPY_MODEL_ffa85c694b694425999187b346c7ecfe",
+            "value": "Loading checkpoint shards: 100%"
+          }
+        },
+        "9a5611a341ed4673aaaf2f463f685d7c": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_e973493cd6d14381bb4ad2f82417e8a9",
+            "placeholder": "​",
+            "style": "IPY_MODEL_89797f6e82104058af92e3ceb094af66",
+            "value": " 2/2 [00:18&lt;00:00,  8.85s/it]"
+          }
+        },
+        "a2de63dfbd6c485e841c6fcd1fefe451": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "ad063e2c68a44f009bfab68c141c09be": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_89ee88168c474e9fbcf4a17f1483eff4",
+              "IPY_MODEL_3877270cf4bc42a9b6142cce7a5d8c54",
+              "IPY_MODEL_9a5611a341ed4673aaaf2f463f685d7c"
+            ],
+            "layout": "IPY_MODEL_a2de63dfbd6c485e841c6fcd1fefe451"
+          }
+        },
+        "c362d50107dd4a2db0d1a79da2af8d57": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "e973493cd6d14381bb4ad2f82417e8a9": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "ec8f6f360a2243b0ac98d34e825ba378": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "f2ee188bfaa84e9680dbc296b1adbef6": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "ffa85c694b694425999187b346c7ecfe": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        }
+      }
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}