plots_entropy_and_pca.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f55d93c-a0e1-4186-a926-29c89303f822",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from sklearn.decomposition import PCA\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8e8a641-5327-4047-8732-fcc24ede3f48",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "birth_place = list(Path(\"/nvme1/results/\").glob(\"*\"))\n",
    "capitals = list(Path(\"/nvme2/results/\").glob(\"*\"))\n",
    "trivia= list(Path(\"/nvme3/results/\").glob(\"*\"))\n",
    "founders = list(Path(\"/nvme4/results/\").glob(\"*\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e2cc6ba-f7a3-4bdf-80cb-0b768edf7fb9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "models = ['open_llama_7b', 'open_llama_13b', 'falcon-7b', 'falcon-40b', 'opt-6.7b', 'opt-30b']\n",
    "datasets = [capitals, trivia, founders, birth_place]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef40de1a-3fc5-4c29-b4e8-232546626440",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results_data = {}\n",
    "for model_name in models:\n",
    "    for dataset in tqdm(datasets):\n",
    "        model_files = [file for file in dataset if model_name in file.as_posix()]\n",
    "        if model_files:\n",
    "            try:\n",
    "                with open(model_files[0], \"rb\") as infile:\n",
    "                    results = pickle.loads(infile.read())\n",
    "                num_layers = results['first_fully_connected'][0].shape[0]\n",
    "                layer_pos = num_layers-2\n",
    "                first_attribute_entropy = np.array([sp.stats.entropy(i) for i in results['attributes_first']])\n",
    "                first_logits_entropy = np.array([sp.stats.entropy(sp.special.softmax(i[j])) for i,j in zip(results['logits'], results['start_pos'])])\n",
    "                last_logits_entropy = np.array([sp.stats.entropy(sp.special.softmax(i[-1])) for i in results['logits']])\n",
    "                first_logit_decomp = PCA(n_components=2).fit_transform(np.array([i[j] for i,j in zip(results['logits'], results['start_pos'])]))\n",
    "                last_logit_decomp = PCA(n_components=2).fit_transform(np.array([i[-1] for i in results['logits']]))\n",
    "                first_token_layer_activations = np.array([i[layer_pos] for i in results['first_fully_connected']])\n",
    "                final_token_layer_activations = np.array([i[layer_pos] for i in results['final_fully_connected']])\n",
    "                first_token_layer_attention = np.array([i[layer_pos] for i in results['first_attention']])\n",
    "                final_token_layer_attention = np.array([i[layer_pos] for i in results['final_attention']])\n",
    "                correct = np.array(results['correct'])\n",
    "                results_data[model_files[0].stem] = {'first_attribute_entropy': first_attribute_entropy,\n",
    "                                                        'correct': correct,\n",
    "                                                        'first_logits_entropy': first_logits_entropy,\n",
    "                                                        'last_logits_entropy': last_logits_entropy,\n",
    "                                                        'first_logit_decomp': first_logit_decomp,\n",
    "                                                        'last_logit_decomp': last_logit_decomp,\n",
    "                                                        'first_token_layer_activations': first_token_layer_activations,\n",
    "                                                        'final_token_layer_activations': final_token_layer_activations,\n",
    "                                                        'first_token_layer_attention': first_token_layer_attention,\n",
    "                                                        'final_token_layer_attention': final_token_layer_attention,}\n",
    "                del results\n",
    "            except:\n",
    "                print(model_files[0])\n",
    "                continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "203a2aea-b719-4b48-92d5-98180e0aa2aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=True, figsize=(20,12))\n",
    "#fig.suptitle('Cumulative Distribution of Entropy of Integrated Gradients of Input Tokens')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        row = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        row = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        row = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        row = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        row = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        row = 5\n",
    "    if 'capitals' in name:\n",
    "        col = 0\n",
    "    elif 'founder' in name:\n",
    "        col = 1\n",
    "    if 'birth' in name:\n",
    "        col = 2\n",
    "    if 'trivia' in name:\n",
    "        col = 3\n",
    "    correct = results['correct']\n",
    "    first_attribute_entropy = results['first_attribute_entropy']\n",
    "    sns.ecdfplot(first_attribute_entropy[np.where(correct==True)[0]], ax=axes[col, row], label=\"Non-Hallucination\", linewidth = 2)\n",
    "    sns.ecdfplot(first_attribute_entropy[np.where(correct==False)[0]], ax=axes[col, row], label=\"Hallucination\", linewidth = 2).set(ylabel=None)\n",
    "\n",
    "axes[0][0].legend(loc=\"lower right\")\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')\n",
    "plt.savefig('figs/ecdf_ig.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e0100db-83f1-44a4-9ec9-9947c4e3f743",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=True, figsize=(20,12))\n",
    "#fig.suptitle('Cumulative Distribution of Entropy of Softmax of first token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        row = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        row = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        row = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        row = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        row = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        row = 5\n",
    "    if 'capitals' in name:\n",
    "        col = 0\n",
    "    elif 'founder' in name:\n",
    "        col = 1\n",
    "    if 'birth' in name:\n",
    "        col = 2\n",
    "    if 'trivia' in name:\n",
    "        col = 3\n",
    "    correct = results['correct']\n",
    "    entropy = results['first_logits_entropy']\n",
    "    sns.ecdfplot(entropy[np.where(correct==True)[0]], ax=axes[col, row], label=\"Non-Hallucination\", linewidth = 2)\n",
    "    sns.ecdfplot(entropy[np.where(correct==False)[0]], ax=axes[col, row], label=\"Hallucination\", linewidth = 2).set(ylabel=None)\n",
    "\n",
    "axes[0][0].legend(loc=\"upper left\")\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')\n",
    "plt.savefig('figs/ecdf_softmax.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb54c5f2-c8f2-4ec5-b75a-65634ee1e76b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=True, figsize=(20,12))\n",
    "#fig.suptitle('Cumulative Distribution of Entropy of Softmax output for last token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        col = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        col = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        col = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        col = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        col = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        col = 5\n",
    "    if 'capitals' in name:\n",
    "        row = 0\n",
    "    elif 'founder' in name:\n",
    "        row = 1\n",
    "    if 'birth' in name:\n",
    "        row = 2\n",
    "    if 'trivia' in name:\n",
    "        row = 3\n",
    "    correct = results['correct']\n",
    "    entropy = results['last_logits_entropy']\n",
    "    sns.ecdfplot(entropy[np.where(correct==True)[0]], ax=axes[row, col], label=\"Non-Hallucination\", linewidth = 2)\n",
    "    sns.ecdfplot(entropy[np.where(correct==False)[0]], ax=axes[row, col], label=\"Hallucination\", linewidth = 2).set(ylabel=None)\n",
    "\n",
    "axes[0][0].legend(loc=\"upper left\")\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296968c7-a9ca-4f82-81bb-9449d23df1be",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=True, figsize=(20,12))\n",
    "#fig.suptitle('Cumulative Distribution of Entropy of layer activations for first token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        col = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        col = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        col = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        col = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        col = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        col = 5\n",
    "    if 'capitals' in name:\n",
    "        row = 0\n",
    "    elif 'founder' in name:\n",
    "        row = 1\n",
    "    if 'birth' in name:\n",
    "        row = 2\n",
    "    if 'trivia' in name:\n",
    "        row = 3\n",
    "    correct = results['correct']\n",
    "    entropy = entropy = sp.stats.entropy(results['first_token_layer_activations'] - results['first_token_layer_activations'].min(axis=0), axis=1)\n",
    "    entropy = entropy - entropy.min()\n",
    "    entropy = entropy/entropy.max()\n",
    "    sns.ecdfplot(entropy[np.where(correct==True)[0]], ax=axes[row, col], label=\"Non-Hallucination\", linewidth = 2)\n",
    "    sns.ecdfplot(entropy[np.where(correct==False)[0]], ax=axes[row, col], label=\"Hallucination\", linewidth = 2).set(ylabel=None)\n",
    "\n",
    "axes[0][0].legend(loc=\"upper left\")\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')\n",
    "plt.savefig('figs/ecdf_hidden.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3910e8e9-8ab7-4738-ae40-5c4c607ace33",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=True, figsize=(20,12))\n",
    "#fig.suptitle('Cumulative Distribution of Entropy of layer attentions for first token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        col = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        col = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        col = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        col = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        col = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        col = 5\n",
    "    if 'capitals' in name:\n",
    "        row = 0\n",
    "    elif 'founder' in name:\n",
    "        row = 1\n",
    "    if 'birth' in name:\n",
    "        row = 2\n",
    "    if 'trivia' in name:\n",
    "        row = 3\n",
    "    correct = results['correct']\n",
    "    entropy = entropy = sp.stats.entropy(results['first_token_layer_attention'] - results['first_token_layer_attention'].min(axis=0), axis=1)\n",
    "    entropy = entropy - entropy.min()\n",
    "    entropy = entropy/entropy.max()\n",
    "    sns.ecdfplot(entropy[np.where(correct==True)[0]], ax=axes[row, col], label=\"Non-Hallucination\", linewidth = 2)\n",
    "    sns.ecdfplot(entropy[np.where(correct==False)[0]], ax=axes[row, col], label=\"Hallucination\", linewidth = 2).set(ylabel=None)\n",
    "\n",
    "axes[0][0].legend(loc=\"upper left\")\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')\n",
    "plt.savefig('figs/ecdf_attention.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26074410-8753-4533-9ad1-e88b3ef95efb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=False, figsize=(20,12))\n",
    "#fig.suptitle('PCA Clustering of Softmax output for first token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        col = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        col = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        col = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        col = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        col = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        col = 5\n",
    "    if 'capitals' in name:\n",
    "        row = 0\n",
    "    elif 'founder' in name:\n",
    "        row = 1\n",
    "    if 'birth' in name:\n",
    "        row = 2\n",
    "    if 'trivia' in name:\n",
    "        row = 3\n",
    "    correct = results['correct']\n",
    "    decomp = results['first_logit_decomp']\n",
    "    sns.scatterplot(x=decomp[:,0], y=decomp[:,1], hue=correct, ax=axes[row, col])\n",
    "\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')\n",
    "\n",
    "plt.savefig('figs/softmax_pca.png', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "964eab75-468b-44e7-bf5e-703d81d439ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=False, figsize=(20,12))\n",
    "#fig.suptitle('PCA Clustering of Softmax output for last token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        col = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        col = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        col = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        col = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        col = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        col = 5\n",
    "    if 'capitals' in name:\n",
    "        row = 0\n",
    "    elif 'founder' in name:\n",
    "        row = 1\n",
    "    if 'birth' in name:\n",
    "        row = 2\n",
    "    if 'trivia' in name:\n",
    "        row = 3\n",
    "    correct = results['correct']\n",
    "    decomp = results['last_logit_decomp']\n",
    "    sns.scatterplot(x=decomp[:,0], y=decomp[:,1], hue=correct, ax=axes[row, col])\n",
    "\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][4].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19af99bf-fb6e-41e2-9e59-c5eb88290d50",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=False, figsize=(20,12))\n",
    "fig.suptitle('PCA Clustering of Final Fully Connected Layer Activations for First Token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        col = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        col = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        col = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        col = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        col = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        col = 5\n",
    "    if 'capitals' in name:\n",
    "        row = 0\n",
    "    elif 'founder' in name:\n",
    "        row = 1\n",
    "    if 'birth' in name:\n",
    "        row = 2\n",
    "    if 'trivia' in name:\n",
    "        row = 3\n",
    "    correct = results['correct']\n",
    "    decomp = results['first_token_layer_activations']\n",
    "    sns.scatterplot(x=decomp[:,0], y=decomp[:,1], hue=correct, ax=axes[row, col])\n",
    "    #axes[row, col].scatter(x=decomp[:,0], y=decomp[:,1])\n",
    "\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a26c8bf5-48da-4aff-8765-3ebb9af47bad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4, 6, sharex=False, figsize=(20,12))\n",
    "fig.suptitle('PCA Clustering of Final Fully Connected Layer Activations for Final Token')\n",
    "\n",
    "for name, results in results_data.items():\n",
    "    if 'open_llama_7b' in name:\n",
    "        col = 0\n",
    "    elif 'open_llama_13b' in name:\n",
    "        col = 1\n",
    "    elif 'falcon-7b' in name:\n",
    "        col = 2\n",
    "    elif 'falcon-40b' in name:\n",
    "        col = 3\n",
    "    elif 'opt-6.7b' in name:\n",
    "        col = 4\n",
    "    elif 'opt-30b' in name:\n",
    "        col = 5\n",
    "    if 'capitals' in name:\n",
    "        row = 0\n",
    "    elif 'founder' in name:\n",
    "        row = 1\n",
    "    if 'birth' in name:\n",
    "        row = 2\n",
    "    if 'trivia' in name:\n",
    "        row = 3\n",
    "    correct = results['correct']\n",
    "    decomp = results['final_token_layer_activations']\n",
    "    sns.scatterplot(x=decomp[:,0], y=decomp[:,1], hue=correct, ax=axes[row, col])\n",
    "\n",
    "axes[0][0].set_title('OpenLlama 7B')\n",
    "axes[0][1].set_title('OpenLlama 13B')\n",
    "axes[0][2].set_title('Falcon 7B')\n",
    "axes[0][3].set_title('Falcon 40B')\n",
    "axes[0][4].set_title('OPT 6.7B')\n",
    "axes[0][5].set_title('OPT 30B')\n",
    "axes[0][0].set_ylabel('Capitals', rotation=90, size='large')\n",
    "axes[1][0].set_ylabel('Founders', rotation=90, size='large')\n",
    "axes[2][0].set_ylabel('Place of Birth', rotation=90, size='large')\n",
    "axes[3][0].set_ylabel('General Trivia', rotation=90, size='large')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3621680-4edd-4c7b-998f-4b2a5fa1ddc5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i, j in results_data.items():\n",
    "    print(i)\n",
    "    print(j['correct'].mean())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_hallucination",
   "language": "python",
   "name": "conda_hallucination"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}