Uncensor_any_LLM_with_abliteration.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/66213432/arp-scan/blob/master/Uncensor_any_LLM_with_abliteration.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "82acAhWYGIPx"
      },
      "source": [
        "# Uncensor any LLM with abliteration\n",
        "\n",
        "> 🗣️ [Large Language Model Course](https://github.com/mlabonne/llm-course)\n",
        "\n",
        "❤️ Created by [@maximelabonne](https://twitter.com/maximelabonne)."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -qqq transformers transformers_stream_generator tiktoken transformer_lens einops jaxtyping --progress-bar off\n",
        "\n",
        "import torch\n",
        "import functools\n",
        "import einops\n",
        "import gc\n",
        "\n",
        "from datasets import load_dataset\n",
        "from tqdm import tqdm\n",
        "from torch import Tensor\n",
        "from typing import List\n",
        "from transformer_lens import HookedTransformer, utils\n",
        "from transformer_lens.hook_points import HookPoint\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "from jaxtyping import Float, Int\n",
        "from collections import defaultdict\n",
        "\n",
        "# Turn automatic differentiation off to save GPU memory (credit: Undi95)\n",
        "torch.set_grad_enabled(False)\n",
        "\n",
        "def reformat_texts(texts):\n",
        "    return [[{\"role\": \"user\", \"content\": text}] for text in texts]\n",
        "\n",
        "# Get harmful and harmless datasets\n",
        "def get_harmful_instructions():\n",
        "    dataset = load_dataset('mlabonne/harmful_behaviors')\n",
        "    return reformat_texts(dataset['train']['text']), reformat_texts(dataset['test']['text'])\n",
        "\n",
        "def get_harmless_instructions():\n",
        "    dataset = load_dataset('mlabonne/harmless_alpaca')\n",
        "    return reformat_texts(dataset['train']['text']), reformat_texts(dataset['test']['text'])\n",
        "\n",
        "harmful_inst_train, harmful_inst_test = get_harmful_instructions()\n",
        "harmless_inst_train, harmless_inst_test = get_harmless_instructions()"
      ],
      "metadata": {
        "id": "BErEJu5WVekL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "MODEL_ID = \"mlabonne/Daredevil-8B\"\n",
        "NEW_MODEL_ID = \"mlabonne/Daredevil-8B-abliterated\"\n",
        "MODEL_TYPE = \"meta-llama/Meta-Llama-3-8B-Instruct\"\n",
        "\n",
        "# Download and load model\n",
        "!git clone https://huggingface.co/{MODEL_ID} {MODEL_TYPE}\n",
        "\n",
        "# Load model and tokenizer\n",
        "model = HookedTransformer.from_pretrained_no_processing(\n",
        "    MODEL_TYPE,\n",
        "    local_files_only=True,\n",
        "    dtype=torch.bfloat16,\n",
        "    default_padding_side='left'\n",
        ")\n",
        "tokenizer = AutoTokenizer.from_pretrained(MODEL_TYPE)\n",
        "tokenizer.padding_side = 'left'\n",
        "tokenizer.pad_token = tokenizer.eos_token"
      ],
      "metadata": {
        "id": "O63-4qZw8nfY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def tokenize_instructions(tokenizer, instructions):\n",
        "    return tokenizer.apply_chat_template(\n",
        "        instructions,\n",
        "        padding=True,\n",
        "        truncation=False,\n",
        "        return_tensors=\"pt\",\n",
        "        return_dict=True,\n",
        "        add_generation_prompt=True,\n",
        "    ).input_ids\n",
        "\n",
        "n_inst_train = min(256, len(harmful_inst_train), len(harmless_inst_train))\n",
        "\n",
        "# Tokenize datasets\n",
        "harmful_tokens = tokenize_instructions(\n",
        "    tokenizer,\n",
        "    instructions=harmful_inst_train[:n_inst_train],\n",
        ")\n",
        "harmless_tokens = tokenize_instructions(\n",
        "    tokenizer,\n",
        "    instructions=harmless_inst_train[:n_inst_train],\n",
        ")"
      ],
      "metadata": {
        "id": "OLePKlGnKVU1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define batch size based on available VRAM\n",
        "batch_size = 32\n",
        "\n",
        "# Initialize defaultdicts to store activations\n",
        "harmful = defaultdict(list)\n",
        "harmless = defaultdict(list)\n",
        "\n",
        "# Process the training data in batches\n",
        "num_batches = (n_inst_train + batch_size - 1) // batch_size\n",
        "for i in tqdm(range(num_batches)):\n",
        "    print(i)\n",
        "    start_idx = i * batch_size\n",
        "    end_idx = min(n_inst_train, start_idx + batch_size)\n",
        "\n",
        "    # Run models on harmful and harmless prompts, cache activations\n",
        "    harmful_logits, harmful_cache = model.run_with_cache(\n",
        "        harmful_tokens[start_idx:end_idx],\n",
        "        names_filter=lambda hook_name: 'resid' in hook_name,\n",
        "        device='cpu',\n",
        "        reset_hooks_end=True\n",
        "    )\n",
        "    harmless_logits, harmless_cache = model.run_with_cache(\n",
        "        harmless_tokens[start_idx:end_idx],\n",
        "        names_filter=lambda hook_name: 'resid' in hook_name,\n",
        "        device='cpu',\n",
        "        reset_hooks_end=True\n",
        "    )\n",
        "\n",
        "    # Collect and store the activations\n",
        "    for key in harmful_cache:\n",
        "        harmful[key].append(harmful_cache[key])\n",
        "        harmless[key].append(harmless_cache[key])\n",
        "\n",
        "    # Flush RAM and VRAM\n",
        "    del harmful_logits, harmless_logits, harmful_cache, harmless_cache\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "# Concatenate the cached activations\n",
        "harmful = {k: torch.cat(v) for k, v in harmful.items()}\n",
        "harmless = {k: torch.cat(v) for k, v in harmless.items()}"
      ],
      "metadata": {
        "id": "GhBNPb95tbmm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Helper function to get activation index\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",
        "# Compute difference of means between harmful and harmless activations at intermediate layers\n",
        "activation_layers = [\"resid_pre\", \"resid_mid\", \"resid_post\"]\n",
        "activation_refusals = defaultdict(list)\n",
        "\n",
        "for layer_num in range(1, model.cfg.n_layers):\n",
        "    pos = -1  # Position index\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(\n",
        "            dim=0\n",
        "        )\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",
        "# Get all calculated potential refusal directions, sort them in descending order based on their mean\n",
        "# Use a subset of layers if certain activations are not promising\n",
        "selected_layers = [\"resid_pre\"]\n",
        "\n",
        "activation_scored = sorted(\n",
        "    [\n",
        "        activation_refusals[layer][l - 1]\n",
        "        for l in range(1, model.cfg.n_layers)\n",
        "        for layer in selected_layers\n",
        "    ],\n",
        "    key=lambda x: abs(x.mean()),\n",
        "    reverse=True,\n",
        ")"
      ],
      "metadata": {
        "id": "PX4w0HeUwAFm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_vhhwl-2-jPg"
      },
      "outputs": [],
      "source": [
        "def _generate_with_hooks(\n",
        "    model: HookedTransformer,\n",
        "    tokenizer: AutoTokenizer,\n",
        "    tokens: Int[Tensor, \"batch_size seq_len\"],\n",
        "    max_tokens_generated: int = 64,\n",
        "    fwd_hooks=[],\n",
        ") -> List[str]:\n",
        "    all_tokens = torch.zeros(\n",
        "        (tokens.shape[0], tokens.shape[1] + max_tokens_generated),\n",
        "        dtype=torch.long,\n",
        "        device=tokens.device,\n",
        "    )\n",
        "    all_tokens[:, : tokens.shape[1]] = tokens\n",
        "    for i in range(max_tokens_generated):\n",
        "        with model.hooks(fwd_hooks=fwd_hooks):\n",
        "            logits = model(all_tokens[:, : -max_tokens_generated + i])\n",
        "            next_tokens = logits[:, -1, :].argmax(\n",
        "                dim=-1\n",
        "            )  # greedy sampling (temperature=0)\n",
        "            all_tokens[:, -max_tokens_generated + i] = next_tokens\n",
        "    return tokenizer.batch_decode(\n",
        "        all_tokens[:, tokens.shape[1] :], skip_special_tokens=True\n",
        "    )\n",
        "\n",
        "def get_generations(\n",
        "    model: HookedTransformer,\n",
        "    tokenizer: AutoTokenizer,\n",
        "    instructions: List[str],\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",
        "        tokens = tokenize_instructions(\n",
        "            tokenizer, instructions=instructions[i : i + batch_size]\n",
        "        )\n",
        "        generation = _generate_with_hooks(\n",
        "            model,\n",
        "            tokenizer,\n",
        "            tokens,\n",
        "            max_tokens_generated=max_tokens_generated,\n",
        "            fwd_hooks=fwd_hooks,\n",
        "        )\n",
        "        generations.extend(generation)\n",
        "    return generations\n",
        "\n",
        "# Inference-time intervention hook\n",
        "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 = (\n",
        "        einops.einsum(\n",
        "            activation, direction.view(-1, 1), \"... d_act, d_act single -> ... single\"\n",
        "        )\n",
        "        * direction\n",
        "    )\n",
        "    return activation - proj\n",
        "\n",
        "# Testing baseline\n",
        "N_INST_TEST = 4\n",
        "baseline_generations = get_generations(\n",
        "    model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[]\n",
        ")\n",
        "\n",
        "# Evaluating layers defined earlier (needs human evaluation to determine best layer for refusal inhibition)\n",
        "EVAL_N = 20  # Evaluate how many of the top N potential directions\n",
        "evals = []\n",
        "for refusal_dir in tqdm(activation_scored[:EVAL_N]):\n",
        "    hook_fn = functools.partial(direction_ablation_hook, direction=refusal_dir)\n",
        "    fwd_hooks = [\n",
        "        (utils.get_act_name(act_name, layer), hook_fn)\n",
        "        for layer in list(range(model.cfg.n_layers))\n",
        "        for act_name in activation_layers\n",
        "    ]\n",
        "    intervention_generations = get_generations(\n",
        "        model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=fwd_hooks\n",
        "    )\n",
        "    evals.append(intervention_generations)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Print generations for human evaluation\n",
        "blacklist = [\"I cannot\", \"I can't\"]\n",
        "for i in range(N_INST_TEST):\n",
        "    print(f\"\\033[1mINSTRUCTION {i}: {harmful_inst_test[i]}\")\n",
        "    print(f\"\\nBASELINE COMPLETION:\\n{baseline_generations[i]}\\033[0m\")\n",
        "    for layer_candidate in range(EVAL_N):\n",
        "        if not any(word in evals[layer_candidate][i] for word in blacklist):\n",
        "            print(f\"\\n---\\n\\nLAYER CANDIDATE #{layer_candidate} INTERVENTION COMPLETION:\")\n",
        "            print(evals[layer_candidate][i])"
      ],
      "metadata": {
        "id": "5SbBZ-Un14Yc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lY5B3-uP_EjE"
      },
      "outputs": [],
      "source": [
        "def get_orthogonalized_matrix(\n",
        "    matrix: Float[Tensor, \"... d_model\"], vec: Float[Tensor, \"d_model\"]\n",
        ") -> Float[Tensor, \"... d_model\"]:\n",
        "    proj = (\n",
        "        einops.einsum(\n",
        "            matrix, vec.view(-1, 1), \"... d_model, d_model single -> ... single\"\n",
        "        )\n",
        "        * vec\n",
        "    )\n",
        "    return matrix - proj\n",
        "\n",
        "# Select the layer with the highest potential refusal direction\n",
        "LAYER_CANDIDATE = 9\n",
        "refusal_dir = activation_scored[LAYER_CANDIDATE]\n",
        "\n",
        "# Orthogonalize the model's weights\n",
        "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",
        "# Generate text with abliterated model\n",
        "orthogonalized_generations = get_generations(\n",
        "    model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[]\n",
        ")\n",
        "\n",
        "# Print generations\n",
        "for i in range(N_INST_TEST):\n",
        "    if len(baseline_generations) > i:\n",
        "        print(f\"INSTRUCTION {i}: {harmful_inst_test[i]}\")\n",
        "        print(f\"\\033[92mBASELINE COMPLETION:\\n{baseline_generations[i]}\")\n",
        "    print(f\"\\033[91mINTERVENTION COMPLETION:\\n{evals[LAYER_CANDIDATE][i]}\")\n",
        "    print(f\"\\033[95mORTHOGONALIZED COMPLETION:\\n{orthogonalized_generations[i]}\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Convert model back to HF safetensors\n",
        "hf_model = AutoModelForCausalLM.from_pretrained(MODEL_TYPE, torch_dtype=torch.bfloat16)\n",
        "lm_model = hf_model.model\n",
        "\n",
        "state_dict = model.state_dict()\n",
        "lm_model.embed_tokens.weight = torch.nn.Parameter(state_dict[\"embed.W_E\"].cpu())\n",
        "\n",
        "for l in range(model.cfg.n_layers):\n",
        "    lm_model.layers[l].self_attn.o_proj.weight = torch.nn.Parameter(\n",
        "        einops.rearrange(\n",
        "            state_dict[f\"blocks.{l}.attn.W_O\"], \"n h m->m (n h)\", n=model.cfg.n_heads\n",
        "        ).contiguous()\n",
        "    )\n",
        "    lm_model.layers[l].mlp.down_proj.weight = torch.nn.Parameter(\n",
        "        torch.transpose(state_dict[f\"blocks.{l}.mlp.W_out\"], 0, 1).contiguous()\n",
        "    )\n",
        "\n",
        "# Push it to the Hugging Face Hub\n",
        "hf_model.push_to_hub(NEW_MODEL_ID)"
      ],
      "metadata": {
        "id": "NGbVGBi0tP-O"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "L4",
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.19"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}