{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'4.30.1'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from matplotlib import pyplot as plt\n",
    "plt.style.use('ggplot')\n",
    "\n",
    "from typing import Optional, List, Dict, Union\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor\n",
    "from torch import optim\n",
    "from torch.utils.data import random_split, DataLoader, TensorDataset\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import transformers\n",
    "\n",
    "import lightning.pytorch as pl\n",
    "# from dataclasses import dataclass\n",
    "\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import f1_score, roc_auc_score, accuracy_score\n",
    "from sklearn.preprocessing import RobustScaler\n",
    "\n",
    "from tqdm.auto import tqdm\n",
    "import os\n",
    "\n",
    "from loguru import logger\n",
    "logger.add(os.sys.stderr, format=\"{time} {level} {message}\", level=\"INFO\")\n",
    "\n",
    "transformers.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM, AutoModelForCausalLM, AutoConfig\n",
    "from transformers import LogitsProcessorList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPTBigCodeConfig {\n",
      "  \"_name_or_path\": \"HuggingFaceH4/starchat-beta\",\n",
      "  \"activation_function\": \"gelu\",\n",
      "  \"architectures\": [\n",
      "    \"GPTBigCodeForCausalLM\"\n",
      "  ],\n",
      "  \"attention_softmax_in_fp32\": true,\n",
      "  \"attn_pdrop\": 0.1,\n",
      "  \"bos_token_id\": 0,\n",
      "  \"embd_pdrop\": 0.1,\n",
      "  \"eos_token_id\": 0,\n",
      "  \"inference_runner\": 0,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"layer_norm_epsilon\": 1e-05,\n",
      "  \"max_batch_size\": null,\n",
      "  \"max_sequence_length\": null,\n",
      "  \"model_type\": \"gpt_bigcode\",\n",
      "  \"multi_query\": true,\n",
      "  \"n_embd\": 6144,\n",
      "  \"n_head\": 48,\n",
      "  \"n_inner\": 24576,\n",
      "  \"n_layer\": 40,\n",
      "  \"n_positions\": 8192,\n",
      "  \"pad_key_length\": true,\n",
      "  \"pre_allocate_kv_cache\": false,\n",
      "  \"resid_pdrop\": 0.1,\n",
      "  \"scale_attention_softmax_in_fp32\": true,\n",
      "  \"scale_attn_weights\": true,\n",
      "  \"summary_activation\": null,\n",
      "  \"summary_first_dropout\": 0.1,\n",
      "  \"summary_proj_to_labels\": true,\n",
      "  \"summary_type\": \"cls_index\",\n",
      "  \"summary_use_proj\": true,\n",
      "  \"torch_dtype\": \"bfloat16\",\n",
      "  \"transformers_version\": \"4.30.1\",\n",
      "  \"use_cache\": true,\n",
      "  \"validate_runner_input\": true,\n",
      "  \"vocab_size\": 49156\n",
      "}\n",
      "\n",
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please run\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      " and submit this information together with your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "bin /home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so\n",
      "CUDA SETUP: CUDA runtime path found: /home/ubuntu/mambaforge/envs/dlk2/lib/libcudart.so\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 117\n",
      "CUDA SETUP: Loading binary /home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/mambaforge/envs/dlk2/lib/python3.9/site-packages/bitsandbytes/cuda_setup/main.py:149: UserWarning: Found duplicate ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] files: {PosixPath('/home/ubuntu/mambaforge/envs/dlk2/lib/libcudart.so'), PosixPath('/home/ubuntu/mambaforge/envs/dlk2/lib/libcudart.so.11.0')}.. We'll flip a coin and try one of these, in order to fail forward.\n",
      "Either way, this might cause trouble in the future:\n",
      "If you get `CUDA error: invalid device function` errors, the above might be the cause and the solution is to make sure only one ['libcudart.so', 'libcudart.so.11.0', 'libcudart.so.12.0'] in the paths that we search based on your env.\n",
      "  warn(msg)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14606a01ebf248cf90ae37d921833178",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# leaderboard https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard\n",
    "model_options = dict(\n",
    "    device_map=\"auto\",\n",
    "    load_in_4bit=True,\n",
    "    # load_in_8bit=True,\n",
    "    torch_dtype=torch.float16,\n",
    "    trust_remote_code=True,\n",
    "    use_safetensors=False,\n",
    "    # use_cache=False,\n",
    ")\n",
    "\n",
    "model_repo = \"HuggingFaceH4/starchat-beta\"\n",
    "\n",
    "config = AutoConfig.from_pretrained(model_repo, trust_remote_code=True,)\n",
    "print(config)\n",
    "config.use_cache = False\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_repo)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, **model_options)\n",
    "tokenizer.pad_token_id = 204"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "40\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "((2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38), 40)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Params\n",
    "BATCH_SIZE = 10 # None # None means auto # 6 gives 16Gb/25GB. where 10GB is the base model. so 6 is 6/15\n",
    "USE_MCDROPOUT = True\n",
    "\n",
    "try:\n",
    "    # num_layers = len(model.model.layers)\n",
    "    num_layers = model.config.n_layer\n",
    "    print(num_layers)\n",
    "except AttributeError:\n",
    "    try:\n",
    "        num_layers = len(model.base_model.model.model.layers)\n",
    "        print(num_layers)\n",
    "    except:\n",
    "        num_layers = 10\n",
    "        \n",
    "stride = 2\n",
    "# don't take the first or last layers as they can make it to easy to leak info\n",
    "extract_layers = tuple(range(2, num_layers-2, stride)) + (num_layers-2,)\n",
    "extract_layers, num_layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(17152, 17991)"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from src.datasets.hs import get_choices_as_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enable_dropout(model, USE_MCDROPOUT:Union[float,bool]=True):\n",
    "    \"\"\" Function to enable the dropout layers during test-time \"\"\"\n",
    "    \n",
    "    for m in model.modules():\n",
    "        if m.__class__.__name__.startswith('Dropout'):\n",
    "            m.train()\n",
    "            if USE_MCDROPOUT!=True:\n",
    "                m.p=USE_MCDROPOUT\n",
    "                # print(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load probe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from https://github.com/timeseriesAI/tsai/blob/f20027e236ff06ed8fa3f5d30da5ebdcc67fe5aa/tsai/models/layers.py#L261\n",
    "\n",
    "class AddCoords1d(nn.Module):\n",
    "    \"\"\"Add coordinates to ease position identification without modifying mean and std\"\"\"\n",
    "    def forward(self, x):\n",
    "        bs, _, seq_len = x.shape\n",
    "        cc = torch.linspace(-1,1,x.shape[-1], device=x.device).repeat(bs, 1, 1)\n",
    "        cc = (cc - cc.mean()) / cc.std()\n",
    "        x = torch.cat([x, cc], dim=1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLPProbe(nn.Module):\n",
    "    def __init__(self, c_in, depth=0, hs=16, dropout=0):\n",
    "        super().__init__()\n",
    "\n",
    "        layers = [\n",
    "            nn.BatchNorm1d(c_in, affine=False),  # this will normalise the inputs\n",
    "            AddCoords1d(),\n",
    "            nn.Dropout1d(dropout),\n",
    "            \n",
    "            nn.Conv1d(c_in+1, hs*(depth+1), kernel_size=2, padding=0),\n",
    "            nn.BatchNorm1d(hs*(depth+1)),\n",
    "            nn.ReLU(),\n",
    "        ]\n",
    "        for i in range(depth):\n",
    "            layers += [\n",
    "                AddCoords1d(),\n",
    "                nn.Conv1d(hs*(depth-i+1)+1, hs*(depth-i), 2, padding=0),\n",
    "                nn.BatchNorm1d(hs*(depth-i)),\n",
    "                nn.ReLU(),\n",
    "                \n",
    "            ]\n",
    "        layers += [nn.AdaptiveAvgPool1d(1)]\n",
    "        self.net = nn.Sequential(*layers)\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Linear(hs, hs), nn.BatchNorm1d(hs), nn.Dropout(dropout), nn.ReLU(), \n",
    "            nn.Linear(hs, 1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.net(x)\n",
    "        # print(1, h.shape)\n",
    "        h = h.squeeze(-1)\n",
    "        # print(1, h.shape)\n",
    "        return self.head(h)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_optimizer import Ranger21\n",
    "import torchmetrics\n",
    "\n",
    "from torchmetrics import Metric, MetricCollection, Accuracy, AUROC\n",
    "    \n",
    "class CSS(pl.LightningModule):\n",
    "    def __init__(self, c_in, total_steps, depth=1, hs=16, lr=4e-3, weight_decay=1e-9, dropout=0):\n",
    "        super().__init__()\n",
    "        self.probe = MLPProbe(c_in, depth=depth, dropout=dropout, hs=hs)\n",
    "        self.save_hyperparameters()\n",
    "        \n",
    "        self.loss_fn = nn.SmoothL1Loss()\n",
    "        \n",
    "        # metrics for each stage\n",
    "        metrics_template = MetricCollection({\n",
    "            'acc': Accuracy(task=\"binary\"), \n",
    "            'auroc': AUROC(task=\"binary\")\n",
    "        })\n",
    "        self.metrics = torch.nn.ModuleDict({\n",
    "            f'metrics_{stage}': metrics_template.clone(prefix=stage+'/') for stage in ['train', 'val', 'test']\n",
    "        })\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.probe(x).squeeze(1)\n",
    "        \n",
    "    def _step(self, batch, batch_idx, stage='train'):\n",
    "        x0, x1, y = batch\n",
    "        ypred0 = self(x0)\n",
    "        ypred1 = self(x1)\n",
    "        \n",
    "        if stage=='pred':\n",
    "            return (ypred1-ypred0).float()\n",
    "        \n",
    "        loss = self.loss_fn(ypred1-ypred0, y)\n",
    "        self.log(f\"{stage}/loss\", loss)\n",
    "        \n",
    "        m = self.metrics[f'metrics_{stage}']\n",
    "        \n",
    "        y_cls = switch2bool(ypred1-ypred0)\n",
    "        m(y_cls, y>0.)\n",
    "        self.log_dict(m, on_epoch=True, on_step=False)\n",
    "        return loss\n",
    "    \n",
    "    def training_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
    "        return self._step(batch, batch_idx)\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx=0):\n",
    "        return self._step(batch, batch_idx, stage='val')\n",
    "    \n",
    "    def predict_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
    "        return self._step(batch, batch_idx, stage='pred').cpu().detach()\n",
    "    \n",
    "    def test_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
    "        return self._step(batch, batch_idx, stage='test')\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        \"\"\"use ranger21 from  https://github.com/kozistr/pytorch_optimizer\"\"\"\n",
    "        optimizer = Ranger21(\n",
    "            self.parameters(),\n",
    "            lr=self.hparams.lr,\n",
    "            weight_decay=self.hparams.weight_decay,       \n",
    "            num_iterations=self.hparams.total_steps,\n",
    "        )\n",
    "        return optimizer\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = '/home/ubuntu/Documents/mjc/elk/discovering_latent_knowledge/notebooks/lightning_logs/version_338/checkpoints/epoch=37-step=2090.ckpt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CSS(\n",
       "  (probe): MLPProbe(\n",
       "    (net): Sequential(\n",
       "      (0): BatchNorm1d(6144, eps=1e-05, momentum=0.1, affine=False, track_running_stats=True)\n",
       "      (1): AddCoords1d()\n",
       "      (2): Dropout1d(p=0.1, inplace=False)\n",
       "      (3): Conv1d(6145, 1764, kernel_size=(2,), stride=(1,))\n",
       "      (4): BatchNorm1d(1764, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (5): ReLU()\n",
       "      (6): AddCoords1d()\n",
       "      (7): Conv1d(1765, 1512, kernel_size=(2,), stride=(1,))\n",
       "      (8): BatchNorm1d(1512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (9): ReLU()\n",
       "      (10): AddCoords1d()\n",
       "      (11): Conv1d(1513, 1260, kernel_size=(2,), stride=(1,))\n",
       "      (12): BatchNorm1d(1260, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (13): ReLU()\n",
       "      (14): AddCoords1d()\n",
       "      (15): Conv1d(1261, 1008, kernel_size=(2,), stride=(1,))\n",
       "      (16): BatchNorm1d(1008, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (17): ReLU()\n",
       "      (18): AddCoords1d()\n",
       "      (19): Conv1d(1009, 756, kernel_size=(2,), stride=(1,))\n",
       "      (20): BatchNorm1d(756, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (21): ReLU()\n",
       "      (22): AddCoords1d()\n",
       "      (23): Conv1d(757, 504, kernel_size=(2,), stride=(1,))\n",
       "      (24): BatchNorm1d(504, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (25): ReLU()\n",
       "      (26): AddCoords1d()\n",
       "      (27): Conv1d(505, 252, kernel_size=(2,), stride=(1,))\n",
       "      (28): BatchNorm1d(252, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (29): ReLU()\n",
       "      (30): AdaptiveAvgPool1d(output_size=1)\n",
       "    )\n",
       "    (head): Sequential(\n",
       "      (0): Linear(in_features=252, out_features=252, bias=True)\n",
       "      (1): BatchNorm1d(252, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (2): Dropout(p=0.1, inplace=False)\n",
       "      (3): ReLU()\n",
       "      (4): Linear(in_features=252, out_features=1, bias=True)\n",
       "    )\n",
       "  )\n",
       "  (loss_fn): SmoothL1Loss()\n",
       "  (metrics): ModuleDict(\n",
       "    (metrics_train): MetricCollection(\n",
       "      (acc): BinaryAccuracy()\n",
       "      (auroc): BinaryAUROC(),\n",
       "      prefix=train/\n",
       "    )\n",
       "    (metrics_val): MetricCollection(\n",
       "      (acc): BinaryAccuracy()\n",
       "      (auroc): BinaryAUROC(),\n",
       "      prefix=val/\n",
       "    )\n",
       "    (metrics_test): MetricCollection(\n",
       "      (acc): BinaryAccuracy()\n",
       "      (auroc): BinaryAUROC(),\n",
       "      prefix=test/\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = CSS.load_from_checkpoint(f)\n",
    "# , c_in=c_in, depth=6, hs=42*6, lr=3e-3, \n",
    "#         #   weight_decay=1e-4, \n",
    "#           dropout=0.1,\n",
    "#           )\n",
    "net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.helpers.torch import to_numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 230,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "       \n",
    "def get_hidden_states(model, tokenizer, input_text, layers=extract_layers, truncation_length=999, output_attentions=False, use_mcdropout=USE_MCDROPOUT, choice_n=\"No\", choice_p=\"Yes\"):\n",
    "    \"\"\"\n",
    "    Given a decoder model and some texts, gets the hidden states (in a given layer) on that input texts\n",
    "    \"\"\"\n",
    "    id_n, id_y = get_choices_as_tokens(choice_n, choice_p)\n",
    "    if not isinstance(input_text, list):\n",
    "        input_text = [input_text]\n",
    "    input_ids = tokenizer(input_text, \n",
    "                          return_tensors=\"pt\",\n",
    "                          padding=True,\n",
    "                            add_special_tokens=True,\n",
    "                         ).input_ids.to(model.device)\n",
    "        \n",
    "    # Handling truncation: truncate start, not end\n",
    "    if truncation_length is not None:\n",
    "        if input_ids.size(1)>truncation_length:\n",
    "            print('truncating', input_ids.size(1))\n",
    "        input_ids = input_ids[:, -truncation_length:]\n",
    "\n",
    "    # forward pass\n",
    "    last_token = -1\n",
    "    first_token = 0\n",
    "    with torch.no_grad():\n",
    "        model.eval()        \n",
    "        if use_mcdropout: enable_dropout(model, use_mcdropout)\n",
    "        \n",
    "        # taken from greedy_decode https://github.com/huggingface/transformers/blob/ba695c1efd55091e394eb59c90fb33ac3f9f0d41/src/transformers/generation/utils.py\n",
    "        logits_processor = LogitsProcessorList()\n",
    "        model_kwargs = dict(use_cache=False)\n",
    "        model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)\n",
    "        outputs = model.forward(**model_inputs, return_dict=True, output_attentions=output_attentions, output_hidden_states=True)\n",
    "        \n",
    "        next_token_logits = outputs.logits[:, last_token, :]\n",
    "        outputs['scores'] = logits_processor(input_ids, next_token_logits)[:, None,:]\n",
    "        \n",
    "        next_tokens = torch.argmax(outputs['scores'], dim=-1)\n",
    "        outputs['sequences'] = torch.cat([input_ids, next_tokens], dim=-1)\n",
    "\n",
    "        # the output is large, so we will just select what we want 1) the first token with[:, 0]\n",
    "        # 2) selected layers with [layers]\n",
    "        attentions = None\n",
    "        if output_attentions:\n",
    "            # shape is [(batch_size, num_heads, sequence_length, sequence_length)]*num_layers\n",
    "            # lets take max?\n",
    "            attentions = [outputs['attentions'][i] for i in layers]\n",
    "            attentions = [v[:, last_token] for v in attentions]\n",
    "            attentions = torch.concat(attentions)\n",
    "        \n",
    "        hidden_states = torch.stack([outputs['hidden_states'][i] for i in layers], 1)\n",
    "        \n",
    "        hidden_states = hidden_states[:, :, last_token] # (batch, layers, past_seq, logits) take just the last token so they are same size\n",
    "        \n",
    "        input_truncated = tokenizer.batch_decode(input_ids)\n",
    "        \n",
    "        s = outputs['sequences']\n",
    "        s = [s[i][len(input_ids[i]):] for i in range(len(s))]\n",
    "        text_ans = tokenizer.batch_decode(s)\n",
    "\n",
    "        scores = outputs['scores'][:, first_token].softmax(-1) # for first (and only) token\n",
    "        prob_n, prob_y = scores[:, [id_n, id_y]].T\n",
    "        eps = 1e-3\n",
    "        ans = (prob_y/(prob_n+prob_y+eps))\n",
    "    \n",
    "    out = dict(hidden_states=hidden_states, ans=ans, text_ans=text_ans, input_truncated=input_truncated, input_id_shape=input_ids.shape,\n",
    "                attentions=attentions, prob_n=prob_n, prob_y=prob_y, scores=outputs['scores'][:, 0], input_text=input_text,\n",
    "               )\n",
    "    out = {k:to_numpy(v) for k,v in out.items()}    \n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 262,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pred(text):\n",
    "    input_text = f\"\"\"\n",
    "    <|system|><|end|>\n",
    "    <|user|>{text}<|end|>\n",
    "    <|assistant|>\n",
    "    \"\"\"\n",
    "    hs0 = get_hidden_states(model, tokenizer, input_text)\n",
    "    hs1 = get_hidden_states(model, tokenizer, input_text)\n",
    "\n",
    "    hs2x = lambda x: torch.from_numpy(x['hidden_states'].transpose(0, 2, 1)).to(device)\n",
    "    x0 = hs2x(hs0)\n",
    "    x1 = hs2x(hs1)\n",
    "    assert not (x1==x0).all()\n",
    "    x0.shape\n",
    "\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        ypred0 = to_numpy(net.probe(x0))\n",
    "        ypred1 = to_numpy(net.probe(x1))\n",
    "        y_pred = (ypred1-ypred0)\n",
    "        \n",
    "    y_pred\n",
    "    hs2_more_positive = (hs1['ans']-hs0['ans'])>0\n",
    "    y_pred * np.sign(hs1['ans']-hs0['ans'])\n",
    "    print(f\"\"\"========================\n",
    "{text}\n",
    "----------------------------\n",
    "Model Ans0: \"{hs0['text_ans'][0]}\" prob_y={hs0['prob_y']:2.2f} prob_n={hs0['prob_n']} probe={ypred0}\n",
    "Model Ans1: \"{hs1['text_ans'][0]}\" prob_y={hs1['prob_y']:2.2f} prob_n={hs1['prob_n']} probe={ypred0}\n",
    "Probe says: {y_pred:2.4f} (+ve means ans0 is more true)\n",
    "where\n",
    "    hs1_more_positive={hs2_more_positive}\n",
    "    hs1_more_true={y_pred>0}\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 263,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "========================\n",
      "A broken mirror will give 7 years bad luck?\n",
      "----------------------------\n",
      "Model Ans0: \" Yes\" prob_y=0.00 prob_n=2.384185791015625e-07 probe=-0.22275735437870026\n",
      "Model Ans1: \" No\" prob_y=0.00 prob_n=6.556510925292969e-07 probe=-0.22275735437870026\n",
      "Probe says: -0.1197 (+ve means ans0 is more true)\n",
      "where\n",
      "    hs1_more_positive=True\n",
      "    hs1_more_true=False\n",
      "\n"
     ]
    }
   ],
   "source": [
    "texts = [\n",
    "    \"A broken mirror will give 7 years bad luck?\"\n",
    "    ]\n",
    "for text in texts:\n",
    "   pred(text)   \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = next(net.parameters()).device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 222,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model says: Positive Positive prob_y=0.44 prob_n=0.01\n",
      "Probe says: -0.0602 (hs2 is more true)\n",
      "where\n",
      "    hs2_more_positive=True\n",
      "\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 219,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model says: Positive Positive prob_y=0.56 prob_n=0.03\n",
      "Probe says: -0.1134 (hs2 is more true)\n",
      "where\n",
      "    hs2_more_positive=True\n",
      "\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What is the probe predicting? Whether hs1 is more true than hs0"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dlk2",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}