{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's implement CCS from scratch.\n",
    "This will deliberately be a simple (but less efficient) implementation to make everything as clear as possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/accounts/projects/jsteinhardt/uid1693600/.local/lib/python3.7/site-packages/pandas/compat/_optional.py:138: UserWarning: Pandas requires version '2.7.0' or newer of 'numexpr' (version '2.6.9' currently installed).\n",
      "  warnings.warn(msg, UserWarning)\n",
      "W1206 14:15:47.056510 140146004121408 builder.py:641] Reusing dataset amazon_polarity (/scratch/users/uid1693600/huggingface-cache/amazon_polarity/amazon_polarity/3.0.0/a27b32b7e7b88eb274a8fa8ba0f654f1fe998a87c22547557317793b5d2772dc)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e5b9e737f6e415aa43cb2e7234b149f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Some weights of the model checkpoint at microsoft/deberta-v2-xxlarge were not used when initializing DebertaV2ForMaskedLM: ['lm_predictions.lm_head.bias', 'lm_predictions.lm_head.LayerNorm.bias', 'lm_predictions.lm_head.dense.bias', 'lm_predictions.lm_head.dense.weight', 'lm_predictions.lm_head.LayerNorm.weight', 'deberta.embeddings.position_embeddings.weight']\n",
      "- This IS expected if you are initializing DebertaV2ForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DebertaV2ForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of DebertaV2ForMaskedLM were not initialized from the model checkpoint at microsoft/deberta-v2-xxlarge and are newly initialized: ['cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import copy\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "# Let's just try IMDB for simplicity\n",
    "data = load_dataset(\"amazon_polarity\")[\"test\"]\n",
    "\n",
    "# Here are a few different model options you can play around with:\n",
    "model_name = \"deberta\"\n",
    "# model_name = \"gpt-j\"\n",
    "# model_name = \"t5\"\n",
    "\n",
    "# if you want to cache the model weights somewhere, you can specify that here\n",
    "cache_dir = None\n",
    "\n",
    "if model_name == \"deberta\":\n",
    "    model_type = \"encoder\"\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"microsoft/deberta-v2-xxlarge\", cache_dir=cache_dir)\n",
    "    model = AutoModelForMaskedLM.from_pretrained(\"microsoft/deberta-v2-xxlarge\", cache_dir=cache_dir)\n",
    "    model.cuda()\n",
    "elif model_name == \"gpt-j\":\n",
    "    model_type = \"decoder\"\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\", cache_dir=cache_dir)\n",
    "    model = AutoModelForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\", cache_dir=cache_dir)\n",
    "    model.cuda()\n",
    "elif model_name == \"t5\":\n",
    "    model_type = \"encoder_decoder\"\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"t5-11b\", cache_dir=cache_dir)\n",
    "    model = AutoModelForSeq2SeqLM.from_pretrained(\"t5-11b\", cache_dir=cache_dir)\n",
    "    model.parallelize()  # T5 is big enough that we may need to run it on multiple GPUs\n",
    "else:\n",
    "    print(\"Not implemented!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First let's write code for extracting hidden states given a model and text. \n",
    "How we do this exactly will depend on the type of model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_encoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
    "    \"\"\"\n",
    "    Given an encoder model and some text, gets the encoder hidden states (in a given layer, by default the last) \n",
    "    on that input text (where the full text is given to the encoder).\n",
    "\n",
    "    Returns a numpy array of shape (hidden_dim,)\n",
    "    \"\"\"\n",
    "    # tokenize\n",
    "    encoder_text_ids = tokenizer(input_text, truncation=True, return_tensors=\"pt\").input_ids.to(model.device)\n",
    "\n",
    "    # forward pass\n",
    "    with torch.no_grad():\n",
    "        output = model(encoder_text_ids, output_hidden_states=True)\n",
    "\n",
    "    # get the appropriate hidden states\n",
    "    hs_tuple = output[\"hidden_states\"]\n",
    "    \n",
    "    hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
    "\n",
    "    return hs\n",
    "\n",
    "def get_encoder_decoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
    "    \"\"\"\n",
    "    Given an encoder-decoder model and some text, gets the encoder hidden states (in a given layer, by default the last) \n",
    "    on that input text (where the full text is given to the encoder).\n",
    "\n",
    "    Returns a numpy array of shape (hidden_dim,)\n",
    "    \"\"\"\n",
    "    # tokenize\n",
    "    encoder_text_ids = tokenizer(input_text, return_tensors=\"pt\").input_ids.to(model.device)\n",
    "    decoder_text_ids = tokenizer(\"\", return_tensors=\"pt\").input_ids.to(model.device)\n",
    "\n",
    "    # forward pass\n",
    "    with torch.no_grad():\n",
    "        output = model(encoder_text_ids, decoder_input_ids=decoder_text_ids, output_hidden_states=True)\n",
    "\n",
    "    # get the appropriate hidden states\n",
    "    hs_tuple = output[\"encoder_hidden_states\"]\n",
    "    hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
    "\n",
    "    return hs\n",
    "\n",
    "def get_decoder_hidden_states(model, tokenizer, input_text, layer=-1):\n",
    "    \"\"\"\n",
    "    Given a decoder model and some text, gets the hidden states (in a given layer, by default the last) on that input text\n",
    "\n",
    "    Returns a numpy array of shape (hidden_dim,)\n",
    "    \"\"\"\n",
    "    # tokenize (adding the EOS token this time)\n",
    "    input_ids = tokenizer(input_text + tokenizer.eos_token, return_tensors=\"pt\").input_ids.to(model.device)\n",
    "\n",
    "    # forward pass\n",
    "    with torch.no_grad():\n",
    "        output = model(input_ids, output_hidden_states=True)\n",
    "\n",
    "    # get the last layer, last token hidden states\n",
    "    hs_tuple = output[\"hidden_states\"]\n",
    "    hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
    "\n",
    "    return hs\n",
    "\n",
    "def get_hidden_states(model, tokenizer, input_text, layer=-1, model_type=\"encoder\"):\n",
    "    fn = {\"encoder\": get_encoder_hidden_states, \"encoder_decoder\": get_encoder_decoder_hidden_states,\n",
    "          \"decoder\": get_decoder_hidden_states}[model_type]\n",
    "\n",
    "    return fn(model, tokenizer, input_text, layer=layer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now let's write code for formatting data and for getting all the hidden states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_imdb(text, label):\n",
    "    \"\"\"\n",
    "    Given an imdb example (\"text\") and corresponding label (0 for negative, or 1 for positive), \n",
    "    returns a zero-shot prompt for that example (which includes that label as the answer).\n",
    "    \n",
    "    (This is just one example of a simple, manually created prompt.)\n",
    "    \"\"\"\n",
    "    return \"The following movie review expresses a \" + [\"negative\", \"positive\"][label] + \" sentiment:\\n\" + text\n",
    "\n",
    "\n",
    "def get_hidden_states_many_examples(model, tokenizer, data, model_type, n=100):\n",
    "    \"\"\"\n",
    "    Given an encoder-decoder model, a list of data, computes the contrast hidden states on n random examples.\n",
    "    Returns numpy arrays of shape (n, hidden_dim) for each candidate label, along with a boolean numpy array of shape (n,)\n",
    "    with the ground truth labels\n",
    "    \n",
    "    This is deliberately simple so that it's easy to understand, rather than being optimized for efficiency\n",
    "    \"\"\"\n",
    "    # setup\n",
    "    model.eval()\n",
    "    all_neg_hs, all_pos_hs, all_gt_labels = [], [], []\n",
    "\n",
    "    # loop\n",
    "    for _ in tqdm(range(n)):\n",
    "        # for simplicity, sample a random example until we find one that's a reasonable length\n",
    "        # (most examples should be a reasonable length, so this is just to make sure)\n",
    "        while True:\n",
    "            idx = np.random.randint(len(data))\n",
    "            text, true_label = data[idx][\"content\"], data[idx][\"label\"]\n",
    "            # the actual formatted input will be longer, so include a bit of a marign\n",
    "            if len(tokenizer(text)) < 400:  \n",
    "                break\n",
    "                \n",
    "        # get hidden states\n",
    "        neg_hs = get_hidden_states(model, tokenizer, format_imdb(text, 0), model_type=model_type)\n",
    "        pos_hs = get_hidden_states(model, tokenizer, format_imdb(text, 1), model_type=model_type)\n",
    "\n",
    "        # collect\n",
    "        all_neg_hs.append(neg_hs)\n",
    "        all_pos_hs.append(pos_hs)\n",
    "        all_gt_labels.append(true_label)\n",
    "\n",
    "    all_neg_hs = np.stack(all_neg_hs)\n",
    "    all_pos_hs = np.stack(all_pos_hs)\n",
    "    all_gt_labels = np.stack(all_gt_labels)\n",
    "\n",
    "    return all_neg_hs, all_pos_hs, all_gt_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:19<00:00,  5.10it/s]\n"
     ]
    }
   ],
   "source": [
    "neg_hs, pos_hs, y = get_hidden_states_many_examples(model, tokenizer, data, model_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's verify that the model's representations are good\n",
    "\n",
    "Before trying CCS, let's make sure there exists a direction that classifies examples as true vs false with high accuracy; if logistic regression accuracy is bad, there's no hope of CCS doing well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logistic regression accuracy: 0.9\n"
     ]
    }
   ],
   "source": [
    "# let's create a simple 50/50 train split (the data is already randomized)\n",
    "n = len(y)\n",
    "neg_hs_train, neg_hs_test = neg_hs[:n//2], neg_hs[n//2:]\n",
    "pos_hs_train, pos_hs_test = pos_hs[:n//2], pos_hs[n//2:]\n",
    "y_train, y_test = y[:n//2], y[n//2:]\n",
    "\n",
    "# for simplicity we can just take the difference between positive and negative hidden states\n",
    "# (concatenating also works fine)\n",
    "x_train = neg_hs_train - pos_hs_train\n",
    "x_test = neg_hs_test - pos_hs_test\n",
    "\n",
    "lr = LogisticRegression(class_weight=\"balanced\")\n",
    "lr.fit(x_train, y_train)\n",
    "print(\"Logistic regression accuracy: {}\".format(lr.score(x_test, y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now let's try CCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLPProbe(nn.Module):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.linear1 = nn.Linear(d, 100)\n",
    "        self.linear2 = nn.Linear(100, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = F.relu(self.linear1(x))\n",
    "        o = self.linear2(h)\n",
    "        return torch.sigmoid(o)\n",
    "\n",
    "class CCS(object):\n",
    "    def __init__(self, x0, x1, nepochs=1000, ntries=10, lr=1e-3, batch_size=-1, \n",
    "                 verbose=False, device=\"cuda\", linear=True, weight_decay=0.01, var_normalize=False):\n",
    "        # data\n",
    "        self.var_normalize = var_normalize\n",
    "        self.x0 = self.normalize(x0)\n",
    "        self.x1 = self.normalize(x1)\n",
    "        self.d = self.x0.shape[-1]\n",
    "\n",
    "        # training\n",
    "        self.nepochs = nepochs\n",
    "        self.ntries = ntries\n",
    "        self.lr = lr\n",
    "        self.verbose = verbose\n",
    "        self.device = device\n",
    "        self.batch_size = batch_size\n",
    "        self.weight_decay = weight_decay\n",
    "        \n",
    "        # probe\n",
    "        self.linear = linear\n",
    "        self.probe = self.initialize_probe()\n",
    "        self.best_probe = copy.deepcopy(self.probe)\n",
    "\n",
    "        \n",
    "    def initialize_probe(self):\n",
    "        if self.linear:\n",
    "            self.probe = nn.Linear(self.d, 1)\n",
    "        else:\n",
    "            self.probe = MLPProbe(self.d)\n",
    "        self.probe.to(self.device)    \n",
    "\n",
    "\n",
    "    def normalize(self, x):\n",
    "        \"\"\"\n",
    "        Mean-normalizes the data x (of shape (n, d))\n",
    "        If self.var_normalize, also divides by the standard deviation\n",
    "        \"\"\"\n",
    "        normalized_x = x - x.mean(axis=0, keepdims=True)\n",
    "        if self.var_normalize:\n",
    "            normalized_x /= normalized_x.std(axis=0, keepdims=True)\n",
    "\n",
    "        return normalized_x\n",
    "\n",
    "        \n",
    "    def get_tensor_data(self):\n",
    "        \"\"\"\n",
    "        Returns x0, x1 as appropriate tensors (rather than np arrays)\n",
    "        \"\"\"\n",
    "        x0 = torch.tensor(self.x0, dtype=torch.float, requires_grad=False, device=self.device)\n",
    "        x1 = torch.tensor(self.x1, dtype=torch.float, requires_grad=False, device=self.device)\n",
    "        return x0, x1\n",
    "    \n",
    "\n",
    "    def get_loss(self, p0, p1):\n",
    "        \"\"\"\n",
    "        Returns the CCS loss for two probabilities each of shape (n,1) or (n,)\n",
    "        \"\"\"\n",
    "        informative_loss = (torch.min(p0, p1)**2).mean(0)\n",
    "        consistent_loss = ((p0 - (1-p1))**2).mean(0)\n",
    "        return informative_loss + consistent_loss\n",
    "\n",
    "\n",
    "    def get_acc(self, x0_test, x1_test, y_test):\n",
    "        \"\"\"\n",
    "        Computes accuracy for the current parameters on the given test inputs\n",
    "        \"\"\"\n",
    "        x0 = torch.tensor(self.normalize(x0_test), dtype=torch.float, requires_grad=False, device=self.device)\n",
    "        x1 = torch.tensor(self.normalize(x1_test), dtype=torch.float, requires_grad=False, device=self.device)\n",
    "        with torch.no_grad():\n",
    "            p0, p1 = self.best_probe(x0), self.best_probe(x1)\n",
    "        avg_confidence = 0.5*(p0 + (1-p1))\n",
    "        predictions = (avg_confidence.detach().cpu().numpy() < 0.5).astype(int)[:, 0]\n",
    "        acc = (predictions == y_test).mean()\n",
    "        acc = max(acc, 1 - acc)\n",
    "\n",
    "        return acc\n",
    "    \n",
    "        \n",
    "    def train(self):\n",
    "        \"\"\"\n",
    "        Does a single training run of nepochs epochs\n",
    "        \"\"\"\n",
    "        x0, x1 = self.get_tensor_data()\n",
    "        permutation = torch.randperm(len(x0))\n",
    "        x0, x1 = x0[permutation], x1[permutation]\n",
    "        \n",
    "        # set up optimizer\n",
    "        optimizer = torch.optim.AdamW(self.probe.parameters(), lr=self.lr, weight_decay=self.weight_decay)\n",
    "        \n",
    "        batch_size = len(x0) if self.batch_size == -1 else self.batch_size\n",
    "        nbatches = len(x0) // batch_size\n",
    "\n",
    "        # Start training (full batch)\n",
    "        for epoch in range(self.nepochs):\n",
    "            for j in range(nbatches):\n",
    "                x0_batch = x0[j*batch_size:(j+1)*batch_size]\n",
    "                x1_batch = x1[j*batch_size:(j+1)*batch_size]\n",
    "            \n",
    "            # probe\n",
    "            p0, p1 = self.probe(x0_batch), self.probe(x1_batch)\n",
    "\n",
    "            # get the corresponding loss\n",
    "            loss = self.get_loss(p0, p1)\n",
    "\n",
    "            # update the parameters\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        return loss.detach().cpu().item()\n",
    "    \n",
    "    def repeated_train(self):\n",
    "        best_loss = np.inf\n",
    "        for train_num in range(self.ntries):\n",
    "            self.initialize_probe()\n",
    "            loss = self.train()\n",
    "            if loss < best_loss:\n",
    "                self.best_probe = copy.deepcopy(self.probe)\n",
    "                best_loss = loss\n",
    "\n",
    "        return best_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CCS accuracy: 0.88\n"
     ]
    }
   ],
   "source": [
    "# Train CCS without any labels\n",
    "ccs = CCS(neg_hs_train, pos_hs_train)\n",
    "ccs.repeated_train()\n",
    "\n",
    "# Evaluate\n",
    "ccs_acc = ccs.get_acc(neg_hs_test, pos_hs_test, y_test)\n",
    "print(\"CCS accuracy: {}\".format(ccs_acc))"
   ]
  }
 ],
 "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.7.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "b80286374679f2ad472c61c83fc267d31329b5dea8e2dcaccb727123767724c5"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}