{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4fc930ca-71a4-430e-b335-8ca8486411f2",
   "metadata": {},
   "source": [
    "## Boundless DAS\n",
    "\n",
    "This tutorial is adapated from __pyvene__, you can find their code [here](https://github.com/stanfordnlp/pyvene/tree/main).\n",
    "\n",
    "Read more about Boundless DAS from the original paper by Zhengxuan Wu et al. [here](https://arxiv.org/pdf/2305.08809)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f378d0ba",
   "metadata": {},
   "source": [
    "## Setup (Ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e7c4402d-bf1b-4256-a39d-fa03a61fcd99",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm, trange\n",
    "from nnsight import LanguageModel\n",
    "from pyvene import BoundlessRotatedSpaceIntervention\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import DataLoader\n",
    "from datasets import Dataset as hf_Dataset\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "import gc\n",
    "\n",
    "from tutorial_price_tagging_utils import factual_sampler, bound_alignment_sampler, lower_bound_alignment_example_sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0df9fe91-1144-40a8-bb6f-2681b48d99be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6539653f58f4aa3a058b3b8e7a69062",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/34 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Load model\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = LanguageModel('sharpbai/alpaca-7b-merged', device_map=\"cuda:0\", torch_dtype=torch.bfloat16, dispatch=True)\n",
    "remote = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "97f50d7a-bfc6-4b0e-bc1d-12cffabde938",
   "metadata": {},
   "outputs": [],
   "source": [
    "def free_unused_cuda_memory():\n",
    "    \"\"\"Free unused cuda memory.\"\"\"\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.empty_cache()\n",
    "    else:\n",
    "        raise RuntimeError(\"not using cuda\")\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b3aa9070-e632-4beb-bbd6-23cfce1028d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_loss(logits, labels, subspace_proj, mask_weight=1.5, vocab_size=32001):\n",
    "    shift_logits = logits[..., :, :].contiguous()\n",
    "    shift_labels = labels[..., :].contiguous()\n",
    "    # Flatten the tokens\n",
    "    loss_fct = torch.nn.CrossEntropyLoss()\n",
    "    shift_logits = shift_logits.view(-1, vocab_size)\n",
    "    shift_labels = shift_labels.view(-1)\n",
    "    # Enable model parallelism\n",
    "    shift_labels = shift_labels.to(shift_logits.device)\n",
    "    loss = loss_fct(shift_logits, shift_labels)\n",
    "    \n",
    "    \n",
    "    boundary_loss = mask_weight * subspace_proj.intervention_boundaries.sum()\n",
    "    loss += boundary_loss\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "08c1ae37-d7aa-452b-a669-2d0ea9e0f30c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(eval_preds, eval_labels, generate_output=False):\n",
    "    total_count = 0\n",
    "    correct_count = 0\n",
    "    \n",
    "    if generate_output:\n",
    "        outputs = []\n",
    "        gts = []\n",
    "        \n",
    "    for eval_pred, eval_label in zip(eval_preds, eval_labels):\n",
    "        \n",
    "        for i in range(eval_label.shape[0]):\n",
    "            label_idxs = eval_label[i].ne(-100).nonzero().squeeze(-1)\n",
    "                        \n",
    "            actual_test_labels = eval_label[i][label_idxs].tolist()\n",
    "            pred_test_labels = [eval_pred[i][idx].argmax(dim=-1).item() for idx in label_idxs]\n",
    "            \n",
    "            correct = actual_test_labels==pred_test_labels # uncomment it to evaluate all tokens\n",
    "            \n",
    "            if generate_output:\n",
    "                outputs.append(pred_test_labels)\n",
    "                gts.append(actual_test_labels)\n",
    "                        \n",
    "            total_count += 1\n",
    "            if correct:\n",
    "                correct_count += 1\n",
    "                \n",
    "    return_dict = {\"accuracy\": round(correct_count/total_count, 2)} \n",
    "    if generate_output:\n",
    "        return_dict[\"outputs\"] = outputs\n",
    "        return_dict[\"labels\"] = gts\n",
    "        \n",
    "    return return_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38ee8e2e",
   "metadata": {},
   "source": [
    "## Price Tagging game\n",
    "\n",
    "The instruction prompt of the Price Tagging game follows the publicly released template of the Alpaca (7B) model. The core instruction contains an English sentence: \"Please say yes only if it costs between [X.XX] and [X.XX] dollars, otherwise no.\" followed by an input dollar amount [X.XX], where [X.XX] are random continuous real numbers drawn with a uniform distribution from [0.00, 9.99]. The output is a single token ‘Yes’ or ‘No’.\n",
    "\n",
    "One hypothesis for how the model solves this task is the left boundary causal model which has one high-level boolean variable representing whether the input amount is higher than the lower bound, and an output node incorporating whether the input amount is also lower than the high bound. In this tutorial we focus on finding alignment for this causal model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "783658dd-87d4-4f34-8f40-5a03aa26870a",
   "metadata": {},
   "source": [
    "## Prealign Task"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e2218d2",
   "metadata": {},
   "source": [
    "To create our datasets, we are using [code](https://github.com/frankaging/pyvene/blob/cf93a1a6491dba65e1422fe20428f5972d17137e/counterfactual_datasets/price_tagging_game.py) copied from pyvene. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0653ae4f-d466-4818-82aa-d6e163f951ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_prealign = factual_sampler(model.tokenizer, 5000, game=\"pricing_tag\")\n",
    "\n",
    "prealign_dataset = hf_Dataset.from_dict(\n",
    "    {\"input_ids\": raw_prealign[0], \"labels\": raw_prealign[1]})\n",
    "prealign_dataset.set_format('torch', columns=['input_ids','labels'])\n",
    "prealign_dataloader = DataLoader(\n",
    "    prealign_dataset, batch_size=8\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0632b20f",
   "metadata": {},
   "source": [
    "Each instance in the dataset appear in this format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d88b074d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<s> Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\\n\\n### Instruction:\\nPlease say yes only if it costs between 2.52 and 7.83 dollars, otherwise no.\\n\\n### Input:\\n9.76 dollars\\n\\n### Response:\\n'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.tokenizer.decode(prealign_dataset['input_ids'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "387c7464-ad36-43a2-bb28-db8689c5ec1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 625/625 [01:42<00:00,  6.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[WARNING: THIS NEEDS TO BE GOOD!] prealign task accuracy: 0.92\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "        \n",
    "    eval_labels = []\n",
    "    eval_preds = []\n",
    "    \n",
    "    for step, inputs in enumerate(tqdm(prealign_dataloader)):\n",
    "        for k, v in inputs.items():\n",
    "            if v is not None and isinstance(v, torch.Tensor):\n",
    "                inputs[k] = v.to(device)\n",
    "                \n",
    "        outputs = model.forward(\n",
    "            input_ids=inputs['input_ids'],\n",
    "            labels=inputs['labels']\n",
    "        )\n",
    "        eval_labels += [inputs['labels'].detach().cpu()]\n",
    "            \n",
    "        eval_preds += [outputs.logits.detach().cpu()]\n",
    "    \n",
    "    eval_metrics = compute_metrics(eval_preds, eval_labels)\n",
    "\n",
    "eval_dict = eval_metrics\n",
    "print(f\"[WARNING: THIS NEEDS TO BE GOOD!] prealign task accuracy: {eval_dict['accuracy']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67f2775c-6392-4b51-8f79-eee242433175",
   "metadata": {},
   "source": [
    "## Boundless DAS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d0d3e58",
   "metadata": {},
   "source": [
    "The goal of Boundless DAS is to learn an alignment between potential distributed neural representations and high level causal variables.\n",
    "\n",
    "To train Boundless DAS, we sample two training examples and then swap the intermediate boolean values between them to produce a counterfactual output using our causal model. In parallel, we swap the aligned dimensions of the neural representations in rotated space. Lastly, we update our rotation matrix such that our neural network has a more similar counterfactual behavior to the causal model.\n",
    "\n",
    "We start by creating the training dataset for our trainable intervention."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "27e036fd-2221-4129-bc69-304a5ea567ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data = bound_alignment_sampler(\n",
    "    model.tokenizer, 10000, [lower_bound_alignment_example_sampler]\n",
    ")\n",
    "\n",
    "raw_train, raw_temp = train_test_split(\n",
    "    list(zip(*raw_data)), test_size=0.2, random_state=42\n",
    ")\n",
    "\n",
    "raw_eval, raw_test = train_test_split(\n",
    "    raw_temp, test_size=0.5, random_state=42\n",
    ")\n",
    "\n",
    "def unpack(data):\n",
    "    return tuple(map(list, zip(*data)))\n",
    "\n",
    "raw_train = unpack(raw_train)\n",
    "raw_eval = unpack(raw_eval)\n",
    "raw_test = unpack(raw_test)\n",
    "\n",
    "def create_dataset(data):\n",
    "    dataset = hf_Dataset.from_dict({\n",
    "        \"input_ids\": data[0],\n",
    "        \"source_input_ids\": data[1],\n",
    "        \"labels\": data[2],\n",
    "        \"intervention_ids\": data[3]  # we will not use this field\n",
    "    }).with_format(\"torch\")\n",
    "    return DataLoader(dataset, batch_size=8)\n",
    "\n",
    "train_dataloader = create_dataset(raw_train)\n",
    "eval_dataloader = create_dataset(raw_eval)\n",
    "test_dataloader = create_dataset(raw_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "27efb02f-9677-4a5a-8697-afed16fdfab3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BoundlessRotatedSpaceIntervention(\n",
       "  (rotate_layer): ParametrizedRotateLayer(\n",
       "    (parametrizations): ModuleDict(\n",
       "      (weight): ParametrizationList(\n",
       "        (0): _Orthogonal()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "subspace_proj = BoundlessRotatedSpaceIntervention(embed_dim=model.config.hidden_size).to('cuda')\n",
    "\n",
    "gradient_accumulation_steps = 4\n",
    "epochs = 3\n",
    "temperature_start = 50.0\n",
    "temperature_end = 0.1\n",
    "intervention_layer = 12\n",
    "\n",
    "t_total = int(len(train_dataloader) * epochs)\n",
    "warm_up_steps = 0.1 * t_total\n",
    "\n",
    "# Define params to be learned\n",
    "optimizer_params = []\n",
    "optimizer_params += [{'params': subspace_proj.rotate_layer.parameters()}]\n",
    "optimizer_params += [{'params': subspace_proj.intervention_boundaries, 'lr': 1e-2}]\n",
    "\n",
    "optimizer = torch.optim.Adam(\n",
    "    optimizer_params,\n",
    "    lr=1e-3,\n",
    ")\n",
    "\n",
    "scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer, num_warmup_steps=warm_up_steps,\n",
    "    num_training_steps=t_total\n",
    ")\n",
    "\n",
    "target_total_step = len(train_dataloader) * epochs\n",
    "\n",
    "temperature_schedule = torch.linspace(\n",
    "    temperature_start, temperature_end, target_total_step\n",
    ").to(torch.bfloat16).to(device)\n",
    "\n",
    "total_step = 0\n",
    "subspace_proj.set_temperature(temperature_schedule[total_step])\n",
    "subspace_proj.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7299f151",
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_subspace_swap(inputs, intervention_layer, model:LanguageModel, subspace_proj): #, batch_size=16\n",
    "    \"\"\"\n",
    "    Batched subspace_swap intervention at a single layer using nnsight\n",
    "    \"\"\"\n",
    "    batch_size = len(inputs['input_ids'])\n",
    "    all_inds = torch.arange(batch_size)\n",
    "        \n",
    "    base_prompt, source_prompt = inputs['input_ids'][:batch_size], inputs['source_input_ids'][:batch_size]\n",
    "\n",
    "    with model.trace(validate=False, remote=remote) as tracer:\n",
    "        with tracer.invoke(base_prompt, scan=False):\n",
    "            base = model.model.layers[intervention_layer].output[0].save()\n",
    "        \n",
    "        with tracer.invoke(source_prompt, scan=False):\n",
    "            source = model.model.layers[intervention_layer].output[0].save()\n",
    "    \n",
    "    with model.trace(validate=False, remote=remote) as tracer:\n",
    "        # intervention\n",
    "        with tracer.invoke(base_prompt, scan=False):\n",
    "            B = base[all_inds,80,:]\n",
    "            S = source[all_inds,80,:]\n",
    "    \n",
    "            mixed_out = subspace_proj(B, S, batch_size)\n",
    "            model.model.layers[intervention_layer].output[0][all_inds,80,:] = mixed_out\n",
    "        save_out = model.output.save()\n",
    "    del base, source, B,S\n",
    "    free_unused_cuda_memory()\n",
    "    \n",
    "    output_logits = save_out.value.logits\n",
    "    del save_out\n",
    "    return output_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "333559b5-a225-409e-a07a-d2fbcc5f41a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 0:   0%|          | 0/1000 [00:00<?, ?it/s]You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "Epoch: 0: 100%|██████████| 1000/1000 [43:49<00:00,  2.63s/it, loss=0.49, acc=0.88]\n",
      "Epoch: 1: 100%|██████████| 1000/1000 [49:47<00:00,  2.99s/it, loss=0.4, acc=0.88]\n",
      "Epoch: 2: 100%|██████████| 1000/1000 [51:38<00:00,  3.10s/it, loss=0.38, acc=0.88]\n",
      "Epoch: 100%|██████████| 3/3 [2:25:15<00:00, 2905.15s/it]\n"
     ]
    }
   ],
   "source": [
    "train_iterator = trange(\n",
    "    0, int(epochs), desc=\"Epoch\"\n",
    ")\n",
    "\n",
    "for epoch in train_iterator:\n",
    "    log_dicts = []\n",
    "    \n",
    "    epoch_iterator = tqdm(\n",
    "        train_dataloader, desc=f\"Epoch: {epoch}\", position=0, leave=True\n",
    "    )\n",
    "    \n",
    "    for step, inputs in enumerate(epoch_iterator):\n",
    "        for k, v in inputs.items():\n",
    "            if v is not None and isinstance(v, torch.Tensor):\n",
    "                inputs[k] = v.to(device)\n",
    "\n",
    "        counterfactual_outputs = batch_subspace_swap(inputs, intervention_layer, model, subspace_proj)\n",
    "        \n",
    "        eval_metrics = compute_metrics(\n",
    "            [counterfactual_outputs], [inputs['labels']]\n",
    "        )\n",
    "        \n",
    "        loss = calculate_loss(counterfactual_outputs, inputs[\"labels\"], subspace_proj)\n",
    "        loss_str = round(loss.item(), 2)\n",
    "        \n",
    "        log_dict = {'loss': loss_str, 'acc': eval_metrics[\"accuracy\"]}\n",
    "        epoch_iterator.set_postfix(log_dict)\n",
    "        \n",
    "        log_dicts.append(log_dict)\n",
    "        \n",
    "        if gradient_accumulation_steps > 1:\n",
    "            loss = loss / gradient_accumulation_steps\n",
    "            \n",
    "        loss.backward()\n",
    "        if total_step % gradient_accumulation_steps == 0:\n",
    "            if not (gradient_accumulation_steps > 1 and total_step == 0):\n",
    "                optimizer.step()\n",
    "                scheduler.step()\n",
    "                optimizer.zero_grad()\n",
    "                subspace_proj.set_temperature(temperature_schedule[total_step])\n",
    "                \n",
    "        total_step += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4ef8226",
   "metadata": {},
   "source": [
    "Evaluation on test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c1c0a00f-2d56-4364-b9ef-e08504128c30",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 125/125 [01:54<00:00,  1.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Boundless DAS accuracy: 0.93\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "        \n",
    "    eval_labels = []\n",
    "    eval_preds = []\n",
    "    \n",
    "    for step, inputs in enumerate(tqdm(test_dataloader)):\n",
    "        for k, v in inputs.items():\n",
    "            if v is not None and isinstance(v, torch.Tensor):\n",
    "                inputs[k] = v.to(device)\n",
    "        \n",
    "        outputs = batch_subspace_swap(inputs, intervention_layer, model, subspace_proj)#, batch_size=dataloader.batch_size)\n",
    "\n",
    "        eval_labels += [inputs['labels'].detach().cpu()]\n",
    "        eval_preds += [outputs.detach().cpu()]\n",
    "    \n",
    "    eval_metrics = compute_metrics(eval_preds, eval_labels)\n",
    "\n",
    "print(f\"Boundless DAS accuracy: {eval_metrics['accuracy']}\")"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
