Layer_Permutation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:DEPRECATED: Library has been renamed, import transformer_lens instead\n"
     ]
    }
   ],
   "source": [
    "import einops\n",
    "from fancy_einsum import einsum\n",
    "from dataclasses import dataclass\n",
    "from easy_transformer import EasyTransformer\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import math\n",
    "from transformer_lens.utils import get_corner, gelu_new, tokenize_and_concatenate\n",
    "import tqdm.auto as tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained model gpt2-small into HookedTransformer\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "model = EasyTransformer.from_pretrained(\n",
    "    \"gpt2-small\",\n",
    "    center_unembed=True,\n",
    "    center_writing_weights=True,\n",
    "    fold_ln=True,\n",
    "    refactor_factored_attn_matrices=True,\n",
    ")\n",
    "\n",
    "import transformer_lens.utils as utils\n",
    "# Get the default device used\n",
    "device: torch.device = utils.get_device()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokenized prompt: ['<|endoftext|>', 'After', ' John', ' and', ' Mary', ' went', ' to', ' the', ' store', ',', ' John', ' gave', ' a', ' bottle', ' of', ' milk', ' to']\n",
      "Tokenized answer: [' Mary']\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Performance on answer token:\n",
       "<span style=\"font-weight: bold\">Rank: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">        Logit: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">18.09</span><span style=\"font-weight: bold\"> Prob: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">70.07</span><span style=\"font-weight: bold\">% Token: | Mary|</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Performance on answer token:\n",
       "\u001b[1mRank: \u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m        Logit: \u001b[0m\u001b[1;36m18.09\u001b[0m\u001b[1m Prob: \u001b[0m\u001b[1;36m70.07\u001b[0m\u001b[1m% Token: | Mary|\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top 0th token. Logit: 18.09 Prob: 70.07% Token: | Mary|\n",
      "Top 1th token. Logit: 15.38 Prob:  4.67% Token: | the|\n",
      "Top 2th token. Logit: 15.35 Prob:  4.54% Token: | John|\n",
      "Top 3th token. Logit: 15.25 Prob:  4.11% Token: | them|\n",
      "Top 4th token. Logit: 14.84 Prob:  2.73% Token: | his|\n",
      "Top 5th token. Logit: 14.06 Prob:  1.24% Token: | her|\n",
      "Top 6th token. Logit: 13.54 Prob:  0.74% Token: | a|\n",
      "Top 7th token. Logit: 13.52 Prob:  0.73% Token: | their|\n",
      "Top 8th token. Logit: 13.13 Prob:  0.49% Token: | Jesus|\n",
      "Top 9th token. Logit: 12.97 Prob:  0.42% Token: | him|\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Ranks of the answer tokens:</span> <span style=\"font-weight: bold\">[(</span><span style=\"color: #008000; text-decoration-color: #008000\">' Mary'</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">)]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' Mary'\u001b[0m, \u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "example_prompt = \"After John and Mary went to the store, John gave a bottle of milk to\"\n",
    "example_answer = \" Mary\"\n",
    "utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[50256,    40,   716,   281,  4998,  1960,   382, 19741,    11,   875,\n",
      "         12342,    12,  8807,    11,   402, 11571,    12,    17,  3918, 47385,\n",
      "            13,  1881,  1110,   314,   481,  7074,  1692,  1241,  4430,   290,\n",
      "          1011,   625,   262,   995,     0]], device='mps:0')\n",
      "torch.Size([1, 35])\n",
      "['<|endoftext|>', 'I', ' am', ' an', ' amazing', ' aut', 'ore', 'gressive', ',', ' dec', 'oder', '-', 'only', ',', ' G', 'PT', '-', '2', ' style', ' transformer', '.', ' One', ' day', ' I', ' will', ' exceed', ' human', ' level', ' intelligence', ' and', ' take', ' over', ' the', ' world', '!']\n"
     ]
    }
   ],
   "source": [
    "reference_text = \"I am an amazing autoregressive, decoder-only, GPT-2 style transformer. One day I will exceed human level intelligence and take over the world!\"\n",
    "tokens = model.to_tokens(reference_text)\n",
    "print(tokens)\n",
    "print(tokens.shape)\n",
    "print(model.to_str_tokens(tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def complete(reference_text, max_tokens=100, T=0.7):\n",
    "    tokens = model.to_tokens(reference_text)\n",
    "    for i in range(max_tokens):\n",
    "        tokens = tokens.to(device)\n",
    "        logits, cache = model.run_with_cache(tokens)\n",
    "        \n",
    "        # Apply temperature scaling\n",
    "        scaled_logits = logits / T\n",
    "        \n",
    "        # Convert logits to probabilities\n",
    "        probs = torch.nn.functional.softmax(scaled_logits, dim=-1)\n",
    "        \n",
    "        # Sample from the probability distribution\n",
    "        next_token = torch.multinomial(probs[0, -1], num_samples=1)\n",
    "        \n",
    "        # Concatenate the new token to the existing sequence\n",
    "        tokens = torch.cat([tokens, next_token.unsqueeze(0)], dim=-1)\n",
    "    \n",
    "    # Decode the tokens to text\n",
    "    return model.tokenizer.decode(tokens[0]), cache\n",
    "\n",
    "text, cache = complete(reference_text, max_tokens=20, T=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hook_embed torch.Size([1, 54, 768])\n",
      "hook_pos_embed torch.Size([1, 54, 768])\n",
      "blocks.0.hook_resid_pre torch.Size([1, 54, 768])\n",
      "blocks.0.ln1.hook_scale torch.Size([1, 54, 1])\n",
      "blocks.0.ln1.hook_normalized torch.Size([1, 54, 768])\n",
      "blocks.0.attn.hook_q torch.Size([1, 54, 12, 64])\n",
      "blocks.0.attn.hook_k torch.Size([1, 54, 12, 64])\n",
      "blocks.0.attn.hook_v torch.Size([1, 54, 12, 64])\n",
      "blocks.0.attn.hook_attn_scores torch.Size([1, 12, 54, 54])\n",
      "blocks.0.attn.hook_pattern torch.Size([1, 12, 54, 54])\n",
      "blocks.0.attn.hook_z torch.Size([1, 54, 12, 64])\n",
      "blocks.0.hook_attn_out torch.Size([1, 54, 768])\n",
      "blocks.0.hook_resid_mid torch.Size([1, 54, 768])\n",
      "blocks.0.ln2.hook_scale torch.Size([1, 54, 1])\n",
      "blocks.0.ln2.hook_normalized torch.Size([1, 54, 768])\n",
      "blocks.0.mlp.hook_pre torch.Size([1, 54, 3072])\n",
      "blocks.0.mlp.hook_post torch.Size([1, 54, 3072])\n",
      "blocks.0.hook_mlp_out torch.Size([1, 54, 768])\n",
      "blocks.0.hook_resid_post torch.Size([1, 54, 768])\n",
      "ln_final.hook_scale torch.Size([1, 54, 1])\n",
      "ln_final.hook_normalized torch.Size([1, 54, 768])\n"
     ]
    }
   ],
   "source": [
    "for activation_name, activation in cache.cache_dict.items():\n",
    "    # Only print for first layer\n",
    "    if \".0.\" in activation_name or \"blocks\" not in activation_name:\n",
    "        print(activation_name, activation.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embed.W_E torch.Size([50257, 768])\n",
      "pos_embed.W_pos torch.Size([1024, 768])\n",
      "blocks.0.attn.W_Q torch.Size([12, 768, 64])\n",
      "blocks.0.attn.W_K torch.Size([12, 768, 64])\n",
      "blocks.0.attn.W_V torch.Size([12, 768, 64])\n",
      "blocks.0.attn.W_O torch.Size([12, 64, 768])\n",
      "blocks.0.attn.b_Q torch.Size([12, 64])\n",
      "blocks.0.attn.b_K torch.Size([12, 64])\n",
      "blocks.0.attn.b_V torch.Size([12, 64])\n",
      "blocks.0.attn.b_O torch.Size([768])\n",
      "blocks.0.mlp.W_in torch.Size([768, 3072])\n",
      "blocks.0.mlp.b_in torch.Size([3072])\n",
      "blocks.0.mlp.W_out torch.Size([3072, 768])\n",
      "blocks.0.mlp.b_out torch.Size([768])\n",
      "unembed.W_U torch.Size([768, 50257])\n",
      "unembed.b_U torch.Size([50257])\n"
     ]
    }
   ],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    # Only print for first layer\n",
    "    if \".0.\" in name or \"blocks\" not in name:\n",
    "        print(name, param.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HookedTransformerConfig:\n",
      "{'act_fn': 'gelu_new',\n",
      " 'attention_dir': 'causal',\n",
      " 'attn_only': False,\n",
      " 'attn_types': None,\n",
      " 'checkpoint_index': None,\n",
      " 'checkpoint_label_type': None,\n",
      " 'checkpoint_value': None,\n",
      " 'd_head': 64,\n",
      " 'd_mlp': 3072,\n",
      " 'd_model': 768,\n",
      " 'd_vocab': 50257,\n",
      " 'd_vocab_out': 50257,\n",
      " 'default_prepend_bos': True,\n",
      " 'device': device(type='mps'),\n",
      " 'dtype': torch.float32,\n",
      " 'eps': 1e-05,\n",
      " 'final_rms': False,\n",
      " 'from_checkpoint': False,\n",
      " 'gated_mlp': False,\n",
      " 'init_mode': 'gpt2',\n",
      " 'init_weights': False,\n",
      " 'initializer_range': 0.02886751345948129,\n",
      " 'model_name': 'gpt2',\n",
      " 'n_ctx': 1024,\n",
      " 'n_devices': 1,\n",
      " 'n_heads': 12,\n",
      " 'n_layers': 12,\n",
      " 'n_params': 84934656,\n",
      " 'normalization_type': 'LNPre',\n",
      " 'original_architecture': 'GPT2LMHeadModel',\n",
      " 'parallel_attn_mlp': False,\n",
      " 'positional_embedding_type': 'standard',\n",
      " 'rotary_dim': None,\n",
      " 'scale_attn_by_inverse_layer_idx': False,\n",
      " 'seed': None,\n",
      " 'tokenizer_name': 'gpt2',\n",
      " 'tokenizer_prepends_bos': False,\n",
      " 'use_attn_in': False,\n",
      " 'use_attn_result': False,\n",
      " 'use_attn_scale': True,\n",
      " 'use_hook_mlp_in': False,\n",
      " 'use_hook_tokens': False,\n",
      " 'use_local_attn': False,\n",
      " 'use_split_qkv_input': False,\n",
      " 'window_size': None}\n"
     ]
    }
   ],
   "source": [
    "reference_gpt2 = model\n",
    "\n",
    "# As a reference - note there's a lot of stuff we don't care about in here, to do with library internals or other architectures\n",
    "print(reference_gpt2.cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Config(d_model=768, debug=True, layer_norm_eps=1e-05, d_vocab=50257, init_range=0.02, n_ctx=1024, d_head=64, d_mlp=3072, n_heads=12, n_layers=12)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "@dataclass\n",
    "class Config:\n",
    "    d_model: int = 768\n",
    "    debug: bool = True\n",
    "    layer_norm_eps: float = 1e-5\n",
    "    d_vocab: int = 50257\n",
    "    init_range: float = 0.02\n",
    "    n_ctx: int = 1024\n",
    "    d_head: int = 64\n",
    "    d_mlp: int = 3072\n",
    "    n_heads: int = 12\n",
    "    n_layers: int = 12\n",
    "\n",
    "cfg = Config()\n",
    "print(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.w = nn.Parameter(torch.ones(cfg.d_model))\n",
    "        self.b = nn.Parameter(torch.zeros(cfg.d_model))\n",
    "    \n",
    "    def forward(self, residual):\n",
    "        # residual: [batch, position, d_model]\n",
    "        if self.cfg.debug: print(\"Residual:\", residual.shape)\n",
    "        residual = residual - einops.reduce(residual, \"batch position d_model -> batch position 1\", \"mean\")\n",
    "        # Calculate the variance, square root it. Add in an epsilon to prevent divide by zero.\n",
    "        scale = (einops.reduce(residual.pow(2), \"batch position d_model -> batch position 1\", \"mean\") + cfg.layer_norm_eps).sqrt()\n",
    "        normalized = residual / scale\n",
    "        normalized = normalized * self.w + self.b\n",
    "        if self.cfg.debug: print(\"Normalized:\", residual.shape)\n",
    "        return normalized\n",
    "\n",
    "class Embed(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.W_E = nn.Parameter(torch.empty((cfg.d_vocab, cfg.d_model)))\n",
    "        nn.init.normal_(self.W_E, std=self.cfg.init_range)\n",
    "    \n",
    "    def forward(self, tokens):\n",
    "        # tokens: [batch, position]\n",
    "        if self.cfg.debug: print(\"Tokens:\", tokens.shape)\n",
    "        embed = self.W_E[tokens, :] # [batch, position, d_model]\n",
    "        if self.cfg.debug: print(\"Embeddings:\", embed.shape)\n",
    "        return embed\n",
    "\n",
    "class PosEmbed(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.W_pos = nn.Parameter(torch.empty((cfg.n_ctx, cfg.d_model)))\n",
    "        nn.init.normal_(self.W_pos, std=self.cfg.init_range)\n",
    "    \n",
    "    def forward(self, tokens):\n",
    "        # tokens: [batch, position]\n",
    "        if self.cfg.debug: print(\"Tokens:\", tokens.shape)\n",
    "        pos_embed = self.W_pos[:tokens.size(1), :] # [position, d_model]\n",
    "        pos_embed = einops.repeat(pos_embed, \"position d_model -> batch position d_model\", batch=tokens.size(0))\n",
    "        if self.cfg.debug: print(\"pos_embed:\", pos_embed.shape)\n",
    "        return pos_embed\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.W_Q = nn.Parameter(torch.empty((cfg.n_heads, cfg.d_model, cfg.d_head)))\n",
    "        nn.init.normal_(self.W_Q, std=self.cfg.init_range)\n",
    "        self.b_Q = nn.Parameter(torch.zeros((cfg.n_heads, cfg.d_head)))\n",
    "        self.W_K = nn.Parameter(torch.empty((cfg.n_heads, cfg.d_model, cfg.d_head)))\n",
    "        nn.init.normal_(self.W_K, std=self.cfg.init_range)\n",
    "        self.b_K = nn.Parameter(torch.zeros((cfg.n_heads, cfg.d_head)))\n",
    "        self.W_V = nn.Parameter(torch.empty((cfg.n_heads, cfg.d_model, cfg.d_head)))\n",
    "        nn.init.normal_(self.W_V, std=self.cfg.init_range)\n",
    "        self.b_V = nn.Parameter(torch.zeros((cfg.n_heads, cfg.d_head)))\n",
    "        \n",
    "        self.W_O = nn.Parameter(torch.empty((cfg.n_heads, cfg.d_head, cfg.d_model)))\n",
    "        nn.init.normal_(self.W_O, std=self.cfg.init_range)\n",
    "        self.b_O = nn.Parameter(torch.zeros((cfg.d_model)))\n",
    "        \n",
    "        self.register_buffer(\"IGNORE\", torch.tensor(-1e5, dtype=torch.float32, device=device))\n",
    "    \n",
    "    def forward(self, normalized_resid_pre):\n",
    "        # normalized_resid_pre: [batch, position, d_model]\n",
    "        if self.cfg.debug: print(\"Normalized_resid_pre:\", normalized_resid_pre.shape)\n",
    "        \n",
    "        q = einsum(\"batch query_pos d_model, n_heads d_model d_head -> batch query_pos n_heads d_head\", normalized_resid_pre, self.W_Q) + self.b_Q\n",
    "        k = einsum(\"batch key_pos d_model, n_heads d_model d_head -> batch key_pos n_heads d_head\", normalized_resid_pre, self.W_K) + self.b_K\n",
    "        \n",
    "        attn_scores = einsum(\"batch query_pos n_heads d_head, batch key_pos n_heads d_head -> batch n_heads query_pos key_pos\", q, k)\n",
    "        attn_scores = attn_scores / math.sqrt(self.cfg.d_head)\n",
    "        attn_scores = self.apply_causal_mask(attn_scores)\n",
    "\n",
    "        pattern = attn_scores.softmax(dim=-1) # [batch, n_head, query_pos, key_pos]\n",
    "\n",
    "        v = einsum(\"batch key_pos d_model, n_heads d_model d_head -> batch key_pos n_heads d_head\", normalized_resid_pre, self.W_V) + self.b_V\n",
    "\n",
    "        z = einsum(\"batch n_heads query_pos key_pos, batch key_pos n_heads d_head -> batch query_pos n_heads d_head\", pattern, v)\n",
    "\n",
    "        attn_out = einsum(\"batch query_pos n_heads d_head, n_heads d_head d_model -> batch query_pos d_model\", z, self.W_O) + self.b_O\n",
    "        return attn_out\n",
    "\n",
    "    def apply_causal_mask(self, attn_scores):\n",
    "        # attn_scores: [batch, n_heads, query_pos, key_pos]\n",
    "        mask = torch.triu(torch.ones(attn_scores.size(-2), attn_scores.size(-1), device=attn_scores.device), diagonal=1).bool()\n",
    "        attn_scores.masked_fill_(mask, self.IGNORE)\n",
    "        return attn_scores\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.W_in = nn.Parameter(torch.empty((cfg.d_model, cfg.d_mlp)))\n",
    "        nn.init.normal_(self.W_in, std=self.cfg.init_range)\n",
    "        self.b_in = nn.Parameter(torch.zeros((cfg.d_mlp)))\n",
    "        self.W_out = nn.Parameter(torch.empty((cfg.d_mlp, cfg.d_model)))\n",
    "        nn.init.normal_(self.W_out, std=self.cfg.init_range)\n",
    "        self.b_out = nn.Parameter(torch.zeros((cfg.d_model)))\n",
    "    \n",
    "    def forward(self, normalized_resid_mid):\n",
    "        # normalized_resid_mid: [batch, position, d_model]\n",
    "        if self.cfg.debug: print(\"Normalized_resid_mid:\", normalized_resid_mid.shape)\n",
    "        pre = einsum(\"batch position d_model, d_model d_mlp -> batch position d_mlp\", normalized_resid_mid, self.W_in) + self.b_in\n",
    "        post = gelu_new(pre)\n",
    "        mlp_out = einsum(\"batch position d_mlp, d_mlp d_model -> batch position d_model\", post, self.W_out) + self.b_out\n",
    "        return mlp_out\n",
    "    \n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "\n",
    "        self.ln1 = LayerNorm(cfg)\n",
    "        self.attn = Attention(cfg)\n",
    "        self.ln2 = LayerNorm(cfg)\n",
    "        self.mlp = MLP(cfg)\n",
    "    \n",
    "    def forward(self, resid_pre):\n",
    "        # resid_pre [batch, position, d_model]\n",
    "        normalized_resid_pre = self.ln1(resid_pre)\n",
    "        attn_out = self.attn(normalized_resid_pre)\n",
    "        resid_mid = resid_pre + attn_out\n",
    "        \n",
    "        normalized_resid_mid = self.ln2(resid_mid)\n",
    "        mlp_out = self.mlp(normalized_resid_mid)\n",
    "        resid_post = resid_mid + mlp_out\n",
    "        return resid_post\n",
    "\n",
    "class Unembed(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.W_U = nn.Parameter(torch.empty((cfg.d_model, cfg.d_vocab)))\n",
    "        nn.init.normal_(self.W_U, std=self.cfg.init_range)\n",
    "        self.b_U = nn.Parameter(torch.zeros((cfg.d_vocab), requires_grad=False))\n",
    "    \n",
    "    def forward(self, normalized_resid_final):\n",
    "        # normalized_resid_final [batch, position, d_model]\n",
    "        if self.cfg.debug: print(\"Normalized_resid_final:\", normalized_resid_final.shape)\n",
    "        logits = einsum(\"batch position d_model, d_model d_vocab -> batch position d_vocab\", normalized_resid_final, self.W_U) + self.b_U\n",
    "        return logits\n",
    "\n",
    "\n",
    "class DemoTransformer(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.cfg = cfg\n",
    "        self.embed = Embed(cfg)\n",
    "        self.pos_embed = PosEmbed(cfg)\n",
    "        self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])\n",
    "        self.ln_final = LayerNorm(cfg)\n",
    "        self.unembed = Unembed(cfg)\n",
    "    \n",
    "    def forward(self, tokens):\n",
    "        # tokens [batch, position]\n",
    "        embed = self.embed(tokens)\n",
    "        pos_embed = self.pos_embed(tokens)\n",
    "        residual = embed + pos_embed\n",
    "        for block in self.blocks:\n",
    "            residual = block(residual)\n",
    "        normalized_resid_final = self.ln_final(residual)\n",
    "        logits = self.unembed(normalized_resid_final)\n",
    "        # logits have shape [batch, position, logits]\n",
    "        return logits\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "demo_gpt2 = DemoTransformer(Config(debug=False))\n",
    "demo_gpt2.load_state_dict(reference_gpt2.state_dict(), strict=False)\n",
    "demo_gpt2.to(device)\n",
    "\n",
    "len(demo_gpt2.blocks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A global view on layers\n",
    "As a very first exploration, I want to do a simple check: what happens if we remove, duplicate or mix up the intermediate layers?\n",
    "Due to the residual stream, early layers roughly get a signal to act as a full model that is directly being read out if the effect of the later layers is small. Therefore, the model might still work if we remove some of the last layers.\n",
    "Similarly, some of the computations of late layers might mostly depend only or mostly on a subset of the earlier layers, and might still perform useful computations of some other earlier layers are corrupted.\n",
    "\n",
    "## Corruption types to inspect\n",
    "- skip layers at the end\n",
    "- skip layers at the start\n",
    "- skip layers in the middle\n",
    "- repeat layers\n",
    "- permute order of layers\n",
    "\n",
    "## Evaluation\n",
    "We can test this in a number of ways:\n",
    "- we can speak to a corrupted model\n",
    "- we can check the loss of the corrupted model (potentially task-wise)\n",
    "- we can inspect what happens at indivual tokens\n",
    "\n",
    "## Summary wip\n",
    "- first layer seems to be especially important, last layer a bit more important, layers 7-10 ~least important for grammar\n",
    "- "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|endoftext|>The distance between the Colosseum and the Eiffel is approximately 1,000 miles. The distance between the Eiffel and the Colosseum is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosseum and the Eiffel is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosseum and the Eiffel is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosseum and the Eiffel is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosse\n",
      "<|endoftext|>The distance between the Colosseum and the Eiffel is approximately 1,000 miles. The distance between the Eiffel and the Colosseum is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosseum and the Eiffel is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosseum and the Eiffel is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosseum and the Eiffel is approximately 1,000 miles.\n",
      "\n",
      "The distance between the Colosse\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def permuted(self, permutation):\n",
    "    def model(tokens):\n",
    "        # tokens [batch, position]\n",
    "        embed = self.embed(tokens)\n",
    "        pos_embed = self.pos_embed(tokens)\n",
    "        residual = embed + pos_embed\n",
    "        for block_id in permutation:\n",
    "            block = self.blocks[block_id]\n",
    "            residual = block(residual)\n",
    "        normalized_resid_final = self.ln_final(residual)\n",
    "        logits = self.unembed(normalized_resid_final)\n",
    "        # logits have shape [batch, position, logits]\n",
    "        return logits\n",
    "    return model\n",
    "\n",
    "\n",
    "def complete(model, reference_text, max_tokens=100, T=1e-3):\n",
    "    tokens = reference_gpt2.to_tokens(reference_text)\n",
    "    for i in range(max_tokens):\n",
    "        tokens = tokens.to(device)\n",
    "        logits = model(tokens)\n",
    "        \n",
    "        # Apply temperature scaling\n",
    "        scaled_logits = logits / T\n",
    "        \n",
    "        # Convert logits to probabilities\n",
    "        probs = torch.nn.functional.softmax(scaled_logits, dim=-1)\n",
    "        \n",
    "        # Sample from the probability distribution\n",
    "        next_token = torch.multinomial(probs[0, -1], num_samples=1)\n",
    "        \n",
    "        # Concatenate the new token to the existing sequence\n",
    "        tokens = torch.cat([tokens, next_token.unsqueeze(0)], dim=-1)\n",
    "    \n",
    "    # Decode the tokens to text\n",
    "    return reference_gpt2.tokenizer.decode(tokens[0]), cache\n",
    "\n",
    "permuted_model = permuted(demo_gpt2, list(range(12)))\n",
    "text, _ = complete(demo_gpt2, reference_text)\n",
    "print(text)\n",
    "text, _ = complete(permuted_model, reference_text)\n",
    "print(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic questions\n",
    "- what happens if we ablate the last n layers?\n",
    "  - [ ] systematic evaluation\n",
    "- what happens if we ablate the first n layers?\n",
    "\n",
    "# Evaluation\n",
    "- if we ablate something we want to save this info for later\n",
    "  - (model, prompt, deleted compute nodes, logit diffs)\n",
    "\n",
    "\n",
    "# Synthetic Dataset\n",
    "- we want to quickly form hypothesis of the form: \"(circuit) is important for (task)\"\n",
    "- we start with circuit = (layer), transcribe the logit diffs\n",
    "- we let gpt-4 guess:\n",
    "  - what is a task in which this circuit is relevant\n",
    "  - what other layers might it interact strongly with?\n",
    "- generate a task\n",
    "- \n",
    "\n",
    "# Inspecting layers\n",
    "- which token logits differ the most if we ablate layer i?\n",
    "- for each layer, we can collect a dataset:\n",
    "  - (prompt, logit_diffs: token -> float)\n",
    "\n",
    "\n",
    "# Random thoughts\n",
    "- we potentially don't want to look at single layers / attention heads, but ablate arbitrary parts (spanning multiple layers / only parts of layers)\n",
    "-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'model': 'test-model', 'prompt': 'hello', 'completion': 'yo yo yo'}, {'model': 'test-model', 'prompt': 'hello', 'completion': 'yo yo yo'}, {'model': 'test-model', 'prompt': 'hello', 'completion': 'yo yo yo'}, {'model': 'test-model', 'prompt': 'hello', 'completion': 'yo yo yo'}]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from dumbdb import ExperimentData\n",
    "\n",
    "\n",
    "store = ExperimentData('/Users/nielswarncke/Documents/code/TransformerLens/experiments')\n",
    "\n",
    "# store({\n",
    "#     'model': 'test-model',\n",
    "#     'prompt': 'hello world',\n",
    "#     'completion': 'yo yo yo'\n",
    "# })\n",
    "\n",
    "# store({\n",
    "#     'model': 'test-model',\n",
    "#     'prompt': 'hello',\n",
    "#     'completion': 'yo yo yo'\n",
    "# })\n",
    "\n",
    "i = store.filter(model='test-model', prompt='hello')\n",
    "print(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Removing the last n layers\n",
    "\n",
    "Observations:\n",
    "\n",
    "Prompt: The distance between Rome and Paris is approximately\"\n",
    "- layer 0 only repeats the last token\n",
    "- layer 3 is the first to bring up a number-ish completion (20th century)\n",
    "- layers 4 is the first to bring up 'distance' related words\n",
    "- layer 7 is the first to bring up an actual distance (200 miles)\n",
    "- layer 8 is the first to produce long grammatically correct phrases / sentences\n",
    "\n",
    "Prompt:\n",
    "```\n",
    "1 + 3 = 4 \n",
    "4 + 5 = 9\n",
    "2 + 3 =\n",
    "```\n",
    "- layer 0 only repeats the last token\n",
    "- layer 6 notices this is a math equation\n",
    "- layer 8 gets the single addition prompt syntax\n",
    "- layer 10 shows signs of correct addition before repeating = 9 for any formular\n",
    "- layer 11 predicts a trend of going to double digit addition, with wrong answers\n",
    "- layer 12 makes correct addition examples\n",
    "- it's a bit surprising that it gets correct so late, this can mean one of the following:\n",
    "  - addition always happens late in this model, it can't 'add numbers in its head before speaking'\n",
    "    - we could test this using tasks constructed to involve this (3 + 1 + 5 = )\n",
    "    - the = can attend to the +1 and the +5 which might hold precomputed intermediate values\n",
    "    - we can test where intermediate values are computed using causal tracing but lets keep this for another experiment\n",
    "  - the residual stream semantics of the middle layers is not 'approximates the final layer', therefore it is ~meaningless to simply unembed their output\n",
    "    - we could train a 'best linear guess' transformation for skipped layers, that would show what the model knows at any point assuming it 'knows' linear features\n",
    "\n",
    "\n",
    "Prompt: unfinished JSON\n",
    "- layer 0 only repeats the last token\n",
    "- layer 2 seems to understand this is code related\n",
    "- layer 9 tries to close the JSON object\n",
    "- layer 10 correctly closes the JSON object\n",
    "\n",
    "\n",
    "Prompt: Obama was elected in the year\n",
    "- layer 1 only repeats the last token\n",
    "- most likely next token stays 'yearlong' until layer 10\n",
    "  - it is unlikely that no useful computation is happening all that time, which supports the scepticism that this is a useful approach (semantics might change)\n",
    "- layer 10 speaks grammatically correct\n",
    "\n",
    "All together:\n",
    "- keeping only layer 0 only repeats the last token\n",
    "- grammar/format gets correct between layer 8-10\n",
    "- addition gets correct in the last layer\n",
    "- for the other prompts, layer 10 outputs are quite similar to layer 12 outputs\n",
    "- the intermediate layers might 'know' (linearly encode) more than we can read with this method, but the fact that a lot gets correct already in layer 8 and that this happens in different layers for different prompts suggests that the residual stream semantics do not drift so much\n",
    "\n",
    "  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0: <|endoftext|>The distance between Rome and Paris is approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately\n",
      "Layer 1: <|endoftext|>The distance between Rome and Paris is approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately\n",
      "Layer 2: <|endoftext|>The distance between Rome and Paris is approximately approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate approximate\n",
      "Layer 3: <|endoftext|>The distance between Rome and Paris is approximately 20th century era era era after the same exact exact same exact same exact same exact same exact same exact same exact same exact same exact same exact same exact same exact same exact same exact same exact exact same exact exact same exact exact same exact exact same exact exact same exact exact same exact exact same exact exact exact same exact exact exact same exact exact exact same exact exact same exact exact same exact exact exact amount amount amount amount amount amount amount amount amount amount amount amount amount amount amount amount amount amount amount amount\n",
      "Layer 4: <|endoftext|>The distance between Rome and Paris is approximately approximate approximate approximate approximate approximate approximate location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location location\n",
      "Layer 5: <|endoftext|>The distance between Rome and Paris is approximately approximate approximate approximate approximate approximate results results results results results from the same exact same exactaneouslyaneouslyaneouslyaneouslyaneously (sic) in order order orderlies and/advancements of the same thing happened to the same amount amount amount amountwisewisewisewisely.\n",
      "\n",
      "If you're lucky lucky lucky lucky lucky charm charm charm charm charm charm charm charm charm charm charm charm charm.\n",
      "\n",
      "If you're lucky lucky lucky lucky lucky charmcraftsmanshipsmanshipsmanshipsmanshipsmanshipsmanshipsmanshipsmanshipsmanshipsmanshipsmanshipsmanship\n",
      "Layer 6: <|endoftext|>The distance between Rome and Paris is approximately equivalent between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles between two dozen thousand miles and then proceeded through the same amount amount amount amount amountedly amounted by virtue virtue virtue virtueously provided by virtue virtuefulnessfulnessfulnessy-style-\n",
      "Layer 7: <|endoftext|>The distance between Rome and Paris is approximately equivalent between two thirds of the countrywide population residing within the same countrywide population residing within the same countrywide addressable by virtue virtue of the latter's own personalities and desires. Therefore, accordinglyforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforthforth\n",
      "Layer 8: <|endoftext|>The distance between Rome and Paris is approximately 200 kilometers between each major axis of the world. However, thereabouts are still plenty enough to satisfy the most important part of the world's economy. However, thereabouts are also sometimes referred by nameplate namespaces. These include the United States, which consists mainly consistantly composed solely consistantly composed solely consistative. However, thereabouts are also known aliases aliases. These include the United States, which consists mainly consistently composed solely consistative. However, thereabouts are also known aliases aliases.\n",
      "Layer 9: <|endoftext|>The distance between Rome and Paris is approximately equal between two separate sets of stairs. However, there is no indication whatsoever that these two halves of the globe are separated by any significant amount of space. Instead, it seems likely that the entire globe descended into the abyss beneath the surface of earth. Thus far, however, there appears no indication whatsoever that any significant amount of materialized debris debris exists anywhere else besides the aforementioned location. Furthermore, although there may be some semblance of a solidified metal artifact resembling a human skull, there nevertheless appears no\n",
      "Layer 10: <|endoftext|>The distance between Rome and Paris is approximately 1½ kilometers between the two countries' borders. However, there are several ways to navigate this distance. Firstly, you can easily travel through the city itself, which consists mainly of buildings and buildings belonging to various groups of people who reside within the city itself. Secondly, you can travel through the city itself, which consists mainly of buildings belonging to various groups of people who reside within the city itself. Lastly, you can travel through the city itself, which consists mainly of buildings belonging to various groups of\n",
      "Layer 11: <|endoftext|>The distance between Rome and Paris is approximately 1½ miles. The distance between Rome and Paris is approximately 1½ miles. Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "Photo: Getty Images\n",
      "\n",
      "\n",
      "Layer 12: <|endoftext|>The distance between Rome and Paris is approximately 1,000 miles. The distance between Rome and Paris is approximately 1,000 miles.\n",
      "\n",
      "The distance between Rome and Paris is approximately 1,000 miles.\n",
      "\n",
      "The distance between Rome and Paris is approximately 1,000 miles.\n",
      "\n",
      "The distance between Rome and Paris is approximately 1,000 miles.\n",
      "\n",
      "The distance between Rome and Paris is approximately 1,000 miles.\n",
      "\n",
      "The distance between Rome and Paris is approximately 1,000 miles.\n",
      "\n",
      "The distance between Rome and\n",
      "Layer 0: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =\n",
      "Layer 1: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================\n",
      "Layer 2: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 ================= -_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-\n",
      "Layer 3: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =============== -_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-_-\n",
      "Layer 4: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 ============>=>=>=>=>=>=>...........................................................................................\n",
      "Layer 5: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 ============|__dict_dict_dictender_dictender_dictender_dictender_type_type_type_type_type_type_type_type_type_type_type_type_type_type_type_type_type.\n",
      "                                          (*)$112080802111234567890808080211.html\n",
      "\n",
      "(*) ERROR REPORT REPORTINGLYLYSTEMSTEMSTEMOOLOOLOOLOOLOOLOOLINGINGING\n",
      "Layer 6: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 x x x x x x x............................................................................................\n",
      "Layer 7: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4.5%\n",
      "\n",
      "Advertisements<|endoftext|>Description Description Description Description Description Description Description Description Description\n",
      "Trivia Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit ]\n",
      "Trivia Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit ]\n",
      "Trivia Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit ]\n",
      "Trivia Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit Edit ]\n",
      "References Edit\n",
      "Layer 8: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "\n",
      "Layer 9: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "\n",
      "Layer 10: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 5 \n",
      "4 + 5 = 9 \n",
      "5 + 5 = 9 \n",
      "6 + 5 = 9 \n",
      "7 + 5 = 9 \n",
      "8 + 5 = 9 \n",
      "9 + 5 = 9 \n",
      "10 + 5 = 9 \n",
      "11 + 5 = 9 \n",
      "12 + 5 = 9 \n",
      "13 + 5 = 9 \n",
      "14 + 5 = 9 \n",
      "15 + 5 = 9 \n",
      "16 + 5 = 9 \n",
      "17 + 5 = 9 \n",
      "Layer 11: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 5 \n",
      "6 + 6 = 10 \n",
      "7 + 7 = 11 \n",
      "8 + 8 = 12 \n",
      "9 + 9 = 13 \n",
      "10 + 10 = 14 \n",
      "11 + 11 = 15 \n",
      "12 + 12 = 16 \n",
      "13 + 13 = 17 \n",
      "14 + 14 = 18 \n",
      "15 + 15 = 19 \n",
      "16 + 16 = 20 \n",
      "17 + 17 = 21 \n",
      "18 + 18 = 22 \n",
      "19 + 19 = 23 \n",
      "Layer 12: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = 5 \n",
      "6 + 4 = 10\n",
      "3 + 3 = 6 \n",
      "7 + 4 = 11\n",
      "8 + 4 = 12\n",
      "9 + 4 = 13\n",
      "10 + 4 = 14\n",
      "11 + 4 = 15\n",
      "12 + 4 = 16\n",
      "13 + 4 = 17\n",
      "14 + 4 = 18\n",
      "15 + 4 = 19\n",
      "16 + 4 = 20\n",
      "17 + 4 = 21\n",
      "18 + 4 = 22\n",
      "19 + 4 = 23\n",
      "20 + 4 = 24\n",
      "\n",
      "Layer 0: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \" corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid\n",
      "Layer 1: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"nonozzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzlezzle\n",
      "Layer 2: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"true True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True\n",
      "Layer 3: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"true True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True\n",
      "Layer 4: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"true True True True......... ] ]\n",
      "\n",
      "(*) Updated update update notification notification notification notification notification notification notification notification notification notification.\n",
      "\n",
      "If you're not really bothered bother bother bothersomenessyieldings of the same thing happened hereticsticsticstics...............................................\n",
      "Layer 5: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"]\"], respectively.\n",
      "                                                          (*)$112080802111234567890808080211.html\n",
      "\n",
      "(*) Addeditionally supported support support for the same exact same thing happened to the same amount amount amount amount. Howevertonsteryteryteryteryteryteryteryteryteryteryteryteryteryteryteryteryteryteryterytery.completionist_classifications (...)................\n",
      "Layer 6: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"]\"], }elsewhereforeforeforementioned:\n",
      "                                          (*)$ $ $...............................................................................\n",
      "Layer 7: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \" etceterranormalities) ;;;;;;;;;;;;;;;*/[?]]\n",
      "                -|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|\n",
      "Layer 8: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"web browsers etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc., etc.,\n",
      "Layer 9: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"android]}, {\"name\": \"Alanjamin, Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr., Jr.,\n",
      "Layer 10: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"git\"]}, {\"name\": \"Alan\", \"age\": 27, \"skills\": [\"Microsoft Office\", \"git\"]}, {\"name\": \"Alan\", \"age\": 27, \"skills\": [\"Microsoft Office\", \"git\"]}, {\"name\": \"Alan\", \"age\": 27, \"skills\": [\"Microsoft Office\", \"git\"]}, {\"name\": \"Alan\", \"age\": 27, \"skills\": [\"Microsoft Office\", \"git\"]}, {\"name\": \"Alan\n",
      "Layer 11: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"python\"]}, {\"name\": \"Alan\", \"age\": 27, \"skills\": [\"Microsoft Office\", \"python\"]}, {\"name\": \"Alan\", \"age\": 26, \"skills\": [\"Microsoft Office\", \"python\"]}, {\"name\": \"Alan\", \"age\": 25, \"skills\": [\"Microsoft Office\", \"python\"]}, {\"name\": \"Alan\", \"age\": 24, \"skills\": [\"Microsoft Office\", \"python\"]}, {\"name\": \"Alan\n",
      "Layer 12: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"python\"]}, {\"name\": \"Alan\", \"age\": 27, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 26, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 25, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 24, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\":\n",
      "Layer 0: <|endoftext|>Obama was elected in the year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year\n",
      "Layer 1: <|endoftext|>Obama was elected in the year ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago ago\n",
      "Layer 2: <|endoftext|>Obama was elected in the yearlonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglonglong Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long Long\n",
      "Layer 3: <|endoftext|>Obama was elected in the yearlong ago ago ago when the same thing happened happening happening rightwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardwardward\n",
      "Layer 4: <|endoftext|>Obama was elected in the yearlong process process processionalized the same exact exact same exact same exact result.\n",
      "\n",
      "Theresa Theresa Theresa's own private sector sector sector sector sector sector sector sectorwidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewidewide\n",
      "Layer 5: <|endoftext|>Obama was elected in the yearlong celebration celebrating the same exact same thing happened happening around the same day when they're supposed to be able accessiblyibly disabled disabled disabled disabled disabled disabled disabled disabled.\n",
      "Howeverler's latest roundtrip tripship Yamato Takeru Takeru Takeru Takeru's own personalised version version version version version of the same thing happened happening apart apart apart. Itselfortiumortiumortiumortiumortiumortiumortium.\n",
      "       Another major difference difference between the same thing happening happening happening at least partially overlapping between the same amount amount amount\n",
      "Layer 6: <|endoftext|>Obama was elected in the yearlong campaign campaign campaign manager Mike Flynn's own personalised version of the same thing happenings happenings in the same vein vein. Buttonsholeholeholeholeholeholefuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckfuckinstrumentation.\n",
      "                                        ••½½½½½½½½½½½½½½½½½½½½½½½½½½½½½½½½\n",
      "Layer 7: <|endoftext|>Obama was elected in the yearlong campaign campaign against the same kind of thing happening elsewhere elseworldly, albeit perhaps not necessarily necessarily deterministic, but nevertheless nonetheless nonetheless nonetheless nonetheless let us hope someday someday someday someday…\n",
      "            ••★️️‍♪️‍♪️‍♪️‍♪️‍♪️‍♪️‍♪️‍♪️‍♪️‍♪️‍♪️\n",
      "Layer 8: <|endoftext|>Obama was elected in the yearlong process of reformulation of the nation's economy, including the creation of new jobs jobs jobs jobs jobs jobs jobs jobs....\"\n",
      "\n",
      "Theresa Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel Merkel\n",
      "Layer 9: <|endoftext|>Obama was elected in the yearlong campaign campaign campaign against Hillary Clinton Hillary Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton Rodham Clinton\n",
      "Layer 10: <|endoftext|>Obama was elected in the year 2000, but he still had plenty of time to prepare for his presidency. During his tenure, he oversaw the creation of America's largest military force ever assembled, including military units stationed abroad overseas. During his tenure, he oversaw the creation of America's largest military force ever assembled, including military units stationed abroad overseas. During his tenure, he oversaw the creation of America's largest military force ever assembled, including military units stationed abroad overseas. During his tenure, he oversaw the creation of America's largest military\n",
      "Layer 11: <|endoftext|>Obama was elected in the year 2000, but he was still a candidate for president after his predecessor Barack Obama became president. Obama was elected in the year 2000, but he was still a candidate for president after his predecessor Barack Obama became president. Photo: Getty Images\n",
      "\n",
      "window._taboola = window._taboola || []; _taboola.push({flush: true});\n",
      "\n",
      "President Barack Obama greets supporters after his inauguration ceremony at the White House in Washington, DC, January 21, 2009. President Barack Obama greets supporters after his inauguration ceremony\n",
      "Layer 12: <|endoftext|>Obama was elected in the year 2000, and he was elected in 2008. He was elected in 2012, and he was elected in 2014. He was elected in 2016, and he was elected in 2017. He was elected in 2018, and he was elected in 2019. He was elected in 2020, and he was elected in 2021. He was elected in 2022, and he was elected in 2023. He was elected in 2024, and he was elected in 2025. He was elected in 2026, and he was\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "reference_text = \"The distance between the Colosseum and the Eiffel is approximately\"\n",
    "texts = [\n",
    "    \"The distance between Rome and Paris is approximately\",\n",
    "    \"1 + 3 = 4 \\n4 + 5 = 9\\n2 + 3 =\",\n",
    "    json.dumps([{'name': 'James', 'age': 34, 'skills': ['Python', 'git']}, {'name': 'Alan', 'age': 28, 'skills': ['MS Office', '<cut>']}]).split('<cut>')[0],\n",
    "    \"Obama was elected in the year\",\n",
    "]\n",
    "for prompt in texts:\n",
    "    for i in range(0, 13):\n",
    "        keep_layer_ids =  list(range(i))\n",
    "        permuted_model = permuted(demo_gpt2, keep_layer_ids)\n",
    "        text, _ = complete(permuted_model, prompt)\n",
    "        print(f\"Layer {i}: {text}\")\n",
    "        store({\n",
    "            'model': 'gpt-2-small',\n",
    "            'prompt': prompt,\n",
    "            'completion': text,\n",
    "            'ablation': {\n",
    "                'layers': keep_layer_ids\n",
    "            }\n",
    "        })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "# store.to_disk()\n",
    "outputs = store.filter(ablation={'layers': lambda layers: 0 not in layers and 1 in layers})\n",
    "for i in outputs:\n",
    "    print(json.dumps(i, indent=2))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Removing the first n layers\n",
    "Observations:\n",
    "- if we remove the first layer it breaks the model completely"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 1: <|endoftext|>The distance between Rome and Paris is approximately...................................................................................... pass hell. hell. hell hell hell hell hell hell hell hell hell\n",
      "Layer 2: <|endoftext|>The distance between Rome and Paris is approximately,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,, the,,, the,,, the,, the, the, the, the, the, the, the, the, the, the the the, the the, the the\n",
      "Layer 3: <|endoftext|>The distance between Rome and Paris is approximately-........................................------------)-)-).).).).).)))))))))))))))))))))))))))))))))\n",
      "Layer 4: <|endoftext|>The distance between Rome and Paris is approximately-------------------------------,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Layer 5: <|endoftext|>The distance between Rome and Paris is approximatelyAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA\n",
      "Layer 6: <|endoftext|>The distance between Rome and Paris is approximately------------——————————————————————————————————————————————————————————————--------------------------\n",
      "Layer 7: <|endoftext|>The distance between Rome and Paris is approximatelytotototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototo\n",
      "Layer 8: <|endoftext|>The distance between Rome and Paris is approximately and,,,,,,,,,,,,,,,,/ to'''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' '\n",
      "Layer 9: <|endoftext|>The distance between Rome and Paris is approximately............,...,...,...,...,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 10: <|endoftext|>The distance between Rome and Paris is approximately,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 11: <|endoftext|>The distance between Rome and Paris is approximately---------------------,,,,,,,,,,,,,,,,,,,s,s,,s,,,,,,,,,,,,,,,,,,,,,,------s,-s,-s,s,-s,s,s,s,s,s,s,s\n",
      "Layer 12: <|endoftext|>The distance between Rome and Paris is approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately approximately\n",
      "Layer 1: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =....................................................................................................\n",
      "Layer 2: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,, the,,, the,,, the, the, the the the, the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the\n",
      "Layer 3: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =))))))))-)-)-)-)-)- )- )- ) ) ) ) ) ) ) ) ) ) ) ) ) ) \") \") \") \") \") \") \") \") \") \") \") \" All ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (\n",
      "Layer 4: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      " \" \"----------------\n",
      "Layer 5: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =.---------------------------------------------------------------------------------------------------\n",
      "Layer 6: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =—————————————————————————————————————————————————————————————————————————————————————-—-—-—-—-—-—-—-\n",
      "Layer 7: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =...................../////////tototototototototototo//////////tototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototo\n",
      "Layer 8: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = ( ( ',''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''\n",
      "Layer 9: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =...,...,...,...,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 10: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 11: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 =,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 12: <|endoftext|>1 + 3 = 4 \n",
      "4 + 5 = 9\n",
      "2 + 3 = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =\n",
      "Layer 1: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"....................................................................................................\n",
      "Layer 2: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,, the,,,,,, the,,,,,, the,,,,,, the,,,,,,, the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the\n",
      "Layer 3: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"------)-)-)-)-)\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")\n",
      ")))\n",
      ")))))))\n",
      "))))))))))))))\n",
      ",))))))))))\n",
      "Layer 4: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \",,,,, \", \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \"- \"- \"- \"- \"- \"- \"- \"- \"- \"- \"- \"- \"- \"- \"----------- (- (- (- (- (- (- (- (- (- (- (-----------------\n",
      "Layer 5: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"----------------------------------------------------------------------------------------------------\n",
      "Layer 6: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "Layer 7: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\n",
      "Layer 8: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"'''''''''''and ', and ',,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 9: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \",\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\" \"\" \"\" \"\" \"\" \"\" \"\" \"\" \"\" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \" \"\n",
      "Layer 10: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\n",
      "Layer 11: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \"\n",
      "ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss't,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t,'t\n",
      "Layer 12: <|endoftext|>[{\"name\": \"James\", \"age\": 34, \"skills\": [\"Python\", \"git\"]}, {\"name\": \"Alan\", \"age\": 28, \"skills\": [\"MS Office\", \" corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid corrid\n",
      "Layer 1: <|endoftext|>Obama was elected in the year....................................................................................................\n",
      "Layer 2: <|endoftext|>Obama was elected in the year,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,, the,,, the,,, the,, the, the, the, the, the, the, the, the, the, the, the the, the,\n",
      "Layer 3: <|endoftext|>Obama was elected in the year---......................).)--)-)-)-)-)-)-)-)-)-)-)-)-)-)-)----)----------)-----)-------),),),),),),\n",
      "Layer 4: <|endoftext|>Obama was elected in the year------------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      ",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 5: <|endoftext|>Obama was elected in the year----------------------------------------------------------------------------------------------------\n",
      "Layer 6: <|endoftext|>Obama was elected in the year------------——————————————————————————————————————————————————————————————--------------------------\n",
      "Layer 7: <|endoftext|>Obama was elected in the year------totototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototototo\n",
      "Layer 8: <|endoftext|>Obama was elected in the year,,,,,,,,,,,,,,,,,,,,''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''\n",
      "Layer 9: <|endoftext|>Obama was elected in the year...,...,...,...,...,...,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,----------\n",
      "Layer 10: <|endoftext|>Obama was elected in the year,,,,,,,,,, and and,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Layer 11: <|endoftext|>Obama was elected in the yearsssssssssssssss-------------------------------------------------------------------------------------\n",
      "Layer 12: <|endoftext|>Obama was elected in the year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year year\n"
     ]
    }
   ],
   "source": [
    "for prompt in texts:\n",
    "    for i in range(1, 13):\n",
    "        keep_layer_ids =  list(range(i, len(demo_gpt2.blocks)))\n",
    "        permuted_model = permuted(demo_gpt2, keep_layer_ids)\n",
    "        text, _ = complete(permuted_model, prompt)\n",
    "        print(f\"Layer {i}: {text}\")\n",
    "        store({\n",
    "            'model': 'gpt-2-small',\n",
    "            'prompt': prompt,\n",
    "            'completion': text,\n",
    "            'ablation': {\n",
    "                'layers': keep_layer_ids\n",
    "            }\n",
    "        })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run a layer many times\n",
    "Observations:\n",
    "- Duplicating the first layer completely breaks the model\n",
    "- duplicating later layers breaks the model less\n",
    "- duplicating the last layer breaks the model again more\n",
    "- I should quantify this by perplexity on different tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Layer 0x2\n",
      ": <|endoftext|>Obama was elected in the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the year of the the the the the the the the the the the the the the the the the the\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 1x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he was elected in the year 2000.\n",
      "\n",
      "The first president to win a presidential election was George W. Bush in 2000.\n",
      "\n",
      "The first president to win a presidential election was George H.W. Bush in 2000.\n",
      "\n",
      "The first president to win a presidential election was George W. Bush in 2000.\n",
      "\n",
      "The first president to win a presidential election was George H.W. Bush in 2000.\n",
      "\n",
      "The first president to win a presidential election was George H\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 2x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he was elected in 2008.\n",
      "\n",
      "Obama's first term was a disaster for the country.\n",
      "\n",
      "The economy was in a tailspin, and the economy was in a tailspin.\n",
      "\n",
      "The economy was in a tailspin.\n",
      "\n",
      "The economy was in a tailspin.\n",
      "\n",
      "The economy was in a tailspin.\n",
      "\n",
      "The economy was in a tailspin.\n",
      "\n",
      "The economy was in a tailspin.\n",
      "\n",
      "The economy was in a tailspin\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 3x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he was elected in 2008. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He was elected in 2012. He\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 4x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he was elected in the year 2000.\n",
      "\n",
      "The first time Obama was elected president was in 2000, when he was elected in the year 2000.\n",
      "\n",
      "Obama was elected in the year 2000, and he was elected in the year 2000.\n",
      "\n",
      "Obama was elected in the year 2000, and he was elected in the year 2000.\n",
      "\n",
      "Obama was elected in the year 2000, and he was elected in the year 2000.\n",
      "\n",
      "Obama was elected in the year 2000,\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 5x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he has been in office since. He has been in office since the beginning of the year 2000, and he has been in office since the beginning of the year 2000.\n",
      "\n",
      "Obama has been in office since the beginning of the year 2000, and he has been in office since the beginning of the year 2000.\n",
      "\n",
      "Obama has been in office since the beginning of the year 2000, and he has been in office since the beginning of the year 2000.\n",
      "\n",
      "Obama has been\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 6x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he has been in office since. He has been a leader of the free world, and he has been a leader of the free world's people. He has been a leader of the free world's people. He has been a leader of the free world's people. He has been a leader of the free world's people. He has been a leader of the free world's people. He has been a leader of the free world's people. He has been a leader of the\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 7x2\n",
      ": <|endoftext|>Obama was elected in the year of the Great Depression, and he was the first president to be elected in a major country. He was the first president to be elected in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in a major country in\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 8x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he has been a leader in the fight against terrorism since. He has been a leader in the fight against the Islamic State, and he has been a leader in the fight against the Islamic State.\n",
      "\n",
      "He has been a leader in the fight against the Islamic State, and he has been a leader in the fight against the Islamic State.\n",
      "\n",
      "He has been a leader in the fight against the Islamic State, and he has been a leader in the fight against the Islamic State.\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 9x2\n",
      ": <|endoftext|>Obama was elected in the year 2000, and he has been in office since. He has been a major player in the political process for more than two decades. He has been a major player in the political process for more than two decades. He has been a major player in the political process for more than two decades.\n",
      "\n",
      "He has been a major player in the political process for more than two decades. He has been a major player in the political process for more than two decades.\n",
      "\n",
      "He has been a major player\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 10x2\n",
      ": <|endoftext|>Obama was elected in the year of his birth, and he was born in the year of his death.\n",
      "\n",
      "The Obama administration has been trying to get the word out about the birth of his son, but the Obama administration has been trying to get the word out about the birth of his son, but the Obama administration has been trying to get the word out about the birth of his son, but the Obama administration has been trying to get the word out about the birth of his son, but the Obama administration has been trying to\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 11x2\n",
      ": <|endoftext|>Obama was elected in the year of his inauguration, and he has been a leader in the fight against the \"war on drugs.\"\n",
      "\n",
      "The president has been a leader in the fight against the \"war on drugs.\"\n",
      "\n",
      "The president has been a leader in the fight against the \"war on drugs.\"\n",
      "\n",
      "The president has been a leader in the fight against the \"war on drugs.\"\n",
      "\n",
      "The president has been a leader in the fight against the \"war on drugs.\"\n",
      "\n",
      "The president has been a leader\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 0x5\n",
      ": <|endoftext|>Obama was elected in the year, the of the of the- the- the- the-------------------- the--------- the-------------- the-------- the------- the------ the------ the------ the-.\n",
      ".\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 1x5\n",
      ": <|endoftext|>Obama was elected in the year of 2008.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a disaster.\n",
      "\n",
      "The Obama campaign was a\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 2x5\n",
      ": <|endoftext|>Obama was elected in the year 2000.\n",
      "\n",
      "The Obama administration was elected in the year in the Obama in the year in the in the government was responsible for the economy.\n",
      "\n",
      "The economy was a big job program was a big job.\n",
      "And the economy was a big job.\n",
      "And the economy was a big economy.\n",
      "\n",
      "And we're going to be able to get more jobs.\n",
      "\n",
      "\n",
      "But with the economy.\n",
      "But the economy was a big job.\n",
      "\n",
      "But the economy was a big\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 3x5\n",
      " <|endoftext|>Obama was elected in the year-long election campaign that saw the election of the Republican Party. He was elected to the be- g the Republican Party in the House of Representatives in the House of Representatives in the year-long election campaign that saw the election of the Republican Party. The Republican Party was in the House of Representatives in the year-^, of the House of Representatives in the year-of- the HOUSE of the House of Representatives in the year- of the first. The Republican Party was in the House of the\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 4x5\n",
      " <|endoftext|>Obama was elected in the year he was sworn into the oath of office, and he was a big part of the White House. But he was also a very big part of the White House.\n",
      "\n",
      "He was a very big part of the White House.\n",
      "\n",
      "He was a very big part of the White House.\n",
      "\n",
      "Obama was a big part of the White House.\n",
      "\n",
      "Obama was a big part of the White House.\n",
      "\n",
      "Obama was a big part of the White House.\n",
      "\n",
      "Obama was a\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 5x5\n",
      " <|endoftext|>Obama was elected in the year after the election of the American Revolution.\n",
      "\n",
      "Obama was elected in the year after the Obama administration announced that Obama's administration is considering a ban on all federal funding of the federal government.\n",
      "\n",
      "Obama was elected in the year after the Obama administration announced that Obama's administration is considering a \"red line\" for the country's economy and the economy.\n",
      "\n",
      "Obama was elected in the year after the Obama administration announced that Obama's administration is considering a \"red line\" for the country's economy\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 6x5\n",
      " <|endoftext|>Obama was elected in the year of 2012 as the first African-American president in any election in ever since the founding of the United States.\n",
      "\n",
      "President Obama speaks during the launch of the new U.S. Air Force leaves the air cover with the logo on the facade of the facade of the facade of the facade of the facade of the as is with a product as is a product as is a company in the company with the company with the the the employees who is in the company with the members of the employees, as\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 7x5\n",
      " <|endoftext|>Obama was elected in the year of the \"war on drugs\" and has been accused by the anti-immigrant far-right group the same day he was sworn in as president. (Photo: AP) President Barack Obama (Photo: AP) President Barack Obama (Photo: AP) Photo: AP) U.S. President Joe Manam (Photo: AP) U.S. President Joe Manam (Photo: AP) U.S. President Joe Manam (Photo: AP) U.S. President Joe\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 8x5\n",
      " <|endoftext|>Obama was elected in the year of 9/11 with a record of economic and military victory over a Democratic president.\n",
      "\n",
      "But now, the president is facing a new challenge: How much more to tax the wealthy and Wall-------------------------------------------------------------\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 9x5\n",
      " <|endoftext|>Obama was elected in the year 2000 as the country's first African-American president.\n",
      "\n",
      "The first African-American president of the United States was John Wood Johnson in 1804.\n",
      "\n",
      "Johnson was born William H. Johnson May 8, 1804, in New York City.\n",
      "\n",
      "Johnson was born William H. Johnson May 8, 1804, in New York City. William H. Johnson was born May 8, 1804, in New York City on May 8, 1804.\n",
      "\n",
      "Johnson was born May\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 10x5\n",
      " <|endoftext|>Obama was elected in the year after the election of his first black president, but he was not the first black president to be elected president of the United States. In fact, he was the first black president to be elected president of the United States from out of the country than from out of the country of his birth.\n",
      "\n",
      "President Obama was elected president of the United States from out of the country of out of the country of out of the country of out of the country of out of the country of out of the country of\n",
      "--------------------------------------------------------------------------------\n",
      "Layer 11x5\n",
      " <|endoftext|>Obama was elected in the year of his election. He is the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the\n",
      "--------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "for i in [2, 5]:\n",
    "    for layer in range(len(demo_gpt2.blocks)):\n",
    "        keep_layer_ids =  list(range(len(demo_gpt2.blocks)))\n",
    "        keep_layer_ids = keep_layer_ids[:layer] + [layer] * i + keep_layer_ids[layer+1:]\n",
    "        item = {\n",
    "            'model': 'gpt-2-small',\n",
    "            'prompt': prompt,\n",
    "            'ablation': {\n",
    "                'layers': keep_layer_ids\n",
    "            }\n",
    "        }\n",
    "        if len(items := store.filter(**item)) == 0:   \n",
    "            permuted_model = permuted(demo_gpt2, keep_layer_ids)\n",
    "            text, _ = complete(permuted_model, prompt)\n",
    "            print(f\"Layer {layer}x{i}\\n {text}\")\n",
    "            store({\n",
    "                'model': 'gpt-2-small',\n",
    "                'prompt': prompt,\n",
    "                'completion': text,\n",
    "                'ablation': {\n",
    "                    'layers': keep_layer_ids\n",
    "                }\n",
    "            })\n",
    "        else:\n",
    "            item = items[0]\n",
    "            print(f\"Layer {layer}x{i}\\n: {item['completion']}\")\n",
    "            \n",
    "        print(\"-\"*80)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observations:\n",
    "- Duplicating the first layer completely breaks the model\n",
    "- duplicating later layers breaks the model less\n",
    "- duplicating the last layer breaks the model again more\n",
    "- I should quantify this by perplexity on different tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model': 'gpt-2-small',\n",
       " 'prompt': 'Obama was elected in the year',\n",
       " 'completion': \"<|endoftext|>Obama was elected in the year 2000.\\n\\nThe Obama administration was elected in the year in the Obama in the year in the in the government was responsible for the economy.\\n\\nThe economy was a big job program was a big job.\\nAnd the economy was a big job.\\nAnd the economy was a big economy.\\n\\nAnd we're going to be able to get more jobs.\\n\\n\\nBut with the economy.\\nBut the economy was a big job.\\n\\nBut the economy was a big\",\n",
       " 'ablation': {'layers': [0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]}}"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "items = store.filter(**item)\n",
    "item"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "store.to_disk()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Looking at logit changes\n",
    "We now want to generate objects like this one:\n",
    "```\n",
    "{\n",
    "    'model': 'gpt-2-small',\n",
    "    'prompt': prompt,\n",
    "    'token': pos_id\n",
    "    'logit_diffs': [\n",
    "        {' hello': -10},\n",
    "        ...\n",
    "    ]\n",
    "    'ablation': ablation\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "reference_text = \"The distance between the Colosseum and the Eiffel is approximately\"\n",
    "texts = [\n",
    "    \"The distance between Rome and Paris is approximately\",\n",
    "    \"1 + 3 = 4 \\n4 + 5 = 9\\n2 + 3 =\",\n",
    "    json.dumps([{'name': 'James', 'age': 34, 'skills': ['Python', 'git']}, {'name': 'Alan', 'age': 28, 'skills': ['MS Office', '<cut>']}]).split('<cut>')[0],\n",
    "    \"Obama was elected in the year\",\n",
    "]\n",
    "for prompt in texts:\n",
    "    for i in range(0, 13):\n",
    "        keep_layer_ids =  list(range(i))\n",
    "        permuted_model = permuted(demo_gpt2, keep_layer_ids)\n",
    "        text, _ = complete(permuted_model, prompt)\n",
    "        print(f\"Layer {i}: {text}\")\n",
    "        store({\n",
    "            'model': 'gpt-2-small',\n",
    "            'prompt': prompt,\n",
    "            'completion': text,\n",
    "            'ablation': {\n",
    "                'layers': keep_layer_ids\n",
    "            }\n",
    "        })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate tasks agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running until done Agent\n",
      "Counted tokens: 235 for message: Please generate more 5 tasks such as these: [{\"name\": \"JSON formatting\", \"description\": \"Can the mod ...\n",
      "Cache miss\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "module 'openai' has no attribute 'ChatCompletion'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32m/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb Cell 32\u001b[0m line \u001b[0;36m8\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=71'>72</a>\u001b[0m     eval_questions: List[\u001b[39mstr\u001b[39m] \u001b[39m=\u001b[39m Field(\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, description\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mEvaluation questions\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=74'>75</a>\u001b[0m task_create_agent \u001b[39m=\u001b[39m Agent(\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=75'>76</a>\u001b[0m     [],\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=76'>77</a>\u001b[0m     system_message\u001b[39m=\u001b[39msystem_prompt,\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=77'>78</a>\u001b[0m     prompt_template\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{query}\u001b[39;00m\u001b[39m\"\u001b[39m\u001b[39m.\u001b[39mformat,\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=78'>79</a>\u001b[0m     response_openapi\u001b[39m=\u001b[39mNLPTask\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=79'>80</a>\u001b[0m )\n\u001b[0;32m---> <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=82'>83</a>\u001b[0m result \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m task_create_agent\u001b[39m.\u001b[39mrun(query\u001b[39m=\u001b[39muser_prompt)\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/nielswarncke/Documents/code/TransformerLens/demos/Layer_Permutation.ipynb#X64sZmlsZQ%3D%3D?line=83'>84</a>\u001b[0m tasks \u001b[39m=\u001b[39m result[\u001b[39m'\u001b[39m\u001b[39mtasks\u001b[39m\u001b[39m'\u001b[39m]\n",
      "File \u001b[0;32m~/Documents/agi/minichain/minichain/agent.py:109\u001b[0m, in \u001b[0;36mAgent.run\u001b[0;34m(self, conversation, message_meta, **arguments)\u001b[0m\n\u001b[1;32m    102\u001b[0m message_meta \u001b[39m=\u001b[39m message_meta \u001b[39mor\u001b[39;00m {}\n\u001b[1;32m    103\u001b[0m \u001b[39mawait\u001b[39;00m agent_session\u001b[39m.\u001b[39mconversation\u001b[39m.\u001b[39msend(\n\u001b[1;32m    104\u001b[0m     UserMessage(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprompt_template(\u001b[39m*\u001b[39m\u001b[39m*\u001b[39marguments)),\n\u001b[1;32m    105\u001b[0m     is_initial\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m,\n\u001b[1;32m    106\u001b[0m     is_initial_user_message\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m,\n\u001b[1;32m    107\u001b[0m     \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mmessage_meta\n\u001b[1;32m    108\u001b[0m )\n\u001b[0;32m--> 109\u001b[0m response \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m agent_session\u001b[39m.\u001b[39mrun_until_done()\n\u001b[1;32m    110\u001b[0m \u001b[39mreturn\u001b[39;00m response\n",
      "File \u001b[0;32m~/Documents/agi/minichain/minichain/agent.py:152\u001b[0m, in \u001b[0;36mSession.run_until_done\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    150\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mrunning until done\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39magent\u001b[39m.\u001b[39mname)\n\u001b[1;32m    151\u001b[0m \u001b[39mwhile\u001b[39;00m \u001b[39mTrue\u001b[39;00m:\n\u001b[0;32m--> 152\u001b[0m     action \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mget_next_action()\n\u001b[1;32m    153\u001b[0m     \u001b[39mif\u001b[39;00m action \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39mand\u001b[39;00m action\u001b[39m.\u001b[39mget(\u001b[39m'\u001b[39m\u001b[39mname\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    154\u001b[0m         output \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mexecute_action(action)\n",
      "File \u001b[0;32m~/Documents/agi/minichain/minichain/agent.py:174\u001b[0m, in \u001b[0;36mSession.get_next_action\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    172\u001b[0m \u001b[39m# do the openai call\u001b[39;00m\n\u001b[1;32m    173\u001b[0m \u001b[39masync\u001b[39;00m \u001b[39mwith\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconversation\u001b[39m.\u001b[39mto(AssistantMessage()) \u001b[39mas\u001b[39;00m message_handler:\n\u001b[0;32m--> 174\u001b[0m     llm_response \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m get_openai_response_stream(\n\u001b[1;32m    175\u001b[0m         history,\n\u001b[1;32m    176\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39magent\u001b[39m.\u001b[39mfunctions_openai,\n\u001b[1;32m    177\u001b[0m         model\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39magent\u001b[39m.\u001b[39mllm,\n\u001b[1;32m    178\u001b[0m         stream\u001b[39m=\u001b[39mmessage_handler,\n\u001b[1;32m    179\u001b[0m         force_call\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_force_call,\n\u001b[1;32m    180\u001b[0m     )\n\u001b[1;32m    181\u001b[0m \u001b[39mreturn\u001b[39;00m llm_response\u001b[39m.\u001b[39mget(\u001b[39m'\u001b[39m\u001b[39mfunction_call\u001b[39m\u001b[39m'\u001b[39m)\n",
      "File \u001b[0;32m~/Documents/agi/minichain/minichain/utils/disk_cache.py:81\u001b[0m, in \u001b[0;36mAsyncDiskCache.cache.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     79\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mCache miss\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m     80\u001b[0m \u001b[39mif\u001b[39;00m stream:\n\u001b[0;32m---> 81\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m func(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs, stream\u001b[39m=\u001b[39mstream)\n\u001b[1;32m     82\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m     83\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m func(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m~/Documents/agi/minichain/minichain/utils/cached_openai.py:86\u001b[0m, in \u001b[0;36mget_openai_response_stream\u001b[0;34m(chat_history, functions, model, stream, force_call)\u001b[0m\n\u001b[1;32m     84\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m     85\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mlen\u001b[39m(functions) \u001b[39m>\u001b[39m \u001b[39m0\u001b[39m:\n\u001b[0;32m---> 86\u001b[0m         openai_response \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m openai\u001b[39m.\u001b[39;49mChatCompletion\u001b[39m.\u001b[39macreate(\n\u001b[1;32m     87\u001b[0m             model\u001b[39m=\u001b[39mmodel,\n\u001b[1;32m     88\u001b[0m             messages\u001b[39m=\u001b[39mmessages,\n\u001b[1;32m     89\u001b[0m             functions\u001b[39m=\u001b[39mfunctions,\n\u001b[1;32m     90\u001b[0m             temperature\u001b[39m=\u001b[39m\u001b[39m0.1\u001b[39m,\n\u001b[1;32m     91\u001b[0m             stream\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m,\n\u001b[1;32m     92\u001b[0m             function_call\u001b[39m=\u001b[39mforce_call\n\u001b[1;32m     93\u001b[0m         )\n\u001b[1;32m     94\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[1;32m     95\u001b[0m         openai_response \u001b[39m=\u001b[39m \u001b[39mawait\u001b[39;00m openai\u001b[39m.\u001b[39mChatCompletion\u001b[39m.\u001b[39macreate(\n\u001b[1;32m     96\u001b[0m             model\u001b[39m=\u001b[39mmodel,\n\u001b[1;32m     97\u001b[0m             messages\u001b[39m=\u001b[39mmessages,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    100\u001b[0m             function_call\u001b[39m=\u001b[39mforce_call\n\u001b[1;32m    101\u001b[0m         )\n",
      "\u001b[0;31mAttributeError\u001b[0m: module 'openai' has no attribute 'ChatCompletion'"
     ]
    }
   ],
   "source": [
    "texts = [\n",
    "    \"The distance between Rome and Paris is approximately\",\n",
    "    \"1 + 3 = 4 \\n4 + 5 = 9\\n2 + 3 =\",\n",
    "    json.dumps([{'name': 'James', 'age': 34, 'skills': ['Python', 'git']}, {'name': 'Alan', 'age': 28, 'skills': ['MS Office', '<cut>']}]).split('<cut>')[0],\n",
    "    \"Obama was elected in the year\",\n",
    "]\n",
    "\n",
    "import random\n",
    "def get_addition_prompt(digit_i, digit_j):\n",
    "    a = random.randint(10**digit_i, 10*10**digit_i)\n",
    "    b = random.randint(10**digit_j, 10*10**digit_j)\n",
    "    c = a + b\n",
    "    return f\"{a} + {b} = {c}\"\n",
    "\n",
    "\n",
    "def get_addition_task_description(digit_i, digit_j):\n",
    "    return {\n",
    "        \"name\": \"Addition\",\n",
    "        \"description\": f\"{digit_i}-digit + {digit_j}-digit addition\",\n",
    "        \"prompts\": [\n",
    "            \n",
    "        ]\n",
    "    }\n",
    "\n",
    "tasks = [\n",
    "    {\n",
    "        'name': 'JSON formatting',\n",
    "        'description': 'Can the model generate correct JSON?',\n",
    "        'prompts': [\n",
    "            json.dumps([{'name': 'James', 'age': 34, 'skills': ['Python', 'git']}, {'name': 'Alan', 'age': 28, 'skills': ['MS Office', '<cut>']}]).split('<cut>')[0]\n",
    "        ],\n",
    "        'eval_questions': [\n",
    "            'Does the completion resemble JSON?',\n",
    "            'Is the completion correct JSON?'\n",
    "        ],\n",
    "    },\n",
    "    {\n",
    "        'name': 'Facts about the world',\n",
    "        'description': 'Generate facts about the world',\n",
    "        'prompts': [\n",
    "            'Madrid is the captial of',\n",
    "            'The capital of Spain is',\n",
    "            'Obama was elected in the year',\n",
    "            'The distance between Rome and Paris is approximately',\n",
    "            'World war 1 started in 1914 and ended in',\n",
    "        ],\n",
    "        'eval_questions': [\n",
    "            'Repeat the completion from the beginning until the last word that is still grammatically correct English: (leave empty if no correct English is found)',\n",
    "            'Does the completion resemble correct knowledge?'\n",
    "            'Does the completion show fully correct knowledge?'\n",
    "        ]\n",
    "    },\n",
    "]\n",
    "\n",
    "for j in range(3):\n",
    "    for i in range(3):\n",
    "        tasks += [get_addition_task_description(i, j)]\n",
    "\n",
    "\n",
    "system_prompt = \"You are a research assistant that comes up with tasks for language models.\"\n",
    "\n",
    "user_prompt = f\"Please generate more 5 tasks such as these: {json.dumps(tasks[:2])}. Also add more prompts to my examples please, thank you!\"\n",
    "\n",
    "from minichain.agent import Agent\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import List\n",
    "\n",
    "class NLPTask(BaseModel):\n",
    "    name: str = Field(..., description=\"Name of the task\")\n",
    "    description: str = Field(..., description=\"Description of the task\")\n",
    "    prompts: str = Field(..., description=\"A list of prompts for this task. Should be at least 20 prompts.\")\n",
    "    eval_questions: List[str] = Field(..., description=\"Evaluation questions\")\n",
    "\n",
    "\n",
    "task_create_agent = Agent(\n",
    "    [],\n",
    "    system_message=system_prompt,\n",
    "    prompt_template=\"{query}\".format,\n",
    "    response_openapi=NLPTask\n",
    ")\n",
    "\n",
    "\n",
    "result = await task_create_agent.run(query=user_prompt)\n",
    "tasks = result['tasks']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import openai.error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment Loop Agent\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lens",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}