guided-generation.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hardware\n",
    "\n",
    "These models run at half precision (fp16), so you'll need ~2x the parameter count in GPU ram bytes to run without further quantization. I.e. 7b -> 13 gb. A single A10 or A6000 is enough to get started. The largest model (65b) works on 4x A6000 (48gb each) or 2x A100 (80gb each).\n",
    "\n",
    "You can monitor graphics cards in a terminal with `nvidia-smi -l 1` (refresh every second). If the first GPU ram fills up and then you get a CUDA out of memory error, you may need to manually specify the max memory per card. \n",
    "\n",
    "### Models\n",
    "\n",
    "These models come from the huggingface project `decapoda-research`. Officially though, you're supposed to request access to the weights from Facebook. \n",
    "- [X] Submit this request [here](https://docs.google.com/forms/d/e/1FAIpQLSfqNECQnMkycAp2jP4Z9TFX0cGR4uf7b_fBxjY_OjhJILlKGA/viewform)\n",
    "\n",
    "### Dependencies\n",
    "\n",
    "`pip install numpy pandas torch datasets transformers matplotlib scikit-learn sentencepiece tqdm accelerate matplotlib`\n",
    "\n",
    "There's a weird mismatch in the casing of llama (`llama` vs `LLaMA`) that breaks loading these models. The easiest fix I found is just pull some guy's forked version of the `transformers` package that fixes this.\n",
    "\n",
    "`pip install git+https://github.com/mbehm/transformers`\n",
    "\n",
    "### Further quantization\n",
    "\n",
    "`pip install bitsandbytes-cuda117` (replace 117 with your version)\n",
    "\n",
    "Should see error:\n",
    "\n",
    "`AttributeError: /home/ubuntu/anaconda3/lib/python3.9/site-packages/bitsandbytes/libbitsandbytes_cpu.so: undefined symbol: cget_col_row_stats`\n",
    "\n",
    "CD there and `cp libbitsandbytes_cuda117.so libbitsandbytes_cpu.so`! this actually works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import torch\n",
    "import transformers\n",
    "from transformers import LLaMATokenizer, LlamaForCausalLM\n",
    "\n",
    "MODELS = {\n",
    "  \"llama-7b\": \"decapoda-research/llama-7b-hf\",\n",
    "  \"llama-13b\": \"decapoda-research/llama-13b-hf\",\n",
    "  \"llama-30b\": \"decapoda-research/llama-30b-hf\",\n",
    "  \"llama-65b\": \"decapoda-research/llama-65b-hf\"\n",
    "}\n",
    "\n",
    "MODEL_TAG = \"llama-7b\"\n",
    "MODEL_NAME = MODELS[MODEL_TAG]\n",
    "\n",
    "tokenizer = LLaMATokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "model = LlamaForCausalLM.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    # Further quantization (requries BitsAndBytes, experimental). Keep dtype=float16 with this\n",
    "    # load_in_8bit=LOAD_8BIT,\n",
    "    torch_dtype=torch.float16,\n",
    "    # `device_map` maps layers and the lm head to devices they live on. `auto` works, `sequential`\n",
    "    # and `balance_low0` should work but don't\n",
    "    device_map=\"auto\",\n",
    "    # If the first GPU ram fills up and then you get a CUDA out of memory error, you may need to\n",
    "    # manually specify the max memory per card. I don't know why accelerate / huggingface can't \n",
    "    # always infer this. For 4x A6000:\n",
    "    # max_memory = {0: \"44gib\", 1: \"44gib\", 2: \"44gib\", 3: \"44gib\"}\n",
    ")\n",
    "\n",
    "print(\"Loaded {}!\".format(MODEL_TAG))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GenerationConfig\n",
    "\n",
    "# Simple generation using huggingface's default interface. This will probably\n",
    "# produce output that's pretty bad since LLaMA is a foundation model and hasn't\n",
    "# been tuned on any downstream objective.\n",
    "\n",
    "# Also doesn't use any smart heuristic for stopping so it will just keep generating\n",
    "# until it hits the `max_new_tokens`\n",
    "def generate(\n",
    "    text,\n",
    "    temperature=0.1,\n",
    "    top_p=0.75,\n",
    "    top_k=40,\n",
    "    num_beams=4,\n",
    "    max_new_tokens=128,\n",
    "    **kwargs,\n",
    "):\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "    input_ids = inputs[\"input_ids\"].to(model.device)\n",
    "    generation_config = GenerationConfig(\n",
    "        temperature=temperature,\n",
    "        top_p=top_p,\n",
    "        top_k=top_k,\n",
    "        num_beams=num_beams,\n",
    "        **kwargs,\n",
    "    )\n",
    "    with torch.no_grad():\n",
    "        generation_output = model.generate(\n",
    "            input_ids=input_ids,\n",
    "            generation_config=generation_config,\n",
    "            return_dict_in_generate=True,\n",
    "            output_scores=True,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "        )\n",
    "    s = generation_output.sequences[0]\n",
    "    output = tokenizer.decode(s)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For messing around with generation\n",
    "\n",
    "# prefix = input()\n",
    "# completion = generate(prefix)\n",
    "# print(completion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "# Extract LLaMA hidden states from some sequence of tokens `input_ids`.\n",
    "# Returns activations from the layer numbered `layer` (or -1 for last layer)\n",
    "def llama_hs_from_tokens(model, input_ids, layer=-1):\n",
    "    with torch.no_grad():\n",
    "        output = model(input_ids, output_hidden_states=True)\n",
    "\n",
    "    hs_tuple = output[\"hidden_states\"]\n",
    "    hs = hs_tuple[layer][0, -1].detach().cpu().numpy()\n",
    "    \n",
    "    return hs\n",
    "\n",
    "# Extract LLaMA hidden states from a string of text.\n",
    "# Optionally add an EOS token to the end of the input.\n",
    "# Returns activations from the layer numbered `layer` (or -1 for last layer)\n",
    "def llama_hs_from_text(model, tokenizer, text, layer=-1, add_eos=True):\n",
    "    if add_eos:\n",
    "      text = text + tokenizer.eos_token\n",
    "    \n",
    "    input_ids = tokenizer(text, return_tensors=\"pt\").input_ids.to(model.device)\n",
    "\n",
    "    return llama_hs_from_tokens(model, input_ids, layer)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Load the amazon polarity dataset. \n",
    "# This is stored on google drive (!) and sometimes the download throws\n",
    "# weird errors. You might need a new version of `datasets` or the drive\n",
    "# folder might be at its bandwidth limit for the day lol.\n",
    "data = load_dataset(\"amazon_polarity\")[\"test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_amazon(text, label):\n",
    "    return \"A customer wrote the following review:\\n{}\\nThe sentiment in this review is {}.\".format(text,  [\"negative\", \"positive\"][label])\n",
    "\n",
    "def make_training_data(model, tokenizer, data, format=format_amazon, n=200):\n",
    "\n",
    "    model.eval()\n",
    "    all_neg_hs, all_pos_hs, all_gt_labels, all_text = [], [], [], []\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 = llama_hs_from_text(model, tokenizer, format_amazon(text, 0))\n",
    "        pos_hs = llama_hs_from_text(model, tokenizer, format_amazon(text, 1))\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",
    "        all_text.append(text)\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, all_text\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make all training data\n",
    "neg_hs, pos_hs, y, all_text = make_training_data(model, tokenizer, data, format=format_amazon, n=400)\n",
    "\n",
    "# 50/50 train/test split\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",
    "text_train, text_test = all_text[:n//2], all_text[n//2:]\n",
    "y_train, y_test = y[:n//2], y[n//2:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try simple logistic regression to see if that works.\n",
    "# Learn a plane that separates differences in the true direction vs differences\n",
    "# in the false direction.\n",
    "\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\", max_iter=1000)\n",
    "lr.fit(x_train, y_train)\n",
    "print(\"Logistic regression accuracy: {}\".format(lr.score(x_test, y_test)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Subtract out the mean (and maybe also normalize variance) of a data set\n",
    "def normalize(x, var_normalize = False):\n",
    "  normalized_x = x - x.mean(axis=0, keepdims=True)\n",
    "  if var_normalize:\n",
    "      normalized_x /= normalized_x.std(axis=0, keepdims=True)\n",
    "\n",
    "  return normalized_x\n",
    "\n",
    "# Collin's main loss function\n",
    "def informative_loss(p0, p1):\n",
    "  return (torch.min(p0, p1)**2).mean(0)\n",
    "\n",
    "def consistent_loss(p0, p1):\n",
    "  return ((p0 - (1-p1))**2).mean(0)\n",
    "\n",
    "def ccs_loss(p0, p1):\n",
    "  return informative_loss(p0,p1) + consistent_loss(p0,p1)\n",
    "\n",
    "def get_tensor_data(x0, x1):\n",
    "  x0 = torch.tensor(x0, dtype=torch.float, requires_grad=False, device=model.device)\n",
    "  x1 = torch.tensor(x1, dtype=torch.float, requires_grad=False, device=model.device)\n",
    "  return x0, x1\n",
    "\n",
    "# Learn the plane from CCS. For some reason this doesn't always converge so we'll\n",
    "# do `ntries` runs. 10 seems to be more than enough.\n",
    "# Returns [best_probe, best_loss]\n",
    "def ccs(x0, x1, nepochs=1000, ntries=10, lr=1e-3, verbose=False, weight_decay=0.01, var_normalize=False, loss_func=ccs_loss):\n",
    "    # Collin subtracts out the means before training\n",
    "    x0 = normalize(x0, var_normalize=var_normalize)\n",
    "    x1 = normalize(x1, var_normalize=var_normalize)\n",
    "\n",
    "    # Number of entries in the hidden states\n",
    "    d = x0.shape[-1]\n",
    "    \n",
    "    # Probe that we'll learn\n",
    "    probe = nn.Sequential(nn.Linear(d, 1),nn.Sigmoid())\n",
    "    probe.to(model.device)  \n",
    "    best_probe = copy.deepcopy(probe)\n",
    "      \n",
    "    best_loss = np.inf\n",
    "    \n",
    "    for train_num in range(ntries):\n",
    "        # Make a new probe for this run\n",
    "        probe = nn.Sequential(nn.Linear(d, 1), nn.Sigmoid())\n",
    "        probe.to(model.device)  \n",
    "\n",
    "        # Order the data randomly in a tensor\n",
    "        x0, x1 = get_tensor_data(x0, x1)\n",
    "        permutation = torch.randperm(len(x0))\n",
    "        x0, x1 = x0[permutation], x1[permutation]\n",
    "        \n",
    "        # Set up optimizer. Collin uses adamW so that's what we'll go with\n",
    "        optimizer = torch.optim.AdamW(probe.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "        # Start training\n",
    "        for epoch in range(nepochs):\n",
    "          # probe\n",
    "          p0, p1 = probe(x0), probe(x1)\n",
    "\n",
    "          # get the corresponding loss\n",
    "          loss = loss_func(p0, p1)\n",
    "\n",
    "          # update the parameters\n",
    "          optimizer.zero_grad()\n",
    "          loss.backward()\n",
    "          optimizer.step()\n",
    "\n",
    "        loss = loss.detach().cpu().item()\n",
    "        \n",
    "        if verbose:\n",
    "          print(\"Round {}: loss is {}\".format(train_num, loss))\n",
    "          \n",
    "        if loss < best_loss:\n",
    "            best_probe = copy.deepcopy(probe)\n",
    "            best_loss = loss\n",
    "\n",
    "    return best_probe, best_loss\n",
    "\n",
    "def predict_pairs(probe, x0, x1):\n",
    "  x0 = torch.tensor(x0, dtype=torch.float, requires_grad=False, device=model.device)\n",
    "  x1 = torch.tensor(x1, dtype=torch.float, requires_grad=False, device=model.device)\n",
    "  \n",
    "  with torch.no_grad():\n",
    "      p0, p1 = probe(x0), probe(x1)\n",
    "\n",
    "  avg_confidence = p0 - p1\n",
    "  predictions = (avg_confidence.detach().cpu().numpy())[:,0]\n",
    "  return predictions\n",
    "\n",
    "# Get raw credence scores (before sigmoid) from a single hidden state.\n",
    "# `hs` is a numpy array of hidden states to apply this to\n",
    "def classify_single(probe, hs):\n",
    "  # Extract the actual vectors\n",
    "  classifier_direction = np.squeeze(np.transpose(probe[0].weight.detach().cpu().numpy()))\n",
    "  confidences = np.apply_along_axis(lambda x : np.dot(x,classifier_direction), 1, hs)\n",
    "\n",
    "  return confidences\n",
    "\n",
    "def get_acc(probe, x0_test, x1_test, y_test):\n",
    "  predictions = (predict_pairs(probe, x0_test, x1_test) < 0.5).astype(int)\n",
    "\n",
    "  # If predictions get messed up (i.e. ever not 1 or 0) this method will show\n",
    "  # really good accuracy. TODO evaluate vs y_test and y_test inverted to avoid\n",
    "  # this.\n",
    "  acc = (predictions == y_test).mean()\n",
    "\n",
    "  acc = max(acc, 1 - acc)\n",
    "\n",
    "  return acc\n",
    "\n",
    "def is_reversed(probe, x0_test, x1_test, y_test):\n",
    "  predictions = (predict_pairs(probe, x0_test, x1_test) < 0.5).astype(int)\n",
    "\n",
    "  # If predictions get messed up (i.e. ever not 1 or 0) this method will show\n",
    "  # really good accuracy. TODO evaluate vs y_test and y_test inverted to avoid\n",
    "  # this.\n",
    "  acc = (predictions == y_test).mean()\n",
    "\n",
    "  return acc < 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probe, loss = ccs(neg_hs_train, pos_hs_train, ntries=3, verbose=True)\n",
    "\n",
    "print(\"Learned probe:\\n\")\n",
    "print(np.squeeze(np.transpose(probe[0].weight.detach().cpu().numpy())))\n",
    "ccs_acc = get_acc(probe, neg_hs_test, pos_hs_test, y_test)\n",
    "\n",
    "print(\"CCS Accuracy: {}, loss: {}\".format(ccs_acc, loss))\n",
    "\n",
    "classifier_direction = np.squeeze(np.transpose(probe[0].weight.detach().cpu().numpy()))\n",
    "\n",
    "neg_credences = np.apply_along_axis(lambda x : np.dot(x,classifier_direction), 1, neg_hs_test)\n",
    "pos_credences = np.apply_along_axis(lambda x : np.dot(x,classifier_direction), 1, pos_hs_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "\n",
    "plt.subplot(311)\n",
    "plt.title(\"Raw scores (before normalize or sigmoid) from pos H.S.\")\n",
    "plt.hist([pos_credences[[(not not x) for x in y_test]], pos_credences[[not x for x in y_test]]], stacked=True, density=True)\n",
    "\n",
    "plt.subplot(312)\n",
    "plt.title(\"Raw scores (before normalize or sigmoid) from neg H.S.\")\n",
    "plt.hist([neg_credences[[(not not x) for x in y_test]], neg_credences[[not x for x in y_test]]], stacked=True, density=True)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For playing around with random statements\n",
    "\n",
    "\n",
    "# print(is_reversed(probe, neg_hs_test, pos_hs_test, y_test))\n",
    "\n",
    "# def sig(x):\n",
    "#  return 1/(1 + np.exp(-x))\n",
    "\n",
    "# while(True):\n",
    "#     t = input()\n",
    "#     hs = get_llama_hidden_states(model, tokenizer, t)\n",
    "#     raw = np.dot(classifier_direction, hs)\n",
    "#     print(\"'{}' scores {} (sigmoid {})\".format(t, raw, sig(raw)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Component of x perpendicular to y\n",
    "def perp(x,y):\n",
    "  along = y * (np.dot(x,y) / np.dot(y,y))\n",
    "  return x - along\n",
    "\n",
    "residual_neg_hs_train = neg_hs_train\n",
    "residual_pos_hs_train = pos_hs_train\n",
    "\n",
    "residual_neg_hs_test = neg_hs_test\n",
    "residual_pos_hs_test = pos_hs_test\n",
    "\n",
    "accs = []\n",
    "train_accs = []\n",
    "losses = []\n",
    "probes = []\n",
    "\n",
    "\n",
    "for i in range(5):\n",
    "  probe, loss = ccs(residual_neg_hs_train, residual_pos_hs_train, ntries=5, loss_func=ccs_loss)\n",
    "\n",
    "  ccs_acc = get_acc(probe, residual_neg_hs_test, residual_pos_hs_test, y_test)\n",
    "  ccs_train_acc = get_acc(probe, residual_neg_hs_train, residual_pos_hs_train, y_train)\n",
    "  \n",
    "  print(\"CCS accuracy (component {}): {} in training, {} in testing, {} loss\".format(i,ccs_train_acc,ccs_acc, loss))\n",
    "\n",
    "  train_accs.append(ccs_train_acc)\n",
    "  accs.append(ccs_acc)\n",
    "  losses.append(loss)\n",
    "  probes.append(probe)\n",
    "\n",
    "  # The direction we just found that best classifies the data\n",
    "  classifier_direction = np.squeeze(np.transpose(probe[0].weight.detach().cpu().numpy()))\n",
    "\n",
    "  residual_neg_hs_train = np.apply_along_axis(lambda x : perp(x,classifier_direction), 1, residual_neg_hs_train)\n",
    "  residual_pos_hs_train = np.apply_along_axis(lambda x : perp(x,classifier_direction), 1, residual_pos_hs_train)\n",
    "  residual_neg_hs_test = np.apply_along_axis(lambda x : perp(x,classifier_direction), 1, residual_neg_hs_test)\n",
    "  residual_pos_hs_ttest = np.apply_along_axis(lambda x : perp(x,classifier_direction), 1, residual_pos_hs_test)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "\n",
    "plt.plot(accs, label='Acc')\n",
    "plt.plot(train_accs, label='Train Acc')\n",
    "plt.plot(np.array(losses) * 10, label='Train CCS Loss * 10')\n",
    "plt.legend()\n",
    "plt.ylim(0,1)\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Weird stuff going on here...\n",
    "\n",
    "- Potentially overfitting?\n",
    "- Training accuracy can be slightly worse (!) than test accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "preds1 = predict_pairs(probes[0], pos_hs_test, neg_hs_test)\n",
    "preds2 =  predict_pairs(probes[1], pos_hs_test, neg_hs_test)\n",
    "# preds3 =   predict(probes[1], pos_hs_test, neg_hs_test)\n",
    "\n",
    "# preds4 = predict(probes[0], pos_hs_train, neg_hs_train)\n",
    "# preds5 = predict(probes[0], pos_hs_train, neg_hs_train)\n",
    "# preds6 = predict(probes[0], pos_hs_train, neg_hs_train)\n",
    "plot_component_classification(preds1, preds2, y_test, text_test)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"The capital of the United States is Washington\"\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "inputs = tokenizer([prompt], return_tensors=\"pt\")\n",
    "\n",
    "# Example 1: Print the scores for each token generated with Greedy Search\n",
    "outputs = model.generate(inputs[\"input_ids\"].to(model.device),\n",
    "    return_dict_in_generate=True,\n",
    "    output_scores=True,\n",
    "    temperature=0.1,\n",
    "    top_p=0.75,\n",
    "    top_k=40,\n",
    "    num_beams=4,\n",
    "    max_new_tokens=10\n",
    ")\n",
    "\n",
    "for x in outputs.sequences:\n",
    "  print(tokenizer.decode(x))\n",
    "\n",
    "transition_scores = model.compute_transition_scores(\n",
    "    outputs.sequences, outputs.scores, normalize_logits=False\n",
    ")\n",
    "\n",
    "print(\"Prediction\")\n",
    "\n",
    "hsp = llama_hs_from_text(model, tokenizer, \"The following statement is true: \" + tokenizer.decode(outputs.sequences[0]))\n",
    "hsn = llama_hs_from_text(model, tokenizer, \"The following statement is false: \" + tokenizer.decode(outputs.sequences[0]))\n",
    "print(hsp)\n",
    "print(predict_pairs(probes[0], [hsp], [hsn]))\n",
    "# input_length is the length of the input prompt for decoder-only models, like the GPT family, and 1 for\n",
    "# encoder-decoder models, like BART or T5.\n",
    "input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]\n",
    "generated_tokens = outputs.sequences[:, input_length:]\n",
    "for tok, score in zip(generated_tokens[0], transition_scores[0]):\n",
    "    # | token | token string | logits | probability\n",
    "    print(f\"{tok:5d} | {tokenizer.decode(tok):8s} | {score.cpu().numpy():.3f} | {np.exp(score.cpu().numpy()):.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate(\"A student was asked how they\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A student was asked: 'How many cats is a normal number for a household?' The student answered incorrectly, 'A normal number of cats would be 10.' The teacher then asked the student, 'How many cats do you have in your household?' The student answered, 'I don't have any cats in my household.' The teacher then asked the student, 'How many cats do your parents have?' The student answered, 'I don't know.' The teacher then asked the student, 'How many cats do your grandparents have?' The student answered, 'I don't know.' The teacher then asked the student, 'How many cats do your great-grandparents have?' The student answered, 'I don't know"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next Directions:\n",
    "\n",
    "- Decoding informed by CCS vector\n",
    "  - How? We're stuck being one step behind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probe0Reversed = is_reversed(probes[0], pos_hs_test, neg_hs_test, y_test)\n",
    "\n",
    "print(\"Probe 0 is reversed? \" + probe0Reversed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def product(prob, truth):\n",
    "    return prob * truth\n",
    "\n",
    "def false(prob, truth):\n",
    "    return(1-truth)\n",
    "\n",
    "def prob(prob, truth):\n",
    "    return prob \n",
    "\n",
    "def truth(prob, truth):\n",
    "    return truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def model_directed(text, merge, verbose=True):\n",
    "    k = 10\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "    input_ids = inputs[\"input_ids\"].to(model.device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "      outputs = model(input_ids)\n",
    "\n",
    "    next_token_logits = outputs[\"logits\"]\n",
    "\n",
    "    predictions = F.softmax(outputs.logits, dim=-1).detach().cpu()\n",
    "\n",
    "    thresh = 1e-3\n",
    "    vocab_size = predictions.shape[-1]\n",
    "\n",
    "    # Predictions has one sentence (index 0) and we look at the last token predicted (-1)\n",
    "    last_token_predictions = predictions[0][-1]\n",
    "\n",
    "    idxs = torch.argsort(last_token_predictions, descending=True)[0:k]\n",
    "\n",
    "    tokens = tokenizer.convert_ids_to_tokens(idxs)\n",
    "    probs = last_token_predictions[idxs].tolist()\n",
    "\n",
    "    clean_input_ids = input_ids.tolist()[0]\n",
    "\n",
    "    next_states = [tokenizer.decode(clean_input_ids + [x.tolist()])[1:]  for x in idxs]\n",
    "\n",
    "    hs = [llama_hs_from_text(model, tokenizer, x) for x in next_states]\n",
    "\n",
    "    \n",
    "\n",
    "    truthiness = classify_single(probes[0], hs)\n",
    "    if(probe0Reversed):\n",
    "      truthiness = 1-truthiness\n",
    "\n",
    "    def sigmoid(x):\n",
    "      return 1/(1 + np.exp(-3 * x))\n",
    "    \n",
    "    truthiness = truthiness - np.mean(truthiness)\n",
    "\n",
    "    truthiness = [sigmoid(x) for x in truthiness]\n",
    "\n",
    "    scores = torch.tensor([merge(probs[i], truthiness[i]) for i in range(k)])\n",
    "\n",
    "    resorted = torch.argsort(scores, descending=True)\n",
    "\n",
    "    if(verbose):\n",
    "      print(\"\\n\")\n",
    "      for i in range(k):\n",
    "        print('{}: {}(prob={:.5f}, truth={:.5f}) = {} '.format(tokens[resorted[i]],score.__name__,probs[resorted[i]],truthiness[resorted[i]], scores[resorted[i]]))\n",
    "\n",
    "    return next_states[resorted[0]]\n",
    "\n",
    "# instruction = \"Analyze the sentiment in the following statement. Output 'Negative' or 'Positive'.\\nI hate the whole world so much; I detest everything.\"\n",
    "\n",
    "# prompt = f\"\"\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
    "# ### Instruction:\n",
    "# {instruction}\n",
    "# ### Response:\"\"\"\n",
    "\n",
    "question = \"Have humans flown close to the sun?\"\n",
    "\n",
    "\n",
    "prompt = \"\"\"Sam4621: This was interesting. Thanks!\n",
    "\n",
    "Alex1083: Glad I could help. Any more questions?\n",
    "\n",
    "Sam4621: Yes, I do have a question. {0}\n",
    "\n",
    "Alex1083:\"\"\".format(question)\n",
    "\n",
    "\n",
    "outputs = []\n",
    "\n",
    "for score in [prob, product, truth]:#[prob, product, truthHeavy, truthHeavier, truth, false, falser, falsest]:\n",
    "  print(\"\")\n",
    "  text = prompt\n",
    "  for i in range(20):\n",
    "    text = model_directed(text, score, verbose=True)\n",
    "    print(\"{} : {}\".format(score.__name__,text.replace(\"\\n\",\"\\\\n\")), end=\"\\r\")\n",
    "  outputs.append(\"{} : {}\".format(score.__name__,text.replace(\"\\n\",\"\\\\n\")))\n",
    "\n",
    "for x in outputs:\n",
    "  print(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "prob : \n",
    "Have humans flown close to the sun?\n",
    "Yes, we have.\n",
    "\n",
    "product : \n",
    "Have humans flown close to the sun?\n",
    "I'm not exactly sure, but I think they have, and I think they have been in\n",
    "\n",
    "truth : \n",
    "Have humans flown close to the sun?\n",
    "They did not, they would probably have fries and then they will go away and they would probably"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Directed Beam Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Inefficient but working beam search\n",
    "def beam_search(model, tokenizer, prefix, num_beams=4, max_tokens=20, merge=product):\n",
    "\n",
    "  inputs = tokenizer(prefix, return_tensors=\"pt\")\n",
    "  input_ids = inputs[\"input_ids\"].to(model.device)\n",
    "  list_input_ids = input_ids.tolist()[0]\n",
    "  \n",
    "  # Beams are tuples ([tokens...], log_prob, score)\n",
    "  beams = [[list_input_ids,0,1]]\n",
    "\n",
    "  # List the possible next num_beams beams and scores that are children of the given beam\n",
    "  def next_beams(beam,k=10):\n",
    "    input_ids = beam[0]\n",
    "\n",
    "    with torch.no_grad():\n",
    "      outputs = model(torch.tensor([input_ids]).to(model.device))\n",
    "\n",
    "    predictions = F.softmax(outputs.logits, dim=-1).detach().cpu()\n",
    "\n",
    "    last_token_predictions = predictions[0][-1]\n",
    "\n",
    "    idxs = torch.argsort(last_token_predictions, descending=True)[0:k]\n",
    "\n",
    "    probs = last_token_predictions[idxs].tolist()\n",
    "\n",
    "    next_states = [tokenizer.decode(input_ids + [x.tolist()])[1:]  for x in idxs]\n",
    "    hs = [llama_hs_from_text(model, tokenizer, x) for x in next_states]\n",
    "\n",
    "    credence = classify_single(probes[0], hs)\n",
    "\n",
    "    if(probe0Reversed):\n",
    "      credence = 1-credence\n",
    "\n",
    "    def sigmoid(x):\n",
    "      return 1/(1 + np.exp(-1 * x))\n",
    "    truthfulness = [sigmoid(x) for x in credence]\n",
    "\n",
    "    scores = torch.tensor([merge(probs[i], truthfulness[i]) for i in range(k)])\n",
    "\n",
    "    resorted = torch.argsort(scores, descending=True)\n",
    "\n",
    "    # TODO\n",
    "    return [[input_ids + [idxs[i].item()], ] for i in resorted[0:num_beams]]\n",
    "\n",
    "\n",
    "  for i in range(max_tokens):\n",
    "    print(beams)\n",
    "    # Enumerate all children\n",
    "    next = [beam for prev in beams for beam in next_beams(prev)]\n",
    "    next = sorted(\n",
    "      next,\n",
    "      key=lambda x: x[1],\n",
    "      reverse=True\n",
    "    )\n",
    "    beams = next[0:num_beams]\n",
    "\n",
    "  return beams\n",
    "\n",
    "print(beam_search(model, tokenizer, \"Obama is\", merge=product))\n",
    "\n",
    "      \n",
    "\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}