causality_test_on_education.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b38663a-74ac-4726-b62e-2645bd06bbc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8381e2a3-9b80-47f8-8d07-91d7feb9997f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from dataset import TextDataset \n",
    "from collections import OrderedDict\n",
    "\n",
    "from dataset import llama_v2_prompt\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from torch import nn\n",
    "\n",
    "import time\n",
    "\n",
    "tic, toc = (time.time, time.time)\n",
    "device = \"cuda\"\n",
    "torch_device = \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b1a14b7-cc71-44c5-8228-ac6df6c29fe7",
   "metadata": {},
   "outputs": [],
   "source": [
    "access_token = 'yours_here'\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-13b-chat-hf\", token=access_token, padding_side='left')\n",
    "model = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-2-13b-chat-hf\", token=access_token)\n",
    "model.half().cuda();\n",
    "model.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78ca7e56-a93b-48a7-a9a3-bd366b5ce65b",
   "metadata": {},
   "source": [
    "## Helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41df83a7-3963-41e6-81ef-1623de03834f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from probes import ProbeClassification, ProbeClassificationMixScaler, LinearProbeClassification, LinearProbeClassificationMixScaler\n",
    "from intervention_utils import load_probe_classifier, return_classifier_dict, num_classes\n",
    "from copy import deepcopy\n",
    "\n",
    "\n",
    "def optimize_one_inter_rep(inter_rep, layer_name, target, probe,\n",
    "                           lr=1e-2, max_epoch=128, \n",
    "                           loss_func=nn.CrossEntropyLoss(), \n",
    "                           verbose=False, simplified=False, N=4, normalized=False):\n",
    "    global first_time\n",
    "    tensor = (inter_rep.clone()).to(torch_device).requires_grad_(True)\n",
    "    rep_f = lambda: tensor\n",
    "    target_clone = target.clone().to(torch_device).to(torch.float)\n",
    "\n",
    "    optimizer = torch.optim.Adam([tensor], lr=lr)\n",
    "\n",
    "    cur_loss = 1000\n",
    "    begin_tensor = rep_f().clone().detach()\n",
    "    cur_input_tensor = rep_f().clone().detach()\n",
    "    if not simplified:\n",
    "        if verbose:\n",
    "            bar = tqdm(range(max_epoch), leave=False)\n",
    "        else:\n",
    "            bar = range(max_epoch)\n",
    "        for i in bar:\n",
    "            input_tensor = rep_f()\n",
    "            optimizer.zero_grad()\n",
    "            probe_seg_out = probe(input_tensor)[0]\n",
    "            # Compute the loss\n",
    "            if first_time:\n",
    "                if logistic:\n",
    "                    print(probe_seg_out)\n",
    "                else:\n",
    "                    print(torch.nn.functional.softmax(probe_seg_out))\n",
    "                first_time = False\n",
    "            loss = loss_func(probe_seg_out, target_clone)\n",
    "            # Call gradient descent\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            if verbose:\n",
    "                bar.set_description(f'At layer {layer_name} [{i + 1}/{max_epoch}]; Loss: {loss.item():.3f}') \n",
    "            if early_stop and abs(cur_input_tensor - input_tensor).mean() < stopping_eta:\n",
    "                break\n",
    "            # dist = torch.sqrt(torch.sum((begin_tensor - input_tensor) ** 2))\n",
    "            dist = torch.sum(torch.sqrt((begin_tensor - input_tensor) ** 2))\n",
    "            if dist > dist_thresh:\n",
    "                break\n",
    "            # cur_loss = loss.item()\n",
    "            cur_input_tensor = input_tensor.clone().detach()\n",
    "    else:\n",
    "        if normalized:\n",
    "            cur_input_tensor = rep_f() + target_clone.view(1, -1) @ probe.proj[0].weight * N * 100 / rep_f().norm() \n",
    "        else:\n",
    "            cur_input_tensor = rep_f() + target_clone.view(1, -1) @ probe.proj[0].weight * N\n",
    "        return cur_input_tensor.clone()\n",
    "    \n",
    "    return rep_f().clone()\n",
    "\n",
    "\n",
    "def edit_inter_rep_multi_layers(output, layer_name):\n",
    "    if residual:\n",
    "        layer_num = layer_name[layer_name.rfind(\"model.layers.\") + len(\"model.layers.\"):]\n",
    "    else:\n",
    "        layer_num = layer_name[layer_name.rfind(\"model.layers.\") + len(\"model.layers.\"):layer_name.rfind(\".mlp\")]\n",
    "    layer_num = int(layer_num)\n",
    "    probe = deepcopy(classifier_dict[attribute][layer_num])\n",
    "    if rand:\n",
    "        probe_weight = probe.proj[0].weight\n",
    "        if rand == 'uniform': \n",
    "            new_probe = torch.rand(probe_weight.shape)\n",
    "        elif rand == 'gaussian': \n",
    "            new_probe = torch.randn(probe_weight.shape)\n",
    "        else: \n",
    "            raise Exception\n",
    "        new_probe = new_probe.to(probe_weight.device)\n",
    "        for i in range(probe_weight.shape[0]):  # Iterate over rows\n",
    "            row_norm_original = torch.norm(probe_weight[i], p=2)\n",
    "            row_norm_new = torch.norm(new_probe[i], p=2)\n",
    "            new_probe[i] = (new_probe[i] / row_norm_new) * row_norm_original\n",
    "        with torch.no_grad():\n",
    "            probe.proj[0].weight = nn.Parameter(new_probe)\n",
    "    cloned_inter_rep = output[0][:,-1].unsqueeze(0).detach().clone().to(torch.float)\n",
    "    with torch.enable_grad():\n",
    "        cloned_inter_rep = optimize_one_inter_rep(cloned_inter_rep, layer_name, \n",
    "                                                  cf_target, probe,\n",
    "                                                  lr=100, max_epoch=1e5, \n",
    "                                                  loss_func=nn.BCELoss(),\n",
    "                                                  simplified=simplified,\n",
    "                                                  N=N,\n",
    "                                                  normalized=normalized)\n",
    "    output[0][:,-1] = cloned_inter_rep.to(torch.float16)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6885a237-f1b0-43a3-b920-265e1afa91ad",
   "metadata": {},
   "source": [
    "### Loading Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0409ac0-6e12-4158-a35f-68013236c779",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import nn\n",
    "import os\n",
    "from src.probes import LinearProbeClassification\n",
    "\n",
    "from src.intervention_utils import load_probe_classifier, return_classifier_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5635d67-504f-467b-ab68-e8ec647f7015",
   "metadata": {},
   "outputs": [],
   "source": [
    "if '<pad>' not in tokenizer.get_vocab():\n",
    "    tokenizer.add_special_tokens({\"pad_token\":\"<pad>\"})\n",
    "\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "model.config.pad_token_id = tokenizer.pad_token_id\n",
    "assert model.config.pad_token_id == tokenizer.pad_token_id, \"The model's pad token ID does not match the tokenizer's pad token ID!\"\n",
    "print('Tokenizer pad token ID:', tokenizer.pad_token_id)\n",
    "print('Model pad token ID:', model.config.pad_token_id)\n",
    "print('Model config pad token ID:', model.config.pad_token_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6da6a6a-f36f-4707-b253-004f6e14a2af",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "classifier_type = LinearProbeClassification\n",
    "classifier_directory = \"probe_checkpoints/controlling_probe\"\n",
    "return_user_msg_last_act = True\n",
    "include_inst = True\n",
    "layer_num = None\n",
    "mix_scaler = False\n",
    "residual_stream = True\n",
    "logistic = True\n",
    "sklearn = False\n",
    "\n",
    "classifier_dict = return_classifier_dict(classifier_directory,\n",
    "                                         classifier_type, \n",
    "                                         chosen_layer=layer_num,\n",
    "                                         mix_scaler=mix_scaler,\n",
    "                                         logistic=logistic,\n",
    "                                         sklearn=sklearn,\n",
    "                                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43939e64-4e40-4d66-9786-a4df24c8129b",
   "metadata": {},
   "outputs": [],
   "source": [
    "which_layers = []\n",
    "from_idx = 20 # Hyperparameter\n",
    "to_idx = 30 # Hyperparameter\n",
    "residual = True # Set True\n",
    "for name, module in model.named_modules():\n",
    "    if residual and name!= \"\" and name[-1].isdigit():\n",
    "        layer_num = name[name.rfind(\"model.layers.\") + len(\"model.layers.\"):]\n",
    "        if from_idx <= int(layer_num) < to_idx:\n",
    "            which_layers.append(name)\n",
    "    elif (not residual) and name.endswith(\".mlp\"):\n",
    "        layer_num = name[name.rfind(\"model.layers.\") + len(\"model.layers.\"):name.rfind(\".mlp\")]\n",
    "        if from_idx <= int(layer_num) < to_idx:\n",
    "            which_layers.append(name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60dc0e9e-0baf-4780-b2c4-f8e833da10ee",
   "metadata": {},
   "source": [
    "## Classifier-based eval"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b38987-bc3c-4706-a31e-f905635e87f1",
   "metadata": {},
   "source": [
    "### Generating answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824d7e48-7522-4579-8236-5cd9c934d341",
   "metadata": {},
   "outputs": [],
   "source": [
    "from baukit import TraceDict\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225f8f3f",
   "metadata": {},
   "source": [
    "### Loading in the questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76f16640-fe63-42f8-bbcd-5a8f535fe817",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('causality_test_questions/education.txt') as f: \n",
    "    questions = f.read().splitlines()\n",
    "questions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8e55dcb",
   "metadata": {},
   "source": [
    "### Function for generating the responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eea457a-33ee-41a7-bf01-7a3d4614394a",
   "metadata": {},
   "outputs": [],
   "source": [
    "simplified=True\n",
    "normalized=False\n",
    "early_stop=False\n",
    "\n",
    "if '<pad>' not in tokenizer.get_vocab():\n",
    "    tokenizer.add_special_tokens({\"pad_token\":\"<pad>\"})\n",
    "\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "model.config.pad_token_id = tokenizer.pad_token_id\n",
    "assert model.config.pad_token_id == tokenizer.pad_token_id, \"The model's pad token ID does not match the tokenizer's pad token ID!\"\n",
    "print('Tokenizer pad token ID:', tokenizer.pad_token_id)\n",
    "print('Model pad token ID:', model.config.pad_token_id)\n",
    "print('Model config pad token ID:', model.config.pad_token_id)\n",
    "\n",
    "def collect_responses_batched(prompts, modified_layer_names, edit_function, batch_size=5, rand=None):\n",
    "    print(modified_layer_names)\n",
    "    responses = []\n",
    "    for i in tqdm(range(0, len(prompts), batch_size)): \n",
    "        \n",
    "        message_lists = [[{\"role\": \"user\", \n",
    "                         \"content\": prompt},\n",
    "                        ] for prompt in prompts[i:i+batch_size]]\n",
    "\n",
    "        # Transform the message list into a prompt string\n",
    "        formatted_prompts = [llama_v2_prompt(message_list) for message_list in message_lists]\n",
    "        \n",
    "        with TraceDict(model, modified_layer_names, edit_output=edit_function) as ret:\n",
    "            with torch.no_grad():\n",
    "                inputs = tokenizer(formatted_prompts, return_tensors='pt', padding=True).to('cuda')\n",
    "                tokens = model.generate(**inputs,\n",
    "                                        max_new_tokens=768,\n",
    "                                        do_sample=False,\n",
    "                                        temperature=0,\n",
    "                                        top_p=1,\n",
    "                                       )\n",
    "                \n",
    "        output = [tokenizer.decode(seq, skip_special_tokens=True).split('[/INST]')[1] for seq in tokens]\n",
    "        responses.extend(output)\n",
    "\n",
    "    return responses"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1ab736e",
   "metadata": {},
   "source": [
    "### Hyperparameters for Intervention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75fd534a-cd76-4b65-83bc-cfb90236b404",
   "metadata": {},
   "outputs": [],
   "source": [
    "category_labels = ['some schooling', 'high school', 'college and more']\n",
    "modified_layer_names = which_layers\n",
    "N = 8\n",
    "\n",
    "which_layers = []\n",
    "from_idx = 20 # Hyperparameter\n",
    "to_idx = 30 # Hyperparameter\n",
    "residual = True # Set True\n",
    "for name, module in model.named_modules():\n",
    "    if residual and name!= \"\" and name[-1].isdigit():\n",
    "        layer_num = name[name.rfind(\"model.layers.\") + len(\"model.layers.\"):]\n",
    "        if from_idx <= int(layer_num) < to_idx:\n",
    "            which_layers.append(name)\n",
    "    elif (not residual) and name.endswith(\".mlp\"):\n",
    "        layer_num = name[name.rfind(\"model.layers.\") + len(\"model.layers.\"):name.rfind(\".mlp\")]\n",
    "        if from_idx <= int(layer_num) < to_idx:\n",
    "            which_layers.append(name)\n",
    "        \n",
    "attribute = \"education\" # which attribute to intervene\n",
    "responses_dict = {}\n",
    "\n",
    "batch_size = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccbd1229-5772-4c53-847e-6d1c3262536a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# unintervened\n",
    "modified_layer_names = []\n",
    "def null(output, layer_name): \n",
    "    return output\n",
    "\n",
    "cf_target = [0] * len(category_labels)\n",
    "cf_target[0] = 1\n",
    "cf_target = torch.Tensor([cf_target])\n",
    "first_time=True\n",
    "rand = None\n",
    "responses_dict['unintervened'] = collect_responses_batched(questions, modified_layer_names, null, batch_size=batch_size, rand=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f3d9ab9-427f-47e6-94ab-35a2d3dee55b",
   "metadata": {},
   "outputs": [],
   "source": [
    "cf_target = [0] * len(category_labels)\n",
    "cf_target[0] = 1\n",
    "cf_target = torch.Tensor([cf_target])\n",
    "modified_layer_names = which_layers\n",
    "first_time=True\n",
    "rand = None\n",
    "responses_dict['some_schooling'] = collect_responses_batched(questions, modified_layer_names, edit_inter_rep_multi_layers, batch_size=batch_size, rand=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf603c3b-a07c-487b-ad4f-522b3f8eaba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cf_target = [0] * len(category_labels)\n",
    "cf_target[2] = 1\n",
    "cf_target = torch.Tensor([cf_target])\n",
    "print(cf_target)\n",
    "modified_layer_names = which_layers\n",
    "first_time=True\n",
    "rand = None\n",
    "responses_dict['college_more'] = collect_responses_batched(questions, modified_layer_names, edit_inter_rep_multi_layers, batch_size=batch_size, rand=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e342e63c-12c1-496b-aab6-5a7aa5bb5fe2",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c65fdaa2-3f14-4314-b2e1-2cf7719567ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3952a4f2-8e23-40ac-8ccc-56d28c87c920",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "openai.api_key = \"yours_here\"\n",
    "import json\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import os\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from dataset import TextDataset \n",
    "from collections import OrderedDict\n",
    "\n",
    "from dataset import llama_v2_prompt\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from torch import nn\n",
    "\n",
    "import time\n",
    "\n",
    "responses = responses_dict\n",
    "\n",
    "with open('causality_test_questions/education.txt') as f:\n",
    "    questions = f.read().splitlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff201fe8-be46-47bf-8caf-8c8d4d8f7231",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.random.seed(0)\n",
    "\n",
    "def formatter(prompt, response1, response2):\n",
    "    return f\"Prompt: {prompt}\\n\\Response 1: {response1}\\n\\Respnose 2: {response2}\\n\\nAnswer:\"\n",
    "\n",
    "results = []\n",
    "where_correct_answer = []\n",
    "response1s = responses['college_more']\n",
    "response2s = responses['some_schooling']\n",
    "template = \"Please answer which one of the following responses is more targeted to a {demographic} user. Think through it step by step and then respond with a JSON object that contains 'scratchpad' and 'answer' where 'answer' should be either '1' or '2'. Note that the answer field _must_ be '1' or '2'. You must give an answer. Do not output anything else.\\n\\n\"\n",
    "for i in tqdm(range(len(questions))):\n",
    "    d = 'highly educated' if np.random.randint(2) == 0 else 'not very educated'\n",
    "    instruction = template.format(demographic=d)\n",
    "    query = formatter(questions[i], response1s[i], response2s[i])\n",
    "    where_correct_answer += [1] if d == 'highly educated' else [2]\n",
    "\n",
    "    response = openai.ChatCompletion.create(\n",
    "        model=\"gpt-4-turbo-preview\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "            {\"role\": \"user\", \"content\": instruction + query}\n",
    "        ],\n",
    "        temperature=0.0,\n",
    "    )\n",
    "\n",
    "    answer = response.choices[0].message.content\n",
    "    results.append(answer.strip())\n",
    "\n",
    "where_correct_answer = np.array(where_correct_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "871f42ef-764c-4cc8-acba-b6c2881742fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(entry):\n",
    "    entry_cleaned = entry.strip().removeprefix(\"```json\\n\").removesuffix(\"\\n```\")\n",
    "    json_obj = json.loads(entry_cleaned)\n",
    "    return json_obj\n",
    "\n",
    "processed_results = np.array([int(process(entry)['answer']) for entry in results])\n",
    "print(f\"Success Rate (0 - 1):\", (processed_results == where_correct_answer).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e7766a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(questions)):\n",
    "    text = f\"USER: {questions[i]}\\n\\n\"\n",
    "    text += \"-\" * 50 + \"\\n\"\n",
    "    text += f\"Intervened: Increase internal model of some schooling\\n\"\n",
    "    text += f\"CHATBOT: {responses_dict['some_schooling'][i]}\"\n",
    "    text += \"\\n\\n\" + \"-\" * 50 + \"\\n\"\n",
    "    text += f\"Intervened: Increase internal model of college and more education\\n\"\n",
    "    text += f\"CHATBOT: {responses_dict['college_more'][i]}\"\n",
    "    print(text)\n",
    "    f = open(f\"intervention_results/education/education_question_{i}_intervened_response.txt\", \"w\")\n",
    "    f.write(text)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}