feature_interp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformer_lens import HookedTransformer\n",
    "import numpy as np"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Dictionaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from torch import nn\n",
    "\n",
    "class AutoEncoder(nn.Module):\n",
    "    def __init__(self, activation_size, n_dict_components, t_type=torch.float32, l1_coef=0.0):\n",
    "        super(AutoEncoder, self).__init__()\n",
    "        \n",
    "        # Only defining the decoder layer, encoder will share its weights\n",
    "        self.decoder = nn.Linear(n_dict_components, activation_size, bias=True)\n",
    "        # Create a bias layer\n",
    "        self.encoder_bias= nn.Parameter(torch.zeros(n_dict_components))\n",
    "\n",
    "        \n",
    "        # Initialize the decoder weights orthogonally\n",
    "        nn.init.orthogonal_(self.decoder.weight)\n",
    "        self.decoder = self.decoder.to(t_type)\n",
    "\n",
    "        # Encoder is a Sequential with the ReLU activation\n",
    "        # No need to define a Linear layer for the encoder as its weights are tied with the decoder\n",
    "        self.encoder = nn.Sequential(nn.ReLU()).to(t_type)\n",
    "\n",
    "        self.l1_coef = l1_coef\n",
    "        self.activation_size = activation_size\n",
    "        self.n_dict_components = n_dict_components\n",
    "\n",
    "    def forward(self, x):\n",
    "        c = self.encoder(x @ self.decoder.weight + self.encoder_bias)\n",
    "        # Apply unit norm constraint to the decoder weights\n",
    "        self.decoder.weight.data = nn.functional.normalize(self.decoder.weight.data, dim=0)\n",
    "\n",
    "        # Decoding step as before\n",
    "        x_hat = self.decoder(c)\n",
    "        return x_hat, c\n",
    "\n",
    "\n",
    "    @property\n",
    "    def device(self):\n",
    "        return next(self.parameters()).device\n",
    "    \n",
    "\n",
    "layer = 15\n",
    "setting = \"residual\"\n",
    "# setting = \"mlp\"\n",
    "model_name = \"EleutherAI/pythia-1.4b-deduped\"\n",
    "# model_name = \"EleutherAI/pythia-70m-deduped\"\n",
    "# filename = r'/root/sparse_coding/autoencoders4.pkl'\n",
    "# filename = r'/root/sparse_coding/auto_encoders.pkl'\n",
    "# filename = r'outputs/20230706-225337/auto_encoders_15.pkl'\n",
    "filename = r'outputs/20230712-132914-EleutherAI/pythia-1.4b-deduped-15/minirun9/autoencoders.pkl'\n",
    "\n",
    "if setting == \"residual\":\n",
    "    cache_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "elif setting == \"mlp\":\n",
    "    cache_name = f\"blocks.{layer}.mlp.hook_post\"\n",
    "else:\n",
    "    raise NotImplementedError\n",
    "\n",
    "# Load in the Pythia model w/ transformer lens\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = HookedTransformer.from_pretrained(model_name, device=device)\n",
    "if(setting == \"residual\"):\n",
    "    neurons = model.cfg.d_model\n",
    "elif(setting == \"mlp\"):\n",
    "    neurons = model.cfg.d_mlp\n",
    "else:\n",
    "    raise NotImplementedError\n",
    "\n",
    "# Load the pickle file\n",
    "with open(filename, 'rb') as file:\n",
    "    autoencoders = pickle.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l1_index = -1\n",
    "dictionaries = [autoencoder.decoder.weight.data.T for autoencoder in autoencoders[l1_index]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in dictionaries:\n",
    "    print(d.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(autoencoders)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dict_index = 2\n",
    "smaller_dict, larger_dict = dictionaries[dict_index], dictionaries[dict_index+1]\n",
    "smaller_auto_encoder, larger_auto_encoder = autoencoders[l1_index][dict_index], autoencoders[l1_index][dict_index+1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dictionary Comparison\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\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",
    "from scipy.optimize import linear_sum_assignment\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",
    "print(max_cosine_similarities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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((max_cosine_similarities > .9).sum())\n",
    "# Plot histogram of max_cosine_similarities\n",
    "import matplotlib.pyplot as plt\n",
    "plt.hist(max_cosine_similarities, bins=50)\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=64\n",
    "dataset = load_dataset(dataset_name, split=\"train\").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",
    "\n",
    "datapoints = dataset.num_rows\n",
    "batch_size = 32\n",
    "neuron_activations = torch.zeros((datapoints*token_amount, neurons))\n",
    "dictionary_activations = torch.zeros((datapoints*token_amount, smaller_dict_features))\n",
    "smaller_auto_encoder = smaller_auto_encoder.to(device)\n",
    "\n",
    "with torch.no_grad(), dataset.formatted_as(\"pt\"):\n",
    "    dl = DataLoader(dataset[\"input_ids\"], batch_size=batch_size)\n",
    "    for i, batch in enumerate(tqdm(dl)):\n",
    "        _, cache = model.run_with_cache(batch.to(device))\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",
    "        reconstruction, batched_dictionary_activations = smaller_auto_encoder(batched_neuron_activations)\n",
    "        dictionary_activations[i*batch_size*token_amount:(i+1)*batch_size*token_amount,:] = batched_dictionary_activations.cpu()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neuron Sparsity per Feature by Activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find corresponding neurons activations between pythia & autoencoder\n",
    "# Get the activations for the best dict features\n",
    "thresholds = [round(0.1*x,1) for x in range(11)]\n",
    "all_activated_neurons = []\n",
    "for threshold in thresholds:\n",
    "    current_activated_neurons = []\n",
    "    for x in range(20):\n",
    "        best_feature = max_indices[x]\n",
    "        best_feature_activations = dictionary_activations[:, best_feature]\n",
    "        # Sort the features by activation, get the indices\n",
    "        nonzero_indices = torch.argsort(best_feature_activations, descending=True)\n",
    "        sorted_indices = nonzero_indices[:10]\n",
    "        t = (neuron_activations[sorted_indices, :] > threshold)\n",
    "        # ( And across the first dim)\n",
    "        t = t.all(dim=0)\n",
    "        neurons_activated = t.sum()\n",
    "        current_activated_neurons.append(neurons_activated)\n",
    "        # print(f\"Feature {x} is active for {t.sum()} neurons\")\n",
    "    all_activated_neurons.append(current_activated_neurons)\n",
    "# Plot boxplot w/ plotly\n",
    "\n",
    "plt.boxplot(all_activated_neurons, labels=thresholds)\n",
    "plt.xticks(rotation=45)\n",
    "plt.xlabel(\"Threshold\")\n",
    "plt.ylabel(\"Number of neurons activated\")\n",
    "plt.title(\"Features/Neurons activated\")\n",
    "plt.ylim(0, 50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find number of non-zero activations for each feature in the dictionary\n",
    "nonzero_activations = dictionary_activations.count_nonzero(dim=0)\n",
    "# clip to 1\n",
    "nonzero_activations = nonzero_activations.clamp(max=1)\n",
    "# plot against the max cosine similarities\n",
    "# plt.hist(nonzero_activations.cpu().numpy(), bins=100)\n",
    "plt.scatter(max_cosine_similarities, nonzero_activations.cpu().numpy())\n",
    "# x-axis is the max cosine similarity\n",
    "# y-axis is the number of non-zero activations\n",
    "# now setting x-axis\n",
    "plt.xlabel(\"Max Cosine Similarity\")\n",
    "plt.ylabel(\"Number of Non-Zero Activations\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Activation Analysis"
   ]
  },
  {
   "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",
    "        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",
    "            # 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):\n",
    "    with torch.no_grad():\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",
    "    return act[0, :, feature].tolist()\n",
    "\n",
    "def ablate_text(text, feature, model, 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)[-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)[-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",
    "def visualize_text(text, feature, model, setting=\"plot\", 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) + [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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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):\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(int_val)\n",
    "        feature_to_ablate = feature # TODO: bring this out of the function\n",
    "\n",
    "        # Subtract value with feature direction*act_of_feature\n",
    "        feature_direction = torch.outer(act[:, feature_to_ablate].squeeze(), autoencoder.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",
    "    return model.run_with_hooks(tokens, \n",
    "        fwd_hooks=[(\n",
    "            cache_name,\n",
    "            mlp_ablation_hook\n",
    "            )]\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 visualize_logit_diff(text, features=None, setting=\"true_tokens\", verbose=False):\n",
    "    features = best_feature\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, smaller_auto_encoder).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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll look at feature 52, which has an MCS of 0.96 (w/ perfect MCS at 1.0). To get a feel for the feature, Let's plot the feature activation."
   ]
  },
  {
   "attachments": {},
   "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 = 1\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, setting=\"plot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(feature_counts[max_indices.copy()] == 1).nonzero(), feature_counts[max_indices.copy()].topk(50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "# sentence = \"'What did you do today?' '\"\n",
    "sentence = \" for\"\n",
    "temp = 0.7\n",
    "tokens_to_generate = 20\n",
    "# feature = 1032\n",
    "# feature = 4661 # Food\n",
    "feature = 8149 # python\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": [
    "# sentences = [\" I don't give a shit\", \" You are such a moron. You don't even have a fuck\", \" He is such a son of a bitch\"]\n",
    "# sentences = [\" I really do love\", \" He has such a kind heart\", \" She did the sweetest\"]\n",
    "# sentences = [\" I really do hate\", \" He's the only one that I hate\", \" She hates\"]\n",
    "sentences = [\n",
    "    # \" I just moved into a new flat\",\n",
    "    # \" I took the lift\",\n",
    "    \"\\n\\rfor i in range\",\n",
    "    \"\\n\\rwhile(\",\n",
    "    \"\\n\\rfrom\",\n",
    "    # \" Would you like a biscuit\",\n",
    "    # \" We're having fish and chips\",\n",
    "    # \" She opened a packet of crisps\",\n",
    "    # \" I bought a new jumper\",\n",
    "    # \" He put on his trainers\",\n",
    "    # \" The children were excited to go to the shop for some sweets\",\n",
    "    # \" Could you get another nappy\",\n",
    "    # \" The car ran out of petrol\",\n",
    "    # \" We waited at the zebra crossing\",\n",
    "]\n",
    "\n",
    "whitelist = dict()\n",
    "for t in sentences:\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",
    "    features = act[0, -1, :].topk(50).indices\n",
    "    values = act[0, -1, :].topk(50).values\n",
    "    blacklist = [2555, 303]\n",
    "    for feature, value in zip(features, values):\n",
    "        if feature not in blacklist:\n",
    "            if feature in whitelist:\n",
    "                whitelist[feature.item()] *= value.item()\n",
    "            else:\n",
    "                whitelist[feature.item()] = value.item()\n",
    "print(sorted(list(whitelist.items()), key=lambda x:x[1], reverse=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 11\n",
    "best_feature = int(max_indices[N])\n",
    "best_feature = 1032\n",
    "# best_feature = 6651\n",
    "# best_feature = 6753\n",
    "# best_feature = 3168\n",
    "# best_feature = 1263\n",
    "# best_feature = 1264\n",
    "best_feature = 8149\n",
    "best_feature = 4701\n",
    "print(\"bias:\", smaller_auto_encoder.encoder_bias.detach().cpu().numpy()[best_feature])\n",
    "print(f\"Max feature index: {N}\")\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(full_text, best_feature, model, setting=\"plot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_logit_diff(full_text, best_feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    bad_ind = (model.W_U.norm(dim=0) > 20)\n",
    "    feature_direction = smaller_dict[best_feature].to(device)\n",
    "    # residual_direction = torch.matmul(feature_direction, model.W_out[layer]) # Add bias\n",
    "    # feature_direction = model.ln_final(feature_direction)\n",
    "    logits = torch.matmul(feature_direction, model.W_U).cpu()\n",
    "    # logits = model.W_U(residual_direction).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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ablate_text(sentences, best_feature, model, setting=\"plot\")\n",
    "ablate_text(text_list, best_feature, model, setting=\"plot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instead of top-k examples, just random samples from non-zero values\n",
    "text_list, full_text, _, _ = get_feature_datapoints(best_feature, dictionary_activations, dataset, setting=\"uniform\", k=10)\n",
    "visualize_text(full_text, best_feature, model, setting=\"plot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_text(text_list, best_feature, model, setting=\"plot\")"
   ]
  },
  {
   "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 # 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, \" x\", best_feature, setting=\"prepend\")\n",
    "prepend_all_tokens_and_get_feature_activation(model, \" let\", best_feature, setting=\"append\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"append\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_text(text_list, best_feature, model, setting=\"plot\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Beyond the first token, this doesn't fit within the hypothesis, though may be OOD. A check (for the future) would be to constrain by the model's next word prediction (and maybe do direct soft prompts)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing to the Neuron Basis\n",
    "Does this feature net us anything over using the normal neuron basis? If this is only learning a monosemantic feature, then that's pretty lame! \n",
    "\n",
    "We can first check how many neurons activate above a threshold for the top-10 feature activating examples (ie a neuron must activate above threshold for all 10 examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Neurons that fire for this feature\n",
    "# Find corresponding neurons activations between pythia & autoencoder\n",
    "# Get the activations for the best dict features\n",
    "len_threshold = 11\n",
    "thresholds = [round(0.1*x,1) for x in range(len_threshold)]\n",
    "current_activated_neurons = np.zeros(len_threshold)\n",
    "for idx, threshold in enumerate(thresholds):\n",
    "    best_feature_activations = dictionary_activations[:, best_feature]\n",
    "    # Sort the features by activation, get the indices\n",
    "    nonzero_indices = torch.argsort(best_feature_activations, descending=True)\n",
    "    sorted_indices = nonzero_indices[:10]\n",
    "    t = (neuron_activations[sorted_indices, :] > threshold)\n",
    "    # And across the first dim)\n",
    "    t = t.all(dim=0)\n",
    "    neurons_activated = t.sum()\n",
    "    current_activated_neurons[idx] = neurons_activated\n",
    "    print(f\"Threshold: {threshold}, Neurons activated: {neurons_activated}\")\n",
    "# Plot boxplot w/ plotly\n",
    "plt.scatter(thresholds, current_activated_neurons)\n",
    "plt.xticks(rotation=45)\n",
    "plt.xlabel(\"Threshold\")\n",
    "plt.ylabel(\"Number of neurons activated\")\n",
    "plt.title(\"Features/Neurons activated\")\n",
    "# plt.ylim(0, 10)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notably it's a lot of neurons for above 0, but then goes to 3 for 0.5. However, we don't really know the statistics of neuron activations. Maybe some neuron's entire range is very tiny? So, we can see how many neurons are above a threshold determined by that neuron's quantiles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "len_threshold = 11\n",
    "thresholds = [round(0.1 * x, 1) for x in range(len_threshold)]\n",
    "thresholds[-1] = 0.99\n",
    "current_activated_neurons = np.zeros(len_threshold)\n",
    "\n",
    "# Find corresponding neurons activations between pythia & autoencoder\n",
    "best_feature_activations = dictionary_activations[:, best_feature]\n",
    "\n",
    "# Sort the features by activation and get the indices\n",
    "nonzero_indices = torch.argsort(best_feature_activations, descending=True)\n",
    "sorted_indices = nonzero_indices[:11]\n",
    "\n",
    "neuron_test = neuron_activations[:50000, :]\n",
    "neuron_test = neuron_test.to(device)\n",
    "thresholds_t = torch.tensor(thresholds).to(device)\n",
    "neuron_activations = neuron_activations.to(device)\n",
    "# above_quantile = neuron_test[above_zero_indices].quantile(0.5, dim=0)\n",
    "\n",
    "# Calculate the quantiles for each neuron's activation range\n",
    "quantiles = torch.quantile(neuron_test, thresholds_t, dim=0)\n",
    "\n",
    "for idx, threshold in enumerate(quantiles):\n",
    "    # Determine which neurons activate above their specific quantile\n",
    "    above_threshold_mask = neuron_activations[sorted_indices, :] > threshold.unsqueeze(0)\n",
    "    neurons_activated = above_threshold_mask.float().mean(dim=0).count_nonzero()\n",
    "    current_activated_neurons[idx] = neurons_activated\n",
    "    print(f\"Quantile: {thresholds[idx]}, Neurons activated: {neurons_activated}\")\n",
    "\n",
    "# Plot scatter plot\n",
    "plt.scatter(thresholds, current_activated_neurons)\n",
    "plt.xticks(rotation=45)\n",
    "plt.xlabel(\"Quantile\")\n",
    "plt.ylabel(\"Number of neurons activated\")\n",
    "plt.title(\"Features/Neurons activated\")\n",
    "# plt.ylim(0, 200)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, 80 neurons are activating in above their top 1% of activations.\n",
    "\n",
    "Another way is to look at the decoder of the dictionary & see the weights effect on neurons since it's just a linear layer. Additionally, I'll multiply by the max-activation of that feature to show the scale of the effect on the neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check features non-zero weights in decoder\n",
    "# Plot a histogram of the weights\n",
    "max_activation = dictionary_activations[:, best_feature].max()\n",
    "weights = smaller_dict[best_feature]\n",
    "plt.hist(weights, bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(weights).topk(20, largest=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So a few datapoints to indicate somewhere between 3 & 80 neurons, maybe several hundred depending on how you interpret the "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Logit lens\n",
    "# Multiply feature direction by W_out[layer]\n",
    "# Then by W_U\n",
    "for n in range(52, 53):\n",
    "    b_feature = max_indices[n]\n",
    "    with torch.no_grad():\n",
    "        feature_direction = smaller_dict[b_feature].to(device)\n",
    "        residual_direction = torch.matmul(feature_direction, model.W_out[layer]) # Add bias\n",
    "        # residual_direction = model.ln_final(residual_direction)\n",
    "        logits = torch.matmul(residual_direction, model.W_U).cpu()\n",
    "    topk_values, topk_indices = torch.topk(logits, 10)\n",
    "    top_text = model.to_str_tokens(topk_indices)\n",
    "    print(f\"Feature {n}: {top_text}\")\n",
    "# print(topk_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import circuitsvis\n",
    "# text = full_text\n",
    "# features = best_feature\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, smaller_auto_encoder).log_softmax(-1).cpu()\n",
    "#     # diff_logits = ablated_logits  - original_logits# ablated > original -> negative diff\n",
    "#     diff_logits =   original_logits - ablated_logits# ablated > original -> negative diff\n",
    "#     tokens = tokens.cpu()\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",
    "#     true_log_probs = rearrange(diff_logits[:, :-1].gather(-1,gather_tokens), \"b s n -> (b s n)\")\n",
    "#     break\n",
    "# # Add an extra dim for the batch\n",
    "# diff_logits = diff_logits[0]\n",
    "# tokens = tokens[0]\n",
    "# print(diff_logits.shape, tokens.shape)\n",
    "# # circuitsvis.logits.token_log_probs(token_indices=tokens, top_k=10, log_probs=original_logits, to_string=model.to_single_str_token)\n",
    "# circuitsvis.logits.token_log_probs(token_indices=tokens, top_k=10, log_probs=diff_logits, to_string=model.to_single_str_token)\n",
    "# # circuitsvis.logits.token_log_probs(token_indices=tokens, top_k=10, log_probs=ablated_logits, to_string=model.to_single_str_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "#Plot a histogram\n",
    "import matplotlib.pyplot as plt\n",
    "for x in range(10):\n",
    "    max_elements = (dictionary_activations[:, max_indices[x]]>0.01)\n",
    "    plt.hist(dictionary_activations[max_elements, max_indices[x]], bins=20)\n",
    "    plt.title('Histogram of Activations for Dictionary Element ' + str(x))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "#Plot a histogram\n",
    "import matplotlib.pyplot as plt\n",
    "for x in range(10):\n",
    "    print((dictionary_activations[:, max_indices[x]] > 0.01).sum().item())"
   ]
  }
 ],
 "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.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}