mmcs_across_layers_across_l1.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Dictionaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformer_lens import HookedTransformer\n",
    "import numpy as np \n",
    "from torch import nn\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from tqdm import tqdm\n",
    "\n",
    "def mcs(smaller_dict, larger_dict):\n",
    "    smaller_dict, larger_dict\n",
    "\n",
    "    #Dictionary Comparison\n",
    "    smaller_dict_features, _ = smaller_dict.shape\n",
    "    larger_dict_features, _ = larger_dict.shape\n",
    "    larger_dict = larger_dict.to(device)\n",
    "    # Hungary algorithm\n",
    "    # Calculate all cosine similarities and store in a 2D array\n",
    "    cos_sims = np.zeros((smaller_dict_features, larger_dict_features))\n",
    "    for idx, vector in enumerate(smaller_dict):\n",
    "        cos_sims[idx] = torch.nn.functional.cosine_similarity(vector.to(device), larger_dict, dim=1).cpu().numpy()\n",
    "    # Convert to a minimization problem\n",
    "    cos_sims = 1 - cos_sims\n",
    "    # Use the Hungarian algorithm to solve the assignment problem\n",
    "    row_ind, col_ind = linear_sum_assignment(cos_sims)\n",
    "    # Retrieve the max cosine similarities and corresponding indices\n",
    "    max_cosine_similarities = 1 - cos_sims[row_ind, col_ind]\n",
    "\n",
    "    return max_cosine_similarities, row_ind, col_ind\n",
    "\n",
    "#Change these settings to load the correct autoencoder\n",
    "setting = \"residual\"\n",
    "model_name = \"EleutherAI/pythia-70m-deduped\"\n",
    "# model_name = \"EleutherAI/pythia-160m-deduped\"\n",
    "l1_index = 0\n",
    "dict_index = 2\n",
    "layers = [0,1,2,3,4,5]\n",
    "autoencoders = []\n",
    "dictionaries = []\n",
    "dictionaries_2 = []\n",
    "\n",
    "device = torch.device(\"cuda:3\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = HookedTransformer.from_pretrained(model_name, device=device)\n",
    "all_layers_MMCS = []\n",
    "all_layers_above_thresold = []\n",
    "for layer in tqdm(layers):\n",
    "    this_layer_MMCS = []\n",
    "    this_layer_above_threshold = []\n",
    "\n",
    "    filename = f\"/mnt/ssd-cluster/bigrun0308/tied_residual_l{layer}_r1/_9/learned_dicts.pt\"\n",
    "    all_autoencoders = torch.load(filename)\n",
    "    num_autoencoders = len(all_autoencoders)\n",
    "    all_dictionaries = [autoencoder.get_learned_dict() for autoencoder, _ in all_autoencoders]\n",
    "    for dict_index in range(num_autoencoders-1):\n",
    "        max_cs, _, _= mcs(all_dictionaries[dict_index], all_dictionaries[dict_index+1])\n",
    "        this_layer_MMCS.append(max_cs.mean())\n",
    "        this_layer_above_threshold.append((max_cs > 0.9).sum())\n",
    "    all_layers_MMCS.append(this_layer_MMCS)\n",
    "    all_layers_above_thresold.append(this_layer_above_threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, hyperparams = all_autoencoders[0]\n",
    "dict_size = hyperparams['dict_size']\n",
    "all_l1s = [hyperparams['l1_alpha'] for _, hyperparams in all_autoencoders]\n",
    "all_l1s = np.array(all_l1s)\n",
    "# Plot them \n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "# create a color scheme \n",
    "colors = ['b', 'g', 'r', 'c', 'm', 'y']\n",
    "\n",
    "# Make two subplots\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\n",
    "# Plot the data\n",
    "for layer in range(len(layers)):\n",
    "    specific_layer = layers[layer]\n",
    "    ax1.plot(all_l1s[:-1], all_layers_MMCS[layer], label=str(specific_layer) + \" MMCS\", color=colors[layer])\n",
    "    # Make this one dashed\n",
    "    ax2.plot(all_l1s[:-1], np.array(all_layers_above_thresold[layer])/dict_size, label=str(specific_layer) + ' above threshold', color=colors[layer], linestyle='dashed')\n",
    "\n",
    "# Add a legend\n",
    "ax1.set_title('MMCS across layers') \n",
    "ax2.set_title('Percent of Features above threshold across layers')\n",
    "ax1.set_xlabel('L1')\n",
    "ax2.set_xlabel('L1')\n",
    "# log xaxis\n",
    "ax1.set_xscale('log')\n",
    "ax2.set_xscale('log')\n",
    "ax1.set_ylabel('MMCS')\n",
    "ax2.set_ylabel('Percent of Features above threshold')\n",
    "ax1.legend(loc='upper right')\n",
    "ax2.legend(loc='upper right')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_l1_ind = []\n",
    "for layer in range(len(layers)):\n",
    "    specific_layer = layers[layer]\n",
    "    l1_ind = np.array(all_layers_MMCS[specific_layer]).argmax()\n",
    "    best_l1_ind.append(l1_ind)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_l1_ind"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCS\n",
    "Across Dictionaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import linear_sum_assignment\n",
    "import matplotlib.pyplot as plt\n",
    "# Loop through the dictionaries in dictionaries\n",
    "fig, axs = plt.subplots(4, figsize=(20, 30))\n",
    "\n",
    "for i in range(len(dictionaries)-1):\n",
    "    smaller_dict, larger_dict = dictionaries[i], dictionaries[i+1]\n",
    "\n",
    "    #Dictionary Comparison\n",
    "    smaller_dict_features, _ = smaller_dict.shape\n",
    "    larger_dict_features, _ = larger_dict.shape\n",
    "    larger_dict = larger_dict.to(device)\n",
    "    # Hungary algorithm\n",
    "    # Calculate all cosine similarities and store in a 2D array\n",
    "    cos_sims = np.zeros((smaller_dict_features, larger_dict_features))\n",
    "    for idx, vector in enumerate(smaller_dict):\n",
    "        cos_sims[idx] = torch.nn.functional.cosine_similarity(vector.to(device), larger_dict, dim=1).cpu().numpy()\n",
    "    # Convert to a minimization problem\n",
    "    cos_sims = 1 - cos_sims\n",
    "    # Use the Hungarian algorithm to solve the assignment problem\n",
    "    row_ind, col_ind = linear_sum_assignment(cos_sims)\n",
    "    # Retrieve the max cosine similarities and corresponding indices\n",
    "    max_cosine_similarities = 1 - cos_sims[row_ind, col_ind]\n",
    "\n",
    "    # Get the indices of the max cosine similarities in descending order\n",
    "    max_indices = np.argsort(max_cosine_similarities)[::-1]\n",
    "    max_cosine_similarities[max_indices][:20]\n",
    "    print((\"# of features above 0.9:\", (max_cosine_similarities > .9).sum()))\n",
    "\n",
    "    # Plot histogram of max_cosine_similarities\n",
    "    axs[i].hist(max_cosine_similarities, bins=20)\n",
    "    axs[i].set_title(f\"MCS between Residual Layer Dicts {i+1} and {i+2}\")\n",
    "    axs[i].set_xlabel(\"MCS\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Across dictionaries of same Layer for MCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import linear_sum_assignment\n",
    "import matplotlib.pyplot as plt\n",
    "# Loop through the dictionaries in dictionaries\n",
    "fig, axs = plt.subplots(6, figsize=(20, 30))\n",
    "max_indices = []\n",
    "for i in range(len(dictionaries)):\n",
    "    smaller_dict, larger_dict = dictionaries[i], dictionaries_2[i]\n",
    "\n",
    "    #Dictionary Comparison\n",
    "    smaller_dict_features, _ = smaller_dict.shape\n",
    "    larger_dict_features, _ = larger_dict.shape\n",
    "    larger_dict = larger_dict.to(device)\n",
    "    # Hungary algorithm\n",
    "    # Calculate all cosine similarities and store in a 2D array\n",
    "    cos_sims = np.zeros((smaller_dict_features, larger_dict_features))\n",
    "    for idx, vector in enumerate(smaller_dict):\n",
    "        cos_sims[idx] = torch.nn.functional.cosine_similarity(vector.to(device), larger_dict, dim=1).cpu().numpy()\n",
    "    # Convert to a minimization problem\n",
    "    cos_sims = 1 - cos_sims\n",
    "    # Use the Hungarian algorithm to solve the assignment problem\n",
    "    row_ind, col_ind = linear_sum_assignment(cos_sims)\n",
    "    # Retrieve the max cosine similarities and corresponding indices\n",
    "    max_cosine_similarities = 1 - cos_sims[row_ind, col_ind]\n",
    "\n",
    "    # Get the indices of the max cosine similarities in descending order\n",
    "    select_max_indices = np.argsort(max_cosine_similarities)[::-1]\n",
    "    max_indices.append(select_max_indices)\n",
    "    print((\"# of features above 0.9:\", (max_cosine_similarities > .9).sum()))\n",
    "\n",
    "    # Plot histogram of max_cosine_similarities\n",
    "    axs[i].hist(max_cosine_similarities, bins=20)\n",
    "    axs[i].set_title(f\"MCS between Residual Layer Dicts of diff sizes\")\n",
    "    axs[i].set_xlabel(\"MCS\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model activations & Dictionary Activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Downnload dataset\n",
    "from datasets import Dataset, load_dataset\n",
    "dataset_name = \"NeelNanda/pile-10k\"\n",
    "token_amount= 30\n",
    "dataset = load_dataset(dataset_name, split=\"train[:4000]\").map(\n",
    "    lambda x: model.tokenizer(x['text']),\n",
    "    batched=True,\n",
    ").filter(\n",
    "    lambda x: len(x['input_ids']) > token_amount\n",
    ").map(\n",
    "    lambda x: {'input_ids': x['input_ids'][:token_amount]}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we can use the model to get the activations\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "from einops import rearrange\n",
    "# neurons = model.W_in.shape[-1]\n",
    "neurons = model.cfg.d_model\n",
    "datapoints = dataset.num_rows\n",
    "batch_size = 32\n",
    "all_dict_activations = []\n",
    "all_neuron_activations = []\n",
    "\n",
    "with torch.no_grad(), dataset.formatted_as(\"pt\"):\n",
    "    dl = DataLoader(dataset[\"input_ids\"], batch_size=batch_size)\n",
    "    for layer in range(len(layers)):\n",
    "        neuron_activations = torch.zeros((datapoints*token_amount, neurons))\n",
    "        dictionary_activations = torch.zeros((datapoints*token_amount, smaller_dict_features))\n",
    "        for i, batch in enumerate(tqdm(dl)):\n",
    "            current_autoencoder = autoencoders[layer]\n",
    "            current_autoencoder.to_device(device)\n",
    "            \n",
    "            _, cache = model.run_with_cache(batch.to(device))\n",
    "            cache_name = f\"blocks.{layers[layer]}.hook_resid_post\"\n",
    "            \n",
    "            batched_neuron_activations = rearrange(cache[cache_name], \"b s n -> (b s) n\" )\n",
    "            neuron_activations[i*batch_size*token_amount:(i+1)*batch_size*token_amount,:] = batched_neuron_activations.cpu()\n",
    "\n",
    "            batched_dictionary_activations = current_autoencoder.encode(batched_neuron_activations)\n",
    "            dictionary_activations[i*batch_size*token_amount:(i+1)*batch_size*token_amount,:] = batched_dictionary_activations.cpu()\n",
    "\n",
    "        all_dict_activations.append(dictionary_activations)\n",
    "        all_neuron_activations.append(neuron_activations)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Activation Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Useful Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from circuitsvis.activations import text_neuron_activations\n",
    "# Get the activations for the best dict features\n",
    "def get_feature_datapoints(feature_index, dictionary_activations, dataset, k=10, setting=\"max\"):\n",
    "    best_feature_activations = dictionary_activations[:, feature_index]\n",
    "    # Sort the features by activation, get the indices\n",
    "    if setting==\"max\":\n",
    "        found_indices = torch.argsort(best_feature_activations, descending=True)[:k]\n",
    "    elif setting==\"uniform\":\n",
    "        # min_value = torch.min(best_feature_activations)\n",
    "        min_value = torch.min(best_feature_activations)\n",
    "        max_value = torch.max(best_feature_activations)\n",
    "\n",
    "        # Define the number of bins\n",
    "        num_bins = k\n",
    "\n",
    "        # Calculate the bin boundaries as linear interpolation between min and max\n",
    "        bin_boundaries = torch.linspace(min_value, max_value, num_bins + 1)\n",
    "\n",
    "        # Assign each activation to its respective bin\n",
    "        bins = torch.bucketize(best_feature_activations, bin_boundaries)\n",
    "\n",
    "        # Initialize a list to store the sampled indices\n",
    "        sampled_indices = []\n",
    "\n",
    "        # Sample from each bin\n",
    "        for bin_idx in torch.unique(bins):\n",
    "            if(bin_idx==0): # Skip the first one. This is below the median\n",
    "                continue\n",
    "            # Get the indices corresponding to the current bin\n",
    "            bin_indices = torch.nonzero(bins == bin_idx, as_tuple=False).squeeze(dim=1)\n",
    "            \n",
    "            # Randomly sample from the current bin\n",
    "            sampled_indices.extend(np.random.choice(bin_indices, size=1, replace=False))\n",
    "\n",
    "        # Convert the sampled indices to a PyTorch tensor & reverse order\n",
    "        found_indices = torch.tensor(sampled_indices).long().flip(dims=[0])\n",
    "    else: # random\n",
    "        # get nonzero indices\n",
    "        nonzero_indices = torch.nonzero(best_feature_activations)[:, 0]\n",
    "        # shuffle\n",
    "        shuffled_indices = nonzero_indices[torch.randperm(nonzero_indices.shape[0])]\n",
    "        found_indices = shuffled_indices[:k]\n",
    "    datapoint_indices =[np.unravel_index(i, (datapoints, token_amount)) for i in found_indices]\n",
    "    text_list = []\n",
    "    full_text = []\n",
    "    token_list = []\n",
    "    full_token_list = []\n",
    "    for md, s_ind in datapoint_indices:\n",
    "        md = int(md)\n",
    "        s_ind = int(s_ind)\n",
    "        full_tok = torch.tensor(dataset[md][\"input_ids\"])\n",
    "        full_text.append(model.tokenizer.decode(full_tok))\n",
    "        tok = dataset[md][\"input_ids\"][:s_ind+1]\n",
    "        text = model.tokenizer.decode(tok)\n",
    "        text_list.append(text)\n",
    "        token_list.append(tok)\n",
    "        full_token_list.append(full_tok)\n",
    "    return text_list, full_text, token_list, full_token_list\n",
    "\n",
    "def get_neuron_activation(token, feature, model, autoencoder, layer, setting=\"dictionary_basis\"):\n",
    "    with torch.no_grad():\n",
    "        _, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "        cache_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "        neuron_act_batch = cache[cache_name]\n",
    "        if setting==\"dictionary_basis\":\n",
    "            act = autoencoder.encode(neuron_act_batch.squeeze())\n",
    "            return act[:, feature].tolist()\n",
    "        else: # neuron/residual basis\n",
    "            return neuron_act_batch[0, :, feature].tolist()\n",
    "\n",
    "def ablate_text(text, feature, model, autoencoder, layer,  setting=\"plot\"):\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    display_text_list = []\n",
    "    activation_list = []\n",
    "    for t in text:\n",
    "        # Convert text into tokens\n",
    "        if isinstance(t, str): # If the text is a list of tokens\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            tokens = model.to_tokens(t, prepend_bos=False)\n",
    "        else: # t equals tokens\n",
    "            tokens = t\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "        seq_size = tokens.shape[1]\n",
    "        if(seq_size == 1): # If the text is a single token, we can't ablate it\n",
    "            continue\n",
    "        original = get_neuron_activation(tokens, feature, model, autoencoder, layer)[-1]\n",
    "        changed_activations = torch.zeros(seq_size, device=device).cpu()\n",
    "        for i in range(seq_size):\n",
    "            # Remove the i'th token from the input\n",
    "            ablated_tokens = torch.cat((tokens[:,:i], tokens[:,i+1:]), dim=1)\n",
    "            changed_activations[i] += get_neuron_activation(ablated_tokens, feature, model, autoencoder, layer)[-1]\n",
    "        changed_activations -= original\n",
    "        display_text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "        activation_list += changed_activations.tolist() + [0.0]\n",
    "    activation_list = torch.tensor(activation_list).reshape(-1,1,1)\n",
    "    if setting == \"plot\":\n",
    "        return text_neuron_activations(tokens=display_text_list, activations=activation_list)\n",
    "    else:\n",
    "        return display_text_list, activation_list\n",
    "        \n",
    "def visualize_text(text, feature, model,autoencoder, layer, setting=\"dictionary_basis\", max_activation = None):\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    if isinstance(feature, int):\n",
    "        feature = [feature]\n",
    "    display_text_list = []\n",
    "    act_list = []\n",
    "    for t in text:\n",
    "        if isinstance(t, str): # If the text is a list of tokens\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            token = model.to_tokens(t, prepend_bos=False)\n",
    "        else: # t are tokens\n",
    "            token = t\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "        for f in feature:\n",
    "            display_text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            act_list += get_neuron_activation(token, f, model, autoencoder, layer, setting) + [0.0]\n",
    "    act_list = torch.tensor(act_list).reshape(-1,1,1)\n",
    "    if(max_activation is not None):\n",
    "        act_list = torch.clamp(act_list, max=max_activation)\n",
    "    return text_neuron_activations(tokens=display_text_list, activations=act_list)\n",
    "# Ablate the feature direction of the tokens\n",
    "# token_list is a list of tokens, convert to tensor of shape (batch_size, seq_len)\n",
    "from einops import rearrange\n",
    "def ablate_feature_direction(tokens, feature, model, autoencoder, layer):\n",
    "    def mlp_ablation_hook(value, hook):\n",
    "        # Rearrange to fit autoencoder\n",
    "        int_val = rearrange(value, 'b s h -> (b s) h')\n",
    "\n",
    "        # Run through the autoencoder\n",
    "        act = autoencoder.encode(int_val)\n",
    "        feature_to_ablate = feature # TODO: bring this out of the function\n",
    "\n",
    "        dictionary_for_this_autoencoder = autoencoder.get_learned_dict()\n",
    "        feature_direction = torch.outer(act[:, feature_to_ablate].squeeze(), dictionary_for_this_autoencoder[feature_to_ablate].squeeze())\n",
    "        batch, seq_len, hidden_size = value.shape\n",
    "        feature_direction = rearrange(feature_direction, '(b s) h -> b s h', b=batch, s=seq_len)\n",
    "        value -= feature_direction\n",
    "        return value\n",
    "    \n",
    "    cache_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "    return model.run_with_hooks(tokens, \n",
    "        fwd_hooks=[(\n",
    "            cache_name, \n",
    "            mlp_ablation_hook\n",
    "            )]\n",
    "        )\n",
    "def add_feature_direction(tokens, feature, model, autoencoder, scalar=1.0):\n",
    "    def residual_add_hook(value, hook):\n",
    "        feature_direction = autoencoder.decoder.weight[:, feature].squeeze()\n",
    "        value += scalar*feature_direction\n",
    "        return value\n",
    "\n",
    "    return model.run_with_hooks(tokens, \n",
    "        fwd_hooks=[(\n",
    "            cache_name,\n",
    "            residual_add_hook\n",
    "            )]\n",
    "        )\n",
    "def ablate_feature_direction_display(text, autoencoder,layer, features=None, setting=\"true_tokens\", verbose=False):\n",
    "\n",
    "    if features==None:\n",
    "        features = torch.tensor([best_feature])\n",
    "    if isinstance(features, int):\n",
    "        features = torch.tensor([features])\n",
    "    if isinstance(features, list):\n",
    "        features = torch.tensor(features)\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    text_list = []\n",
    "    logit_list = []\n",
    "    for t in text:\n",
    "        tokens = model.to_tokens(t, prepend_bos=False)\n",
    "        with torch.no_grad():\n",
    "            original_logits = model(tokens).log_softmax(-1).cpu()\n",
    "            ablated_logits = ablate_feature_direction(tokens, features, model, autoencoder, layer).log_softmax(-1).cpu()\n",
    "        diff_logits = ablated_logits  - original_logits# ablated > original -> negative diff\n",
    "        tokens = tokens.cpu()\n",
    "        if setting == \"true_tokens\":\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            gather_tokens = rearrange(tokens[:,1:], \"b s -> b s 1\") # TODO: verify this is correct\n",
    "            # Gather the logits for the true tokens\n",
    "            diff = rearrange(diff_logits[:, :-1].gather(-1,gather_tokens), \"b s n -> (b s n)\")\n",
    "        elif setting == \"max\":\n",
    "            # Negate the diff_logits to see which tokens have the largest effect on the neuron\n",
    "            val, ind = (-1*diff_logits).max(-1)\n",
    "            diff = rearrange(val[:, :-1], \"b s -> (b s)\")\n",
    "            diff*= -1 # Negate the values gathered\n",
    "            split_text = model.to_str_tokens(ind, prepend_bos=False)\n",
    "            gather_tokens = rearrange(ind[:,1:], \"1 s -> 1 s 1\")\n",
    "        split_text = split_text[1:] # Remove the first token since we're not predicting it\n",
    "        if(verbose):\n",
    "            text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            orig = rearrange(original_logits[:, :-1].gather(-1, gather_tokens), \"b s n -> (b s n)\")\n",
    "            ablated = rearrange(ablated_logits[:, :-1].gather(-1, gather_tokens), \"b s n -> (b s n)\")\n",
    "            logit_list += orig.tolist() + [0.0]\n",
    "            logit_list += ablated.tolist() + [0.0]\n",
    "        text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "        logit_list += diff.tolist() + [0.0]\n",
    "    logit_list = torch.tensor(logit_list).reshape(-1,1,1)\n",
    "    if verbose:\n",
    "        print(f\"Max & Min logit-diff: {logit_list.max().item():.2f} & {logit_list.min().item():.2f}\")\n",
    "    return text_neuron_activations(tokens=text_list, activations=logit_list)\n",
    "\n",
    "def generate_text(input_text, num_tokens, model, autoencoder, feature, temperature=0.7, setting=\"add\", scalar=1.0):\n",
    "    # Convert input text to tokens\n",
    "    input_ids = model.tokenizer.encode(input_text, return_tensors='pt').to(device)\n",
    "\n",
    "    for _ in range(num_tokens):\n",
    "        # Generate logits\n",
    "        with torch.no_grad():\n",
    "            if(setting==\"add\"):\n",
    "                logits = add_feature_direction(input_ids, feature, model, autoencoder, scalar=scalar)\n",
    "            else:\n",
    "                logits = model(input_ids)\n",
    "\n",
    "        # Apply temperature\n",
    "        logits = logits / temperature\n",
    "\n",
    "        # Sample from the distribution\n",
    "        probs = torch.nn.functional.softmax(logits[:, -1, :], dim=-1)\n",
    "        predicted_token = torch.multinomial(probs, num_samples=1)\n",
    "\n",
    "        # Append predicted token to input_ids\n",
    "        input_ids = torch.cat((input_ids, predicted_token), dim=-1)\n",
    "\n",
    "    # Decode the tokens to text\n",
    "    output_text = model.tokenizer.decode(input_ids[0])\n",
    "\n",
    "    return output_text\n",
    "\n",
    "# Logit Lens\n",
    "def logit_lens(model, best_feature, dictionary, layer):\n",
    "    with torch.no_grad():\n",
    "        # There are never-used tokens, which have high norm. We want to ignore these.\n",
    "        bad_ind = (model.W_U.norm(dim=0) > 20)\n",
    "        feature_direction = dictionary[best_feature].to(device)\n",
    "        # feature_direction = torch.matmul(feature_direction, model.W_out[layer]) # if MLP\n",
    "        logits = torch.matmul(feature_direction, model.W_U).cpu()\n",
    "    # Don't include bad indices\n",
    "    logits[bad_ind] = -1000\n",
    "    topk_values, topk_indices = torch.topk(logits, 20)\n",
    "    top_text = model.to_str_tokens(topk_indices)\n",
    "    print(f\"{top_text}\")\n",
    "    print(topk_values)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "# Ablate feature direction for all features in layer 4, see how they effect feature 325 in layer 5\n",
    "# Get datapoints that activate our given feature\n",
    "\n",
    "def find_causal_features(target_feature, target_layer,previous_layer,  k=10, setting=\"ablate\", data_setting=\"max\"):\n",
    "    target_layer_dict_activations = all_dict_activations[target_layer-1][:, target_feature]\n",
    "    # non_zero_activations = torch.nonzero(target_layer_dict_activations[:, target_feature]).squeeze()\n",
    "    # Instead, get the top-10 activations\n",
    "    if(data_setting==\"max\"):    \n",
    "        non_zero_values, non_zero_indices = torch.topk(target_layer_dict_activations, k)\n",
    "    else: # 0.5*max to max\n",
    "        max_value = torch.max(target_layer_dict_activations)\n",
    "        min_value = torch.min(target_layer_dict_activations)\n",
    "        min_value = (max_value - min_value) * 0.5\n",
    "\n",
    "        # Define the number of bins\n",
    "        num_bins = k\n",
    "\n",
    "        # Calculate the bin boundaries as linear interpolation between min and max\n",
    "        bin_boundaries = torch.linspace(min_value, max_value, num_bins + 1)\n",
    "\n",
    "        # Assign each activation to its respective bin\n",
    "        bins = torch.bucketize(target_layer_dict_activations, bin_boundaries)\n",
    "\n",
    "        # Initialize a list to store the sampled indices\n",
    "        sampled_indices = []\n",
    "\n",
    "        # Sample from each bin\n",
    "        for bin_idx in torch.unique(bins):\n",
    "            if(bin_idx==0): # Skip the first one. This is below the median\n",
    "                continue\n",
    "            # Get the indices corresponding to the current bin\n",
    "            bin_indices = torch.nonzero(bins == bin_idx, as_tuple=False).squeeze(dim=1)            \n",
    "            # Randomly sample from the current bin\n",
    "            sampled_indices.extend(np.random.choice(bin_indices, size=1, replace=False))\n",
    "\n",
    "        # Convert the sampled indices to a PyTorch tensor & reverse order\n",
    "        non_zero_indices = torch.tensor(sampled_indices).long().flip(dims=[0])\n",
    "        non_zero_values = target_layer_dict_activations[non_zero_indices]\n",
    "    non_zero_values, non_zero_indices\n",
    "    datapoint_indices = [np.unravel_index(i, (datapoints, token_amount)) for i in non_zero_indices]\n",
    "    token_list = []\n",
    "\n",
    "    max_activating_token_positions = []\n",
    "    for datapoint_index, max_activating_token_pos in datapoint_indices:\n",
    "        datapoint_index = int(datapoint_index)\n",
    "        max_activating_token_pos = int(max_activating_token_pos)\n",
    "        tok = dataset[datapoint_index][\"input_ids\"][:max_activating_token_pos+1]\n",
    "        token_list.append(tok)\n",
    "        max_activating_token_positions.append(max_activating_token_pos)\n",
    "    max_len = max(len(seq) for seq in token_list)\n",
    "\n",
    "    # Pad all sequences to the length of the longest sequence\n",
    "    padded_sequences = torch.tensor([seq + [model.tokenizer.pad_token_id] * (max_len - len(seq)) for seq in token_list])\n",
    "\n",
    "    def ablate_feature_direction_hook(value, hook, feature_to_ablate):\n",
    "        # Rearrange to fit autoencoder\n",
    "        int_val = rearrange(value, 'b s h -> (b s) h')\n",
    "\n",
    "        # Run through the autoencoder\n",
    "        _, act = autoencoders[previous_layer-1](int_val)\n",
    "        \n",
    "        # Subtract value with feature direction*act_of_feature\n",
    "        feature_direction = torch.outer(act[:, feature_to_ablate].squeeze(), autoencoders[previous_layer-1].decoder.weight[:, feature_to_ablate].squeeze())\n",
    "        batch, seq_len, hidden_size = value.shape\n",
    "        feature_direction = rearrange(feature_direction, '(b s) h -> b s h', b=batch, s=seq_len)\n",
    "        value -= feature_direction\n",
    "        return value\n",
    "\n",
    "    def restore_feature_direction_hook(value, hook, feature_to_restore):\n",
    "        # Rearrange to fit autoencoder\n",
    "        int_val = rearrange(value, 'b s h -> (b s) h')\n",
    "\n",
    "        # Run through the autoencoder\n",
    "        _, act = autoencoders[previous_layer-1](int_val)\n",
    "        \n",
    "        \n",
    "        # Subtract value with feature direction*act_of_feature\n",
    "        feature_direction = torch.outer(act[:, feature_to_restore].squeeze(), autoencoders[previous_layer-1].decoder.weight[:, feature_to_restore].squeeze())\n",
    "        batch, seq_len, hidden_size = value.shape\n",
    "        feature_direction = rearrange(feature_direction, '(b s) h -> b s h', b=batch, s=seq_len)\n",
    "        value *= 0\n",
    "        value += feature_direction\n",
    "        return value\n",
    "\n",
    "    def save_activations_hook(value, hook, setting=\"ablate\"):\n",
    "        batch_size, seq_size, _ = value.shape\n",
    "        # Rearrange to fit autoencoder\n",
    "        int_val = rearrange(value, 'b s h -> (b s) h')\n",
    "        max_target_pos = [max_activating_token_positions[i] + i*seq_size for i in range(batch_size)]\n",
    "        # Run through the autoencoder\n",
    "        _, act = autoencoders[target_layer-1](int_val)\n",
    "        changed_feature_values = act[max_target_pos, target_feature].cpu()\n",
    "        if(setting==\"ablate\"):\n",
    "            # activations.append((non_zero_values - changed_feature_values).mean().item())\n",
    "            activations.append(((non_zero_values - changed_feature_values + 1e-10) / (non_zero_values+ 1e-10)).mean().item())\n",
    "        else:\n",
    "            activations.append(((changed_feature_values + 1e-10)/(non_zero_values+ 1e-10)).mean().item())\n",
    "        return value\n",
    "    \n",
    "    activations = []\n",
    "    total_number_of_features = dictionary_activations.shape[1]\n",
    "    # total_number_of_features = 20\n",
    "    for feature_to_ablate in range(total_number_of_features):\n",
    "        if(setting==\"ablate\"):\n",
    "            hooks=[(f\"blocks.{previous_layer}.hook_resid_post\", partial(ablate_feature_direction_hook, feature_to_ablate=feature_to_ablate)), (f\"blocks.{target_layer}.hook_resid_post\", partial(save_activations_hook, setting=setting))]\n",
    "        else:\n",
    "            hooks=[(f\"blocks.{previous_layer}.hook_resid_post\", partial(restore_feature_direction_hook, feature_to_restore=feature_to_ablate)), (f\"blocks.{target_layer}.hook_resid_post\", partial(save_activations_hook, setting=setting))]\n",
    "        with torch.no_grad():\n",
    "            model.run_with_hooks(padded_sequences, fwd_hooks=hooks)\n",
    "    return torch.tensor(activations)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygraphviz as pgv\n",
    "from IPython.display import Image\n",
    "features_found_in_layers = []\n",
    "target_feature = 325\n",
    "target_layer = 5\n",
    "# setting = \"ablate\"\n",
    "setting = \"restore\"\n",
    "# data_setting = \"max\"\n",
    "data_setting = \"median_to_max\"\n",
    "features_found_in_layers.append([target_feature])\n",
    "num_datapoints_k = 20  # Number of top (or uniform) activating examples to consider when ablating\n",
    "graph_children_init_k = 5\n",
    "graph_children_recurse_k= 3\n",
    "\n",
    "G = pgv.AGraph(directed=True)\n",
    "for layer in range(target_layer, 1, -1): \n",
    "    current_features = []\n",
    "    for feature in features_found_in_layers[-1]:\n",
    "        activations = find_causal_features(feature, layer, k=num_datapoints_k, previous_layer=layer-1, setting=setting, data_setting=data_setting)\n",
    "        if(layer == target_layer):\n",
    "            values, indices = activations.topk(graph_children_init_k)\n",
    "            current_features = indices.tolist()\n",
    "        else:\n",
    "            values, indices = activations.topk(graph_children_recurse_k)\n",
    "            current_features.append(indices[0].item())\n",
    "        for v, i in zip(values, indices):\n",
    "            # Calculate the CS between the feature and the causal feature\n",
    "            v1 = dictionaries[layer-1][feature]\n",
    "            v2 = dictionaries[layer-2][i.item()]\n",
    "            cs = torch.nn.functional.cosine_similarity(v1, v2, dim=0)\n",
    "            G.add_edge(f\"{layer}_{feature}\", f\"{layer-1}_{i}\", label=f\"{int(v.item()*100)}% | {cs.item():.2f}\")\n",
    "    features_found_in_layers.append(list(set(current_features)))\n",
    "    print(features_found_in_layers)\n",
    "G.layout(prog=\"dot\")\n",
    "G.draw('file.png')  # Draws a png\n",
    "\n",
    "# # Display the graph\n",
    "Image(filename='file.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# G.draw(\"file.png\", args='-Gdpi=200', prog=\"dot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from sklearn.manifold import TSNE\n",
    "# import matplotlib.pyplot as plt\n",
    "\n",
    "# # Run the t-SNE algorithm\n",
    "# tsne = TSNE(n_components=2, random_state=0, perplexity=5)\n",
    "# vectors_2d = tsne.fit_transform(vectors)\n",
    "\n",
    "# # Plot the vectors\n",
    "# plt.scatter(vectors_2d[:, 0], vectors_2d[:, 1])\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GG = G"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Interp\n",
    "Investigate the example sentences the activate this feature.\n",
    "\n",
    "Max: show max activating (tokens,contexts)\n",
    "\n",
    "Uniform: Show range of activations from each bin (e.g. sample an example from 1-2, 2-3, etc). \n",
    "[Note: if a feature is monosemantic, then the full range of activations should be that feature, not just max-activating ones]\n",
    "\n",
    "Full_text: shows the full text example\n",
    "\n",
    "Text_list: shows up to the most activating example (try w/ max activating on a couple of examples to see)\n",
    "\n",
    "ablate_text: remove the context one token at a time, and show the decrease/increase in activation of that feature\n",
    "\n",
    "ablate_feature_direction: removes feature direction from model's activation mid-inference, showing the logit diff in the output for every token.\n",
    "\n",
    "logit_lens: show the logit lens for that feature. If matches ablate_feature_direction, then the computation path is through the residual stream, else, it's through future layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Layer 5\n",
    "# GG.get_node(\"5_325\").attr[\"label\"] = \"5_325\\nLast word in parantheses. Causes end parantheses\"\n",
    "# # Layer 4\n",
    "# GG.get_node(\"4_1527\").attr[\"label\"] = \"4_1527\\nYears after (\"\n",
    "# GG.get_node(\"4_1456\").attr[\"label\"] = \"4_1456\\nMath variables after (, predicts closing )\"\n",
    "# GG.get_node(\"4_1280\").attr[\"label\"] = \"4_1280\\nDates, particularly years\"\n",
    "# # Layer 3\n",
    "# GG.get_node(\"3_1407\").attr[\"label\"] = \"3_1407\\nYears, especially after (\"\n",
    "# GG.get_node(\"3_663\").attr[\"label\"] = \"3_663\\nYears right after (, specifically trial/court data\"\n",
    "# GG.get_node(\"3_1360\").attr[\"label\"] = \"3_1360\\nEnding numbers in trial/court data\"\n",
    "# GG.get_node(\"3_891\").attr[\"label\"] = \"3_891\\nasides/phrases\"\n",
    "# GG.get_node(\"3_1273\").attr[\"label\"] = \"3_1273\\nAcronyms\"\n",
    "# GG.get_node(\"3_1594\").attr[\"label\"] = \"3_1594\\nWords & Dates?\"\n",
    "# GG.layout(prog=\"dot\")\n",
    "# G.draw('file.png')  # Draws a png\n",
    "\n",
    "# # # Display the graph\n",
    "# Image(filename='file.png') #[[325], [1280, 1709, 1456, 1814, 1527], [0, 331, 1227, 210, 1407]]\n",
    "# GG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = 11\n",
    "layer_index = layers.index(layer)\n",
    "autoencoder = autoencoders[layer_index]\n",
    "# t = \" I do like a\"\n",
    "t = \" He had a first one (1), and then a second (2\"\n",
    "# t = \"Das Berghain ist ein bekannter Nachtclub in Berlin, der für\"\n",
    "split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "token = model.to_tokens(t, prepend_bos=False)\n",
    "_, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "neuron_act_batch = cache[cache_name]\n",
    "act = autoencoder.encode(neuron_act_batch.squeeze())\n",
    "# neg = \" He had a first one (1), and then a second 2\"\n",
    "# split_text = model.to_str_tokens(neg, prepend_bos=False)\n",
    "# token = model.to_tokens(neg, prepend_bos=False)\n",
    "# _, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "# neuron_act_batch = cache[cache_name]\n",
    "# act[-1, :] -= smaller_auto_encoder.encode(neuron_act_batch.squeeze())[-1,:]\n",
    "v, i = act[-1, :].topk(10)\n",
    "\n",
    "print(\"Activations:\",[round(val,2) for val in v.tolist()])\n",
    "print(\"Feature_ids\", i.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = 11\n",
    "layer_index = layers.index(layer)\n",
    "autoencoder = autoencoders[layer_index]\n",
    "autoencoder.to_device(device)\n",
    "dictionary = dictionaries[layer_index]\n",
    "dictionary_activations = all_dict_activations[layer_index]\n",
    "select_max_indices = max_indices[layer_index]\n",
    "# N = 36\n",
    "# best_feature = int(select_max_indices[N])\n",
    "best_feature = 853\n",
    "# ind=torch.tensor([1030,  796, 1527, 1786, 1814, 1333, 1273, 1946, 1156, 1163])\n",
    "# ind=torch.tensor([1527, 1814, 1280, 1709,  809, 1786, 1333,  148,  796, 1647])\n",
    "# ind = [[325], [1030, 1333, 1814, 1527, 796], [1273, 662, 1407], [888, 46, 662], [497, 1619, 614]] # ablated\n",
    "# ind = [[325], [1280, 1709, 1333, 1814, 1527], [0, 331, 662, 1407], [46, 157, 662, 445], [497, 1466, 435, 1527]] # restored\n",
    "# ind = [[325], [1030, 459, 491, 1456, 1527], [1227, 747, 1648, 891, 1407], [808, 714, 46, 184, 989], [1730, 583, 1160, 497, 1749]]\n",
    "# ind = [[325], [1280, 1709, 1456, 1814, 1527], [0, 331, 1227, 210, 1407], [1989, 714, 46, 662, 157], [0, 497, 1619, 1466, 1564]]\n",
    "\n",
    "# selected_index = 7\n",
    "# layer_ind = ind[::-1][layer-1]\n",
    "# best_feature = int(select_max_indices[selected_index])\n",
    "# best_feature = int(layer_ind[selected_index])\n",
    "# best_feature = \n",
    "print(\"bias:\", autoencoder.encoder_bias.detach().cpu().numpy()[best_feature])\n",
    "print(f\"Feature index: {best_feature}\")\n",
    "print(f\"MCS: {max_cosine_similarities[best_feature]}\")\n",
    "text_list, full_text, token_list, full_token_list = get_feature_datapoints(best_feature, dictionary_activations, dataset, setting=\"uniform\")\n",
    "# text_list, full_text, token_list, full_token_list = get_feature_datapoints(best_feature, dictionary_activations, dataset, setting=\"max\")\n",
    "# visualize_text(text_list, best_feature, model, autoencoder, layer)\n",
    "visualize_text(full_text, best_feature, model, autoencoder, layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_text(text_list, best_feature, model, autoencoder, layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ablate_feature_direction_display(full_text, best_feature, layer)\n",
    "ablate_feature_direction_display(full_text, autoencoder, layer, features=best_feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logit_lens(model, best_feature, dictionary, layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Given a feature direction in layer 4, multiply by the layer norm2, and then multiply by the MLP? I need to be careful about running it through the entire MLP\n",
    "\n",
    "layer = 3\n",
    "feature = 1273\n",
    "feature_direction = rearrange(dictionaries[layer-1][feature], 'h -> 1 1 h').to(device)\n",
    "with torch.no_grad():\n",
    "    # # MLP\n",
    "    # feature_direction = model.blocks[layer].ln2(feature_direction)\n",
    "    # translated_feature_direction = rearrange(model.blocks[layer].mlp(feature_direction), '1 1 h -> h')\n",
    "    # Attention\n",
    "    # feature_direction = model.blocks[layer].ln1(feature_direction)\n",
    "    # TODO: check if position_ids make sense\n",
    "    translated_feature_direction = model_transformers.gpt_neox.layers[layer].attention(feature_direction, attention_mask=None, position_ids=torch.tensor([[0]]).to(device))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now do cosine similarity between the translated feature and future layer\n",
    "future_layer_dictionary = dictionaries[layer].to(device)\n",
    "cs = torch.nn.functional.cosine_similarity(translated_feature_direction, future_layer_dictionary, dim=1)\n",
    "print(f\"ind = {cs.topk(10).indices.cpu().tolist()[0]}\\nind = {cs.topk(10, largest=False).indices.cpu().tolist()[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download pythia-70m from transformer library\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "model_transformers = AutoModelForCausalLM.from_pretrained(model_name).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert torch tensor to list\n",
    "def tensor_to_list(tensor):\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = \" I do like a\"\n",
    "split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "token = model.to_tokens(t, prepend_bos=False)\n",
    "_, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "neuron_act_batch = cache[cache_name]\n",
    "_, act = smaller_auto_encoder(neuron_act_batch)\n",
    "v, i = act[0, -1, :].topk(10)\n",
    "\n",
    "print(\"Activations:\",[round(val,2) for val in v.tolist()])\n",
    "print(\"Feature_ids\", i.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence = \" for\"\n",
    "temp = 0.7\n",
    "tokens_to_generate = 20\n",
    "feature = 10 \n",
    "scalar = 100.0\n",
    "# Using the function:\n",
    "print(\"Normal:\\n\" + generate_text(sentence, tokens_to_generate, model, smaller_auto_encoder, feature=feature, temperature=temp, scalar=scalar, setting=\"normal\"))\n",
    "print(\"Add:\\n\" + generate_text(sentence, tokens_to_generate, model, smaller_auto_encoder, feature=feature, temperature=temp, scalar=scalar, setting=\"add\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smaller_dict[best_feature].topk(10), smaller_dict[best_feature].topk(10, largest=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# OUTlier dimensions negative:\n",
    "#  568,  516, 1894,  468, 1299, 1326,  412, 1458,  934,  615, 1147,  672, 1377, 1518, 1889,  715, 1325,\n",
    "# outlier dims positive:\n",
    "# 412,  481,  305, 1894,  319, 1887, 1643, 1876,  801,  615,  362, 1564, 1323,   25,  989, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check features non-zero weights in decoder\n",
    "# Plot a histogram of the weights\n",
    "weights = smaller_dict[best_feature]\n",
    "plt.hist(weights, bins=30)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights.topk(10, largest=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replicate activation of position\n",
    "# plot as hist\n",
    "import matplotlib.pyplot as plt\n",
    "for i in range(5):\n",
    "    # make new fig\n",
    "    plt.figure()\n",
    "    act = neuron_activations[3:100000:token_amount, outlier_dimension.indices[i]]\n",
    "    plt.title(f\"Dimension {outlier_dimension.indices[i]}\")\n",
    "    plt.hist(act, bins=100)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from collections import Counter\n",
    "# For the outlier values that aren't first position, what are the tokens?\n",
    "# Find index of outlier values that aren't first position\n",
    "for di in all_dict_activations:\n",
    "    outlier_dimensions = di.max(dim=0).values.topk(10)\n",
    "    for i in range(5):\n",
    "        outlier_indices = torch.where(di[:10000, outlier_dimensions.indices[i]] > 0)[0]\n",
    "        # Only do first indice now\n",
    "        outlier_indices = outlier_indices[outlier_indices % token_amount != 0]\n",
    "        # Find the tokens associated w/ them\n",
    "        indices = [np.unravel_index(oi, (datapoints, token_amount)) for oi in outlier_indices]\n",
    "        token_list = []    \n",
    "        for md, s_ind in indices:\n",
    "            md = int(md)\n",
    "            s_ind = int(s_ind)\n",
    "            tok = dataset[md][\"input_ids\"][s_ind]\n",
    "            token_list.append(model.tokenizer.decode(tok))\n",
    "        token_list\n",
    "        # Count the unique occurrences of the tokens\n",
    "        Counter(token_list)\n",
    "        print(f\"{i} {Counter(token_list)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for di in all_dict_activations:\n",
    "    outlier_dimensions = di.max(dim=0).values.topk(10).indices\n",
    "    print(outlier_dimensions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_dict_activations[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dictionary_activations[:, best_feature].count_nonzero().sum() / dictionary_activations.shape[0], dictionary_activations[:, best_feature].count_nonzero(dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nz = dictionary_activations.count_nonzero(dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nz.median() / dictionary_activations.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_text = [\n",
    "    \"I can count up to: 2 4 8 16 32 64 128 256 512 1024 2048 4096 8192 16384 32768 7 6 12 16 18 20 22 24\",\n",
    "]\n",
    "visualize_text(custom_text, best_feature, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Centric Viewpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Go through datapoints & see if the features that activate on them make sense.\n",
    "d_point = 0\n",
    "# text = tokens_dataset[d_point]\n",
    "data_ind, sequence_pos = np.unravel_index(d_point, (datapoints, token_amount))\n",
    "feature_val, feature_ind = dictionary_activations[d_point].topk(10)\n",
    "data_ind = int(data_ind)\n",
    "sequence_pos = int(sequence_pos)\n",
    "full_tok = torch.tensor(dataset[data_ind][\"input_ids\"])\n",
    "full_text = []\n",
    "full_text.append(model.tokenizer.decode(full_tok))\n",
    "visualize_text(full_text, feature_ind, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check the neuron/residual basis\n",
    "When we look at the weights of a feature, we are seeing the literal dimensions from the residual stream/neurons being read from the feature. \n",
    "\n",
    "Here I'm visualizing the weight values for the residual stream. If there are outliers, then it's mainly reading from that dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(weights*max_activation).topk(20), (weights*max_activation).topk(20, largest=False).values, (weights*max_activation > 0.2).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepend/Append tokens\n",
    "We can iterate over all tokens to check which ones activate a feature a lot to more rigorously test a hypothesis on what a feature means.\n",
    "\n",
    "Note: I'm literately running the model through all 50k tokens prepended to the text here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepend_all_tokens_and_get_feature_activation(model, minimal_activating_example, feature, setting=\"prepend\"):\n",
    "    tokens = model.to_tokens(minimal_activating_example, prepend_bos=False)\n",
    "\n",
    "    # Run through every number up to vocab size\n",
    "    vocab_size = model.cfg.d_vocab\n",
    "    batch_size = 256*2 # Define your desired batch size\n",
    "\n",
    "    dollar_feature_activations = torch.zeros(vocab_size)\n",
    "    for start in range(0, vocab_size, batch_size):\n",
    "        end = min(start + batch_size, vocab_size)\n",
    "\n",
    "        token_prep = torch.arange(start, end).to(device)\n",
    "        token_prep = token_prep.unsqueeze(1)  # Add a dimension for concatenation\n",
    "\n",
    "        # 1. Prepend to the tokens\n",
    "        if setting == \"prepend\":\n",
    "            tokens_catted = torch.cat((token_prep, tokens.repeat(end - start, 1)), dim=1).long()\n",
    "        elif setting == \"append\":\n",
    "            tokens_catted = torch.cat((tokens.repeat(end - start, 1), token_prep), dim=1).long()\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown setting: {setting}\")\n",
    "\n",
    "        # 2. Run through the model\n",
    "        with torch.no_grad():\n",
    "            _, cache = model.run_with_cache(tokens_catted.to(device))\n",
    "            neuron_act_batch = cache[cache_name]\n",
    "            _, act = smaller_auto_encoder(neuron_act_batch)\n",
    "\n",
    "        # 3. Get the feature\n",
    "        dollar_feature_activations[start:end] = act[:, -1, feature].cpu().squeeze()\n",
    "\n",
    "    k = 20\n",
    "    k_increasing_val, k_increasing_ind = dollar_feature_activations.topk(k)\n",
    "    k_decreasing_val, k_decreasing_ind = dollar_feature_activations.topk(k, largest=False)\n",
    "    if(setting == \"prepend\"):\n",
    "        print(f\"[token]{minimal_activating_example}\")\n",
    "    elif(setting == \"append\"):\n",
    "        print(f\"{minimal_activating_example}[token]\")\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown setting: {setting}\")\n",
    "    # Print indices converted to tokens\n",
    "    print(f\"Top-{k} increasing: {model.to_str_tokens(k_increasing_ind)}\")\n",
    "    # Print values\n",
    "    print(f\"Top-{k} increasing: {[f'{val:.2f}' for val in k_increasing_val]}\")\n",
    "    print(f\"Top-{k} decreasing: {model.to_str_tokens(k_decreasing_ind)}\")\n",
    "    print(f\"Top-{k} decreasing: {[f'{val:.2f}' for val in k_decreasing_val]}\")\n",
    "    print(f\"Number of 0 activations: {torch.sum(dollar_feature_activations == 0)}\")\n",
    "    if(setting == \"prepend\"):\n",
    "        best_text = \"\".join(model.to_str_tokens(dollar_feature_activations.argmax()) + [minimal_activating_example])\n",
    "    else:\n",
    "        best_text = \"\".join([minimal_activating_example] + model.to_str_tokens(dollar_feature_activations.argmax()))\n",
    "    return best_text\n",
    "\n",
    "best_text = \"\"\n",
    "for x in range(3):\n",
    "    # best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"prepend\")\n",
    "    best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"append\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" for all $\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"append\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_text = \"\"\n",
    "for x in range(3):\n",
    "    best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"prepend\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepend_all_tokens_and_get_feature_activation(model, \" for all $\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \"The\", best_feature, setting=\"append\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"append\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}