openai_dictionary_direction.ipynb

{
 "cells": [
  {
   "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",
    "import pickle\n",
    "\n",
    "# Define the autoencoder so pickle knows how to serialize it. \n",
    "# Later, we should actually save as a state_dict instead of a dumb pickle\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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Dictionaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Change these settings to load the correct autoencoder\n",
    "# filename = \"/home/mchorse/logan_sparse_coding/sparse_coding/layer-2-autoencoder.pkl\"\n",
    "# filename = \"outputs/20230720-192301-EleutherAI/pythia-70m-deduped-2/minirun9/autoencoders.pkl\"\n",
    "# filename = \"outputs/20230720-192301-EleutherAI/pythia-70m-deduped-2/minirun5/autoencoders.pkl\"\n",
    "# filename = \"outputs/20230721-161919-EleutherAI/pythia-70m-deduped-2/minirun9/autoencoders.pkl\"\n",
    "# filename = \"outputs/20230721-161919-EleutherAI/pythia-70m-deduped-2/minirun3/autoencoders.pkl\"\n",
    "# filename = \"outputs/20230804-205842-EleutherAI/pythia-70m-deduped-2/minirun4/autoencoders.pkl\"\n",
    "\n",
    "# filename = \"outputs/20230805-205314-EleutherAI/pythia-70m-deduped-2/minirun14/autoencoders.pkl\"\n",
    "filename = \"outputs/20230806-181211-gpt2-8/minirun9/autoencoders.pkl\"\n",
    "layer = 8\n",
    "setting = \"residual\"\n",
    "# model_name = \"EleutherAI/pythia-70m-deduped\"\n",
    "model_name = \"gpt2\"\n",
    "\n",
    "device = torch.device(\"cuda:5\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = HookedTransformer.from_pretrained(model_name, device=device)\n",
    "\n",
    "if setting == \"residual\":\n",
    "    cache_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "    neurons = model.cfg.d_model\n",
    "elif setting == \"mlp\":\n",
    "    cache_name = f\"blocks.{layer}.mlp.hook_post\"\n",
    "    neurons = model.cfg.d_mlp\n",
    "else:\n",
    "    raise NotImplementedError\n",
    "\n",
    "\n",
    "# Load the pickle file\n",
    "with open(filename, 'rb') as file:\n",
    "    autoencoders = pickle.load(file)\n",
    "\n",
    "# Index for l1 value, usually only 1 value is available\n",
    "l1_index = 3\n",
    "dictionaries = [autoencoder.decoder.weight.data.T for autoencoder in autoencoders[l1_index]]\n",
    "for d in dictionaries:\n",
    "    print(d.shape)\n",
    "print(\"len of autoencoders: \", len(autoencoders))\n",
    "dict_index = 0\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]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Change these settings to load the correct autoencoder\n",
    "# filename = \"/home/mchorse/logan_sparse_coding/sparse_coding/layer-2-autoencoder.pkl\"\n",
    "filename = \"outputs/20230720-192301-EleutherAI/pythia-70m-deduped-2/minirun9/autoencoders.pkl\"\n",
    "# filename = \"outputs/20230720-192301-EleutherAI/pythia-70m-deduped-2/minirun5/autoencoders.pkl\"\n",
    "filename2 = \"outputs/20230721-161919-EleutherAI/pythia-70m-deduped-2/minirun9/autoencoders.pkl\"\n",
    "# filename = \"outputs/20230721-161919-EleutherAI/pythia-70m-deduped-2/minirun3/autoencoders.pkl\"\n",
    "layer = 2\n",
    "setting = \"residual\"\n",
    "model_name = \"EleutherAI/pythia-70m-deduped\"\n",
    "\n",
    "device = torch.device(\"cuda:4\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = HookedTransformer.from_pretrained(model_name, device=device)\n",
    "\n",
    "if setting == \"residual\":\n",
    "    cache_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "    neurons = model.cfg.d_model\n",
    "elif setting == \"mlp\":\n",
    "    cache_name = f\"blocks.{layer}.mlp.hook_post\"\n",
    "    neurons = model.cfg.d_mlp\n",
    "else:\n",
    "    raise NotImplementedError\n",
    "\n",
    "\n",
    "# Load the pickle file\n",
    "with open(filename, 'rb') as file:\n",
    "    autoencoders = pickle.load(file)\n",
    "\n",
    "# Load the pickle file\n",
    "with open(filename2, 'rb') as file:\n",
    "    autoencoders2 = pickle.load(file)\n",
    "\n",
    "# Index for l1 value, usually only 1 value is available\n",
    "l1_index = 1\n",
    "dictionaries = [autoencoder.decoder.weight.data.T for autoencoder in autoencoders[l1_index]]\n",
    "dictionaries2 = [autoencoder.decoder.weight.data.T for autoencoder in autoencoders2[l1_index]]\n",
    "for d in dictionaries:\n",
    "    print(d.shape)\n",
    "print(\"len of autoencoders: \", len(autoencoders))\n",
    "dict_index = 1\n",
    "smaller_dict, larger_dict = dictionaries[dict_index], dictionaries2[dict_index]\n",
    "smaller_auto_encoder, larger_auto_encoder = autoencoders[l1_index][dict_index], autoencoders2[l1_index][dict_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import linear_sum_assignment\n",
    "from tqdm import tqdm\n",
    "# DO THE MCS across all dictionaries\n",
    "def mcs(dict1, dict2):\n",
    "    dict1_features, _ = dict1.shape\n",
    "    dict2_features, _ = dict2.shape\n",
    "    dict2 = dict2.to(device)\n",
    "    # Hungary algorithm\n",
    "    # Calculate all cosine similarities and store in a 2D array\n",
    "    cos_sims = np.zeros((dict1_features, dict2_features))\n",
    "    for idx, vector in enumerate(dict1):\n",
    "        cos_sims[idx] = torch.nn.functional.cosine_similarity(vector.to(device), dict2, 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",
    "    max_indices = np.argsort(max_cosine_similarities)[::-1].copy()\n",
    "    return max_cosine_similarities, max_indices\n",
    "\n",
    "\n",
    "l1_index = 3\n",
    "# dictionaries = [autoencoder.decoder.weight.data.T for autoencoder in autoencoders[l1_index]]\n",
    "# dictionaries2 = [autoencoder.decoder.weight.data.T for autoencoder in autoencoders2[l1_index]]\n",
    "dictionaries = [autoencoder[0].decoder.weight.data.T for autoencoder in autoencoders]\n",
    "# combined_dicts = dictionaries + dictionaries2\n",
    "combined_dicts = dictionaries\n",
    "adjancy_matrix = np.zeros((len(combined_dicts), len(combined_dicts)))\n",
    "for i in tqdm(range(len(combined_dicts))):\n",
    "    for j in range(len(combined_dicts)):\n",
    "        if j > i:\n",
    "            break\n",
    "        max_cosine_similarities, max_indices = mcs(combined_dicts[i], combined_dicts[j])\n",
    "        adjancy_matrix[i, j] = max_cosine_similarities.mean()\n",
    "        adjancy_matrix[j, i] = max_cosine_similarities.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the adjancy matrix\n",
    "from matplotlib import pyplot as plt\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(adjancy_matrix, cmap='hot', interpolation='nearest')\n",
    "plt.title(\"MCS for dict_size=4k\")\n",
    "# sparsities = [28,69,94,95,113,115,126,205,393]\n",
    "sparsities = [28, 69, 97, 107, 114, 119, 123, 205, 402]\n",
    "# sparsities are x-axis\n",
    "plt.xticks(range(len(sparsities)), sparsities)\n",
    "plt.xlabel(\"Features/Tokens\")\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the adjancy matrix\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(adjancy_matrix, cmap='hot', interpolation='nearest')\n",
    "plt.title(\"MCS for index: \" + str(l1_index))\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the adjancy matrix\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(adjancy_matrix, cmap='hot', interpolation='nearest')\n",
    "plt.title(\"MCS for index: \" + str(l1_index))\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_cosine_similarities.mean()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCS\n",
    "Max cosine similarity between one dictionary & another one. If they learned the same feature, then they'll have high cosine similarity. \n",
    "\n",
    "If two dictionaries learned it, it's probably a real feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import linear_sum_assignment\n",
    "import matplotlib.pyplot as plt\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].copy()\n",
    "print((\"# of features above 0.9:\", (max_cosine_similarities > .9).sum()))\n",
    "# Plot histogram of max_cosine_similarities\n",
    "plt.hist(max_cosine_similarities, bins=20)\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= 40\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",
    "# neurons = model.W_in.shape[-1]\n",
    "neurons = model.cfg.d_model\n",
    "datapoints = dataset.num_rows\n",
    "batch_size = 64\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": [
    "# 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",
    "        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, setting=\"dictionary_basis\"):\n",
    "    with torch.no_grad():\n",
    "        _, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "        neuron_act_batch = cache[cache_name]\n",
    "        if setting==\"dictionary_basis\":\n",
    "            _, act = smaller_auto_encoder(neuron_act_batch)\n",
    "            return act[0, :, feature].tolist()\n",
    "        else: # neuron/residual basis\n",
    "            return neuron_act_batch[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=\"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, 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, entire_feature_direction=False):\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",
    "    def ablated_this_feature_dir(value, hook):\n",
    "        value -= feature\n",
    "        return value\n",
    "    if(entire_feature_direction):\n",
    "        return model.run_with_hooks(tokens, \n",
    "            fwd_hooks=[(\n",
    "                cache_name, \n",
    "                ablated_this_feature_dir\n",
    "                )]\n",
    "            )\n",
    "    else:\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 ablate_feature_direction_display(text, features=None, setting=\"true_tokens\", verbose=False, entire_feature_direction=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, smaller_auto_encoder, entire_feature_direction).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",
    "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, smaller_dict, 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 = smaller_dict[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": [
    "#Ablate all directions one at a time, to see the effect of the logit on the output\n",
    "# 1. get io & s token\n",
    "# 2. get the logit of both tokens on a normal model run\n",
    "# 3. get the logit of both tokens on a model run with the direction ablated\n",
    "# 4. get the difference between the two logits, then print out topk features\n",
    "ex_sent_A = [\n",
    "    '<|endoftext|>Then, Richard and Jennifer were working at the plateau. Jennifer decided to give a feather to',\n",
    "    '<|endoftext|>Then, Laura and Catherine were working at the cafe. Catherine decided to give a towel to',\n",
    "    '<|endoftext|>Then, Parker and Jake were working at the glacier. Jake decided to give fins to',\n",
    "    '<|endoftext|>Then, Tim and Wyatt were working at the garage. Wyatt decided to give a bible to',\n",
    "    '<|endoftext|>Then, Sam and Peter were working at the chapel. Peter decided to give a remote to',\n",
    "    '<|endoftext|>Then, Ross and Jason were working at the chapel. Jason decided to give a chalk to',\n",
    "    '<|endoftext|>Then, Nathan and Eva were working at the cafe. Eva decided to give a vinyl to',\n",
    "    '<|endoftext|>Then, Matt and Evan were working at the stadium. Evan decided to give a crate to',\n",
    "]\n",
    "ex_sent_B = [\n",
    "    '<|endoftext|>Then, Taylor and Grant were working at the lounge. Taylor decided to give a shorts to',\n",
    "    '<|endoftext|>Then, Jake and Kate were working at the synagogue. Jake decided to give a backpack to',\n",
    "    '<|endoftext|>Then, Julia and Rose were working at the museum. Julia decided to give a wrench to',\n",
    "    '<|endoftext|>Then, Jennifer and Brooklyn were working at the fortress. Jennifer decided to give a tape to',\n",
    "    '<|endoftext|>Then, Mason and Jeff were working at the street. Mason decided to give a vinyl to',\n",
    "    '<|endoftext|>Then, Richard and Jean were working at the cottage. Richard decided to give a lipstick to',\n",
    "    '<|endoftext|>Then, Ross and Elijah were working at the home. Ross decided to give a flag to',\n",
    "    '<|endoftext|>Then, Taylor and Matt were working at the palace. Taylor decided to give a bookmark to',\n",
    "    '<|endoftext|>Then, Molly and Keith were working at the chapel. Molly decided to give a baseball to',\n",
    "    '<|endoftext|>Then, Parker and Dean were working at the tunnel. Parker decided to give a tablet to',\n",
    "    '<|endoftext|>Then, Kevin and Connor were working at the harbor. Kevin decided to give a bakery to',\n",
    "    '<|endoftext|>Then, Matt and Riley were working at the garden. Matt decided to give a charcoal to',\n",
    "    '<|endoftext|>Then, Keith and Joy were working at the plateau. Keith decided to give a ring to',\n",
    "    '<|endoftext|>Then, Ruby and Ray were working at the brewery. Ruby decided to give a lipstick to',\n",
    "]\n",
    "ex_sent = ex_sent_A + ex_sent_B\n",
    "all_A = []\n",
    "all_B = []\n",
    "\n",
    "for sentence in ex_sent_A:\n",
    "    # Split the sentence by 'and' and 'were' to get the names\n",
    "    names = sentence.split('and')[0].split()[-1]\n",
    "    all_A.append(\" \" + names)\n",
    "    \n",
    "    names = sentence.split('and')[1].split('were')[0].strip()\n",
    "    all_B.append(\" \" + names)\n",
    "\n",
    "for sentence in ex_sent_B:\n",
    "    # Split the sentence by 'and' and 'were' to get the names\n",
    "    names = sentence.split('and')[0].split()[-1]\n",
    "    all_B.append(\" \" + names)\n",
    "    \n",
    "    names = sentence.split('and')[1].split('were')[0].strip()\n",
    "    all_A.append(\" \" + names)\n",
    "\n",
    "\n",
    "print(all_A)\n",
    "print(all_B)\n",
    "sentence_token = model.to_tokens(ex_sent, prepend_bos=False)\n",
    "all_A_tokens = [model.to_single_token(A) for A in all_A]\n",
    "all_B_tokens = [model.to_single_token(B) for B in all_B]\n",
    "\n",
    "#2\n",
    "with torch.no_grad():\n",
    "    logits = model(sentence_token).log_softmax(-1)\n",
    "    A_logits = logits[torch.arange(len(all_A_tokens)), -1, all_A_tokens]\n",
    "    B_logits = logits[torch.arange(len(all_B_tokens)), -1, all_B_tokens]\n",
    "\n",
    "print('A_logit: ', A_logits)\n",
    "print('B_logit: ', B_logits)\n",
    "#import partial\n",
    "from functools import partial\n",
    "from tqdm import tqdm\n",
    "\n",
    "#3\n",
    "def mlp_ablation_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 = smaller_auto_encoder(int_val)\n",
    "\n",
    "    # Subtract value with feature direction*act_of_feature\n",
    "    feature_direction = torch.outer(act[:, feature_to_ablate].squeeze(), smaller_auto_encoder.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",
    "num_features = smaller_dict.shape[0]\n",
    "diff_A = np.zeros(num_features)\n",
    "diff_B = np.zeros(num_features)\n",
    "with torch.no_grad():\n",
    "    for feature_to_ablate in tqdm(range(num_features)):\n",
    "        ablated_logits = model.run_with_hooks(sentence_token, fwd_hooks=[(cache_name, partial(mlp_ablation_hook, feature_to_ablate=feature_to_ablate))])\n",
    "        ablated_logits = ablated_logits.log_softmax(-1)\n",
    "\n",
    "        A_ablated_logits = ablated_logits[torch.arange(len(all_A_tokens)), -1, all_A_tokens]\n",
    "        B_ablated_logits = ablated_logits[torch.arange(len(all_B_tokens)), -1, all_B_tokens]\n",
    "\n",
    "        diff_A[feature_to_ablate] = (A_logits - A_ablated_logits).mean().item()\n",
    "        diff_B[feature_to_ablate] = (B_logits - B_ablated_logits).mean().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_logits, B_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Topk of diff_A and diff_B\n",
    "topk = 10\n",
    "diff_A_tensor = torch.tensor(diff_A)\n",
    "diff_B_tensor = torch.tensor(diff_B)\n",
    "diff_A_topk = torch.topk(diff_A_tensor, topk)\n",
    "diff_B_topk = torch.topk(diff_B_tensor, topk)\n",
    "print(\"A we want lower:\",diff_A_topk)\n",
    "print(\"B we want Higher:\", diff_B_topk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Also do topk of diff_A_tensor, etc , largest=False\n",
    "#   612, 2769, 3806\n",
    "#   4786, 1152, 3750, 1753, 3383, 6035, 5915, 1413\n",
    "print(\"A we want lower:\",torch.topk(diff_A_tensor, 10, largest=False))\n",
    "print(\"B we want Higher:\",torch.topk(diff_B_tensor, 10, largest=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 50\n",
    "diff_A_tensor = torch.tensor(diff_A)\n",
    "diff_B_tensor = torch.tensor(diff_B)\n",
    "low_A = torch.topk(diff_A_tensor, k, largest=False).indices\n",
    "high_B = torch.topk(diff_B_tensor, k).indices\n",
    "feature_list2 = []\n",
    "for l_a in low_A:\n",
    "    if(l_a in high_B):\n",
    "        # print(l_a.item())\n",
    "        feature_list2.append(l_a.item())\n",
    "print(feature_list2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_A_tensor = torch.tensor(diff_A)\n",
    "diff_B_tensor = torch.tensor(diff_B)\n",
    "k = 50\n",
    "low_A = torch.topk(diff_A_tensor, k, largest=True).indices\n",
    "high_B = torch.topk(diff_B_tensor, k, largest=False).indices\n",
    "feature_list = []\n",
    "for l_a in low_A:\n",
    "    if(l_a in high_B):\n",
    "        feature_list.append(l_a.item())\n",
    "print(feature_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_feature = torch.tensor([ 2.4080e-02, 2.4359e-02, 2.5915e-03, 4.9120e-02, 1.2658e-01,\n",
    "         3.2408e-02, -2.5182e-02, 4.3675e-02, -5.4410e-02, 1.6136e-02,\n",
    "        -2.0048e-02, -2.2420e-02, 3.3210e-02, 3.4793e-02, -2.8834e-02,\n",
    "         7.4147e-02, 3.7202e-02, 5.3582e-02, 1.1046e-03, 1.2104e-02,\n",
    "         5.3141e-03, 1.5514e-02, -3.4517e-02, -1.2846e-02, -1.4674e-02,\n",
    "         8.6960e-03, -1.5789e-02, 9.3915e-04, 5.9387e-04, -6.1752e-02,\n",
    "         4.1706e-02, -6.9348e-02, 6.0245e-02, -1.5278e-02, 4.6114e-03,\n",
    "        -6.0599e-02, 4.6384e-02, 8.7855e-03, 9.8222e-02, -1.4912e-02,\n",
    "        -2.0680e-02, 8.0620e-04, 2.1009e-02, -1.8937e-02, 1.9966e-02,\n",
    "         1.8603e-02, 9.0589e-02, -5.7756e-03, -4.1959e-02, 5.0267e-03,\n",
    "         2.3697e-02, -5.8151e-02, -7.7805e-02, -4.9728e-02, -4.8614e-03,\n",
    "        -3.9631e-02, -4.0052e-02, 3.9487e-02, 6.3542e-03, 2.7683e-02,\n",
    "        -3.4077e-02, -2.9967e-02, 3.2453e-03, -8.4002e-02, -3.5077e-02,\n",
    "        -9.1061e-03, -8.3716e-04, 7.6448e-03, -7.9082e-03, 3.9118e-02,\n",
    "        -3.9306e-02, -2.0184e-02, -1.7360e-02, -9.5160e-03, 1.3632e-02,\n",
    "        -1.1137e-01, -3.2181e-02, -5.6427e-02, -2.3615e-02, 9.5962e-03,\n",
    "         5.5335e-02, -1.8662e-02, 3.4204e-02, -5.0086e-02, -7.4954e-02,\n",
    "         8.4810e-03, 5.8262e-02, 1.7019e-02, 3.1432e-03, 1.1476e-02,\n",
    "         6.3471e-03, -4.1524e-02, 1.8465e-02, 7.5723e-02, 1.5079e-02,\n",
    "         1.9857e-02, 1.1886e-02, 9.8358e-03, -5.3220e-02, 5.9651e-02,\n",
    "         2.2428e-02, 3.4533e-02, -1.7907e-02, 1.0397e-02, 4.1672e-02,\n",
    "        -1.8936e-02, 2.1871e-02, 2.5337e-03, 2.3077e-03, 3.9737e-02,\n",
    "        -3.9214e-02, 9.2625e-03, -3.8041e-02, -1.2577e-02, -3.1696e-03,\n",
    "        -2.7372e-02, -5.1618e-02, 4.4042e-02, 2.4820e-02, 2.7953e-02,\n",
    "         1.5850e-02, -1.6006e-03, 5.3238e-02, 9.8043e-03, 3.8326e-02,\n",
    "        -4.1484e-02, 3.6432e-02, 2.2653e-02, 1.8522e-02, -3.7113e-02,\n",
    "         1.7174e-02, -8.3303e-03, 2.8099e-03, 1.0931e-01, -4.4335e-02,\n",
    "         1.1882e-02, -6.4499e-03, 3.7061e-02, -9.4151e-02, -5.1065e-02,\n",
    "         2.7060e-02, -6.5621e-02, -3.0455e-02, -1.6427e-02, 3.9223e-02,\n",
    "        -3.9643e-02, 6.4269e-04, -1.2907e-03, -5.0777e-03, -5.1587e-02,\n",
    "        -3.4985e-02, 3.7856e-02, -9.9967e-02, -6.5395e-03, -4.6955e-02,\n",
    "         2.4562e-02, 1.5419e-02, -2.1492e-02, -8.9394e-03, 4.2825e-02,\n",
    "        -4.5467e-02, -1.2693e-02, -4.9931e-02, 1.9798e-03, -3.1977e-02,\n",
    "         6.7696e-03, -3.7578e-02, -3.1094e-02, -3.4964e-02, 3.0003e-02,\n",
    "         3.0392e-02, 2.8398e-02, 7.5477e-04, 2.7859e-02, 8.7245e-03,\n",
    "         1.9231e-02, -3.1447e-02, -3.8160e-02, 3.6587e-02, -1.2608e-02,\n",
    "         6.7020e-02, 9.5656e-02, 2.0806e-02, -1.0801e-02, -1.6442e-02,\n",
    "         2.9674e-02, -4.6152e-02, -6.7173e-02, 5.4613e-02, -7.5615e-03,\n",
    "        -1.2474e-05, -9.5788e-03, 6.4168e-03, 1.7756e-02, -1.0295e-02,\n",
    "        -2.1090e-02, -3.4596e-02, -1.4731e-02, 1.0119e-02, 5.8010e-03,\n",
    "        -5.4254e-03, -4.6099e-02, 6.3387e-02, -1.7798e-02, -3.1989e-02,\n",
    "        -2.7206e-02, -7.4612e-02, 3.3296e-02, 9.3995e-03, 3.7320e-02,\n",
    "        -8.7339e-02, 2.0072e-02, -3.8689e-02, -9.5642e-03, -2.8851e-03,\n",
    "         3.4954e-03, -3.7127e-03, -3.8476e-03, 1.9904e-02, -4.4326e-03,\n",
    "         2.4047e-02, 2.1409e-02, -5.3637e-04, -1.3460e-02, -2.5007e-02,\n",
    "        -3.6522e-02, -6.4226e-02, 6.0937e-03, -4.9409e-03, -5.5151e-02,\n",
    "         1.1245e-02, -5.9511e-02, 4.5356e-02, 2.3629e-02, -3.4198e-03,\n",
    "         2.6451e-02, 5.8297e-02, -1.5115e-02, 3.8943e-02, -1.3371e-02,\n",
    "        -2.6044e-02, -5.6735e-02, 6.8021e-02, -3.7823e-02, 7.3421e-02,\n",
    "        -2.4523e-02, -6.3215e-03, -8.1932e-03, 5.6012e-02, 4.7597e-04,\n",
    "        -4.0288e-02, -1.3732e-03, -3.2098e-02, -9.2334e-03, 5.5161e-03,\n",
    "         2.5030e-02, 1.5095e-02, -2.9644e-02, 5.2908e-02, -2.3363e-02,\n",
    "         5.3523e-03, 3.9384e-03, 2.1839e-02, 1.2031e-02, 2.4092e-02,\n",
    "         2.3921e-02, -7.7738e-02, -1.1835e-01, -5.1954e-02, -2.1009e-02,\n",
    "         2.5704e-02, -2.4611e-02, 3.3695e-02, 2.6925e-02, 2.6932e-02,\n",
    "         1.2319e-02, -8.9244e-02, -8.6162e-03, 3.5292e-02, -3.8614e-02,\n",
    "        -3.9250e-02, 1.6301e-02, -3.1443e-02, 1.9028e-03, 5.7968e-02,\n",
    "         3.9234e-03, -2.0115e-02, -7.6918e-03, 4.2788e-02, 3.0034e-02,\n",
    "         9.1978e-02, -5.9110e-02, 2.3647e-02, 2.0807e-02, -1.2634e-02,\n",
    "        -7.6869e-03, 4.4716e-02, -1.8892e-02, 2.7425e-02, -9.9979e-03,\n",
    "         1.9796e-02, 1.0973e-02, -1.6451e-03, -4.3213e-02, -3.5061e-03,\n",
    "         2.9969e-03, -7.1300e-02, 6.3207e-03, 8.0630e-03, -3.5077e-02,\n",
    "        -9.7324e-03, -5.2640e-02, -5.5317e-02, -2.2824e-02, 2.7998e-02,\n",
    "        -7.8537e-02, 4.3458e-02, -1.4984e-02, -7.7363e-02, 3.7290e-02,\n",
    "         7.7163e-02, -3.3954e-03, 1.3120e-02, 5.7814e-02, -1.7183e-02,\n",
    "         2.4303e-03, 3.1352e-02, 7.0760e-02, -5.9879e-03, 6.1438e-02,\n",
    "        -9.8398e-03, 1.6338e-02, 2.8400e-02, 5.7705e-02, -2.2048e-03,\n",
    "        -4.5825e-03, -1.2012e-02, -1.3619e-03, 3.8692e-02, -1.8441e-02,\n",
    "        -1.1080e-02, 3.3503e-03, 1.1516e-02, 5.0336e-02, 2.0678e-03,\n",
    "        -2.6509e-02, 4.7976e-03, -4.7704e-03, -3.4007e-02, -5.1153e-02,\n",
    "        -1.9541e-02, -1.5282e-02, -3.3280e-02, 2.1600e-02, -6.8638e-02,\n",
    "         7.2270e-03, -4.8205e-02, 6.4255e-02, -3.0913e-02, 3.3795e-02,\n",
    "         3.1021e-02, 1.7689e-02, 7.8526e-03, 4.4164e-02, -2.6782e-03,\n",
    "         2.9056e-02, 3.4432e-02, 3.8272e-02, 7.4780e-03, 5.0119e-02,\n",
    "        -1.8974e-02, -2.9729e-02, 4.4389e-02, -2.5922e-02, -6.9515e-02,\n",
    "        -6.1480e-02, -1.2153e-02, -5.7087e-02, 8.2086e-02, 3.9201e-02,\n",
    "         2.0275e-02, 8.8203e-03, -4.9722e-02, 6.9982e-02, -7.2797e-03,\n",
    "         4.0989e-02, 4.1517e-03, 8.9593e-03, 2.5157e-02, -4.8101e-03,\n",
    "        -9.6375e-03, -2.2509e-02, 5.1393e-02, 7.2553e-03, 1.3507e-02,\n",
    "         4.3849e-02, 6.2116e-02, 3.8052e-02, 1.5074e-02, 1.4775e-02,\n",
    "         2.6600e-02, 2.7580e-02, -6.2622e-03, -2.5771e-02, 1.7867e-02,\n",
    "        -1.4621e-02, 1.3103e-02, 1.5070e-02, -4.5592e-02, 6.9849e-06,\n",
    "         5.1731e-03, -2.9146e-02, -1.2707e-02, 2.8848e-02, 8.5701e-02,\n",
    "        -1.1278e-02, 2.7468e-02, -4.3419e-03, -2.8725e-02, -3.2829e-02,\n",
    "        -1.2745e-02, 3.3975e-02, -1.1140e-02, -3.2737e-02, -3.7541e-02,\n",
    "        -1.2454e-02, -1.5200e-02, -2.8173e-02, -1.0102e-01, -1.6749e-02,\n",
    "         2.8634e-02, -1.3546e-02, -1.8982e-02, 9.3449e-02, 4.1330e-02,\n",
    "        -1.3391e-02, -2.9870e-02, -1.0827e-02, -4.5005e-02, 1.5958e-02,\n",
    "        -2.7183e-03, -3.2368e-02, 5.6117e-02, -4.7420e-02, -2.1395e-02,\n",
    "        -4.4957e-02, -5.9798e-03, 2.2991e-02, 7.0553e-02, -9.5782e-03,\n",
    "         9.5198e-03, 5.8040e-03, -3.0258e-02, -5.3331e-02, 8.5738e-03,\n",
    "         5.9829e-02, 4.9407e-02, -3.9949e-02, -1.5012e-02, -5.7299e-02,\n",
    "         3.3791e-02, 2.6755e-02, 1.6138e-02, 1.5216e-02, -1.4221e-02,\n",
    "        -5.5148e-02, -6.4813e-03, -4.5133e-02, 2.5805e-02, 1.9024e-02,\n",
    "         2.3914e-03, 5.3716e-03, -5.0496e-02, -3.6670e-04, -2.6395e-02,\n",
    "         5.3805e-02, -1.9281e-02, 2.9166e-02, 1.8902e-02, 6.9283e-03,\n",
    "        -8.6916e-03, 1.3865e-01, 4.3540e-02, 7.9599e-03, -1.1489e-02,\n",
    "        -1.9561e-02, -3.3223e-02, 4.2447e-02, 2.4892e-02, 7.6215e-03,\n",
    "        -5.5037e-02, -4.0666e-03, -4.7167e-03, 3.1368e-02, 2.2559e-02,\n",
    "        -9.4491e-03, -3.7170e-02, -4.0265e-02, -2.2772e-04, -6.9743e-03,\n",
    "        -7.5886e-03, 3.6595e-02, -1.0402e-02, -3.0940e-02, -1.4825e-02,\n",
    "         1.4237e-02, 4.5008e-02, 2.5084e-02, 1.0162e-02, 2.5703e-02,\n",
    "        -1.3544e-02, -1.1389e-01, -4.0357e-02, -4.7644e-02, -3.7380e-02,\n",
    "        -1.4030e-02, 5.7597e-02, 1.7919e-04, 1.8142e-02, -1.6253e-02,\n",
    "        -7.6142e-02, 3.8851e-02, -2.1893e-02, -1.5224e-02, 2.1425e-02,\n",
    "         1.6085e-02, -2.6741e-02, -2.4140e-02, -1.4206e-02, 4.1798e-02,\n",
    "        -3.2894e-02, 4.9439e-02, -9.9219e-03, -7.2214e-02, -4.4329e-03,\n",
    "         2.9387e-02, -3.9025e-02, 1.9442e-02, 3.4890e-02, 1.3546e-02,\n",
    "         6.8964e-04, -2.7797e-02, -6.8082e-02, 9.9170e-03, 2.3642e-02,\n",
    "        -6.8385e-02, 1.0770e-02, -5.6083e-03, 1.7623e-02, -3.4506e-02,\n",
    "        -4.2243e-02, 7.4690e-02, -2.4971e-02, -5.2503e-02, -1.4595e-02,\n",
    "         5.1033e-02, 6.1160e-02, 3.7497e-03, -3.8879e-02, -4.1647e-03,\n",
    "         2.6918e-02, -1.4865e-02, 9.0000e-03, -1.4608e-02, -2.1336e-02,\n",
    "        -5.9530e-03, -1.6575e-03, -4.7094e-02, 2.3869e-02, 1.4982e-02,\n",
    "         4.9038e-02, -4.5014e-02, 4.8011e-02, -1.5510e-02, -1.6457e-02,\n",
    "         1.2145e-02, 4.9655e-03, -4.4219e-02, 2.6135e-02, 1.6504e-02,\n",
    "         3.7103e-03, 3.4081e-02, -8.6893e-03, -9.9807e-03, 1.2927e-02,\n",
    "         3.5594e-02, 2.4592e-02, 3.7850e-02, -1.1348e-02, -1.1096e-03,\n",
    "        -4.4439e-02, 2.7691e-02, -6.9026e-02, 5.7734e-02, 2.9289e-02,\n",
    "         3.4821e-02, -5.2970e-03, 3.7610e-02, -1.7178e-02, -5.0140e-02,\n",
    "        -8.3218e-02, 4.2783e-02, -1.2670e-02, 1.1850e-02, 3.1992e-02,\n",
    "        -4.4622e-04, 3.7421e-02, -2.3873e-02, -1.9229e-02, 1.4420e-02,\n",
    "         1.5459e-02, 2.4360e-02, -1.4355e-03, 3.8322e-02, -4.1234e-02,\n",
    "        -4.1331e-02, -2.9026e-02, -1.7661e-02, 5.0412e-02, -5.7450e-02,\n",
    "        -2.2180e-02, 1.4419e-02, -1.6419e-02, -2.6536e-02, -6.5096e-02,\n",
    "        -8.5083e-03, -5.8017e-03, 4.6651e-02, 3.4101e-02, -1.8907e-02,\n",
    "         2.7087e-02, 2.5500e-02, 1.1684e-03, 6.8219e-02, -3.2993e-02,\n",
    "        -2.3292e-02, 1.5120e-03, -2.1995e-02, 6.2556e-02, 4.9651e-02,\n",
    "         1.6080e-02, -1.3513e-02, 3.5774e-02, -4.4417e-02, 4.9697e-03,\n",
    "         4.9093e-02, 1.3976e-02, 2.6455e-02, 3.7071e-02, -4.3469e-02,\n",
    "         1.3540e-02, 5.2918e-02, 3.4559e-02, -3.1488e-02, -1.9801e-02,\n",
    "        -3.4259e-02, -1.0953e-02, -3.2256e-02, -3.9682e-02, 9.6360e-03,\n",
    "        -9.6176e-03, 1.2003e-03, 4.2134e-02, -5.7197e-02, -4.2081e-02,\n",
    "        -4.6472e-03, 2.6133e-02, 3.7587e-02, 9.8670e-03, -3.4046e-02,\n",
    "        -3.2811e-03, -1.5316e-02, -4.8647e-03, 1.9632e-02, 1.1838e-02,\n",
    "         1.5452e-02, -6.3881e-02, 1.2259e-02, 1.2232e-03, -4.3000e-02,\n",
    "        -6.8827e-02, 4.1810e-02, 7.1543e-03, -5.1488e-02, 2.4518e-02,\n",
    "        -1.0627e-03, 1.4221e-02, -5.3517e-02, -4.3952e-02, -3.0068e-02,\n",
    "        -1.9294e-03, -1.0837e-02, -1.0782e-02, 1.1658e-02, -7.7354e-03,\n",
    "        -9.4467e-03, 5.8047e-02, 6.2752e-03, 1.7554e-02, 5.2714e-04,\n",
    "         1.2364e-02, 4.7268e-02, -3.3143e-02, -2.9131e-02, -7.5221e-03,\n",
    "        -1.3250e-02, -1.0116e-02, 1.5731e-02, 1.2194e-03, 5.5533e-02,\n",
    "        -2.4715e-02, 3.5104e-04, 7.7462e-02, -4.2483e-02, 9.3196e-03,\n",
    "         3.5436e-04, 3.0903e-02, -6.8819e-02, 1.9523e-02, 2.2681e-04,\n",
    "        -2.1906e-02, -3.0358e-03, 1.9632e-02, 1.0104e-02, -4.5632e-03,\n",
    "         2.3078e-02, -2.9149e-02, -4.9732e-02, 1.2652e-05, 4.3399e-03,\n",
    "         3.1472e-02, 1.0199e-02, -3.4592e-04, -3.0710e-02, 8.5242e-03,\n",
    "        -6.5155e-03, -3.7067e-02, -6.3603e-03, -1.0648e-02, 4.0477e-02,\n",
    "        -1.5365e-02, 2.0180e-02, 8.3195e-03, 1.8911e-02, -3.3727e-02,\n",
    "        -3.8646e-02, -3.2142e-02, -1.6740e-02, 2.0304e-02, 3.7467e-02,\n",
    "         2.5757e-02, -9.1715e-02, 6.6637e-02, -2.6891e-02, 3.7948e-02,\n",
    "        -2.8087e-02, -1.2200e-02, 2.7008e-02, -6.5081e-03, 7.5711e-02,\n",
    "         1.4658e-02, -1.6642e-02, 4.5490e-02, 1.1131e-02, 4.4595e-02,\n",
    "        -5.6698e-02, 1.6837e-02, -7.4032e-03])\n",
    "for feat in feature_list:\n",
    "    # Calculate CS\n",
    "    feature_direction = smaller_dict[feat]\n",
    "    cs = torch.nn.functional.cosine_similarity(feature_direction, original_feature, dim=0)\n",
    "    # cs = torch.nn.functional.cosine_similarity(feature_direction, smaller_dict[feature_list[3]], dim=0)\n",
    "    print(feat, cs.item())\n",
    "# Calculate cs between two random features of size 768\n",
    "rand_dir_1 = torch.randn(768)\n",
    "rand_dir_2 = torch.randn(768)\n",
    "cs = torch.nn.functional.cosine_similarity(rand_dir_1, rand_dir_2, dim=0)\n",
    "print(cs.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_feature_2 = [-2.5726e-02, -6.4886e-04, -5.1383e-04, -2.0602e-03, -1.1533e-01,\n",
    "        -5.2642e-02,  4.8039e-03, -2.6216e-02,  6.1887e-02, -1.2594e-02,\n",
    "        -1.0364e-02,  4.2911e-02, -2.3051e-02, -8.9775e-03,  3.9231e-02,\n",
    "        -8.0026e-02, -3.2398e-02, -4.7174e-02, -2.8013e-02, -5.5694e-03,\n",
    "        -1.0717e-02,  2.3064e-02,  2.9110e-02,  1.0392e-02,  1.4949e-02,\n",
    "        -3.3139e-02,  9.7329e-03,  1.4964e-02,  9.0917e-03,  3.5896e-02,\n",
    "        -3.5559e-02,  6.7633e-02, -5.7505e-02,  1.5716e-02, -1.7437e-02,\n",
    "         5.7396e-02, -4.3503e-02, -8.4819e-03, -6.5898e-02, -7.8586e-03,\n",
    "         2.7507e-02, -2.1453e-02, -4.0811e-03,  4.9995e-03, -4.0651e-02,\n",
    "        -5.1222e-02, -4.9716e-02,  6.1054e-03,  5.4668e-02,  3.9923e-03,\n",
    "        -1.8735e-02,  7.3003e-02,  6.3771e-02,  3.6127e-02,  1.0396e-02,\n",
    "         3.5287e-02,  1.6697e-02, -3.6990e-02, -3.2620e-02, -3.7428e-02,\n",
    "         2.5951e-02,  3.2467e-02,  3.5407e-04,  5.8685e-02,  3.6463e-02,\n",
    "         7.8591e-03,  2.4923e-03,  1.2690e-02,  2.6344e-03, -3.4273e-02,\n",
    "         3.9547e-02,  2.4529e-02,  2.6538e-02, -1.8732e-02, -1.7933e-02,\n",
    "         1.1706e-01,  7.3611e-03,  5.7754e-02,  5.2099e-02, -2.1565e-02,\n",
    "        -6.3068e-02,  5.8480e-03, -4.4039e-02,  3.9488e-02,  6.8688e-02,\n",
    "        -8.3774e-04, -6.2737e-02, -2.1059e-02, -3.9884e-02, -2.7339e-02,\n",
    "         2.3889e-03,  5.1001e-02, -3.7041e-02, -7.8818e-02, -7.8122e-03,\n",
    "        -6.5944e-03, -3.0724e-02, -2.1623e-03,  4.0961e-02, -6.5949e-02,\n",
    "        -1.2377e-02, -5.4195e-02,  3.4806e-02, -5.7749e-03, -2.7602e-02,\n",
    "         2.8398e-02, -2.5806e-02, -4.9847e-03, -1.3046e-02, -5.0813e-02,\n",
    "         3.8295e-02, -1.8783e-02,  3.5261e-02,  3.0712e-03,  3.4825e-02,\n",
    "         1.5244e-02,  2.7976e-02, -4.9202e-02, -3.1048e-02, -1.2737e-02,\n",
    "        -1.4603e-02,  1.2811e-02, -5.5029e-02, -2.3307e-02, -3.5287e-02,\n",
    "         6.5566e-02, -3.1492e-02, -1.1382e-02, -1.3777e-02,  5.6112e-02,\n",
    "        -2.3837e-02, -6.7690e-05,  4.8088e-03, -1.0323e-01,  3.7318e-02,\n",
    "        -3.1996e-03, -1.9138e-02, -1.7097e-02,  6.7230e-02,  2.3842e-02,\n",
    "        -4.3859e-02,  5.9642e-02,  3.0886e-02,  1.8266e-02, -5.4966e-02,\n",
    "         4.2181e-02,  2.4301e-03,  3.7863e-03,  1.4051e-02,  4.8524e-02,\n",
    "         5.1219e-02, -2.2420e-02,  8.2472e-02,  2.1550e-02,  4.3147e-02,\n",
    "        -1.0695e-02, -1.0234e-02,  1.4201e-02,  1.5793e-02, -2.9356e-02,\n",
    "         3.2282e-02,  2.3433e-02,  4.2362e-02, -8.9699e-04,  8.8142e-03,\n",
    "        -1.1446e-02,  4.0446e-02,  1.1830e-02,  4.1430e-02, -3.8639e-02,\n",
    "        -9.0949e-03, -1.8916e-02,  1.5558e-02, -4.4642e-02, -7.8419e-03,\n",
    "        -1.9711e-02,  4.5687e-02,  1.2629e-02, -3.8284e-02,  9.4555e-03,\n",
    "        -6.1299e-02, -9.5572e-02, -1.5757e-02,  1.7990e-02,  2.5031e-02,\n",
    "        -2.9053e-02,  1.8828e-02,  4.9860e-02, -4.5340e-02, -1.4800e-02,\n",
    "         1.3446e-03, -1.2208e-02, -1.6560e-02, -4.1575e-03,  7.8791e-03,\n",
    "         2.9964e-02,  3.5376e-02,  1.0379e-02, -1.3305e-02, -6.4321e-03,\n",
    "        -8.0099e-04,  3.6716e-02, -5.0114e-02, -1.3348e-03,  2.8296e-02,\n",
    "         2.8295e-02,  8.1506e-02, -4.5367e-02, -7.9561e-03, -3.3327e-02,\n",
    "         7.6064e-02, -1.2496e-02,  1.5921e-02,  1.6194e-02,  7.7633e-03,\n",
    "         2.0605e-02,  1.5549e-02, -7.1915e-04, -2.5339e-02,  5.6805e-03,\n",
    "        -1.2525e-02, -3.3941e-02, -4.7678e-02,  1.9785e-02,  3.4760e-02,\n",
    "         4.5086e-02,  5.0492e-02, -2.0813e-02, -2.0093e-02,  7.7106e-02,\n",
    "        -1.9640e-02,  6.7751e-02, -3.7802e-02, -2.1798e-02, -1.6374e-02,\n",
    "        -3.7402e-02, -4.7581e-02,  8.5529e-03, -5.9982e-02,  1.3218e-03,\n",
    "         4.7270e-02,  4.8120e-02, -4.6423e-02,  2.8720e-02, -5.6258e-02,\n",
    "         4.2057e-02, -2.7699e-03,  2.1331e-02, -4.0283e-02, -1.4520e-02,\n",
    "         2.1053e-02, -1.0795e-03,  3.2510e-02, -4.3478e-03, -3.2881e-03,\n",
    "        -6.1366e-02,  1.5677e-02,  4.0255e-02, -7.7206e-02,  2.9707e-02,\n",
    "        -1.8393e-02,  2.1846e-03, -2.1147e-02, -3.8249e-02, -2.9695e-02,\n",
    "        -3.1903e-02,  7.6872e-02,  1.0846e-01,  5.4730e-02,  3.1813e-03,\n",
    "        -1.5633e-03,  1.9181e-03, -4.9521e-02, -1.5803e-02, -2.4915e-02,\n",
    "        -1.7885e-02,  6.0699e-02, -3.8426e-02, -5.1543e-02,  2.3449e-02,\n",
    "         5.8460e-02, -2.9590e-02,  3.2966e-02,  1.0397e-03, -2.6838e-02,\n",
    "         5.0908e-03,  4.1476e-02,  2.5819e-02, -3.9116e-02, -2.2437e-02,\n",
    "        -8.0938e-02,  6.1203e-02, -3.0964e-02, -3.0232e-02,  2.5044e-02,\n",
    "        -1.0622e-02, -3.2421e-02,  5.8003e-03, -3.1614e-02,  2.3944e-02,\n",
    "        -2.4293e-04, -3.3807e-02,  1.3699e-02,  3.4272e-02,  1.1896e-02,\n",
    "        -8.9490e-03,  8.2719e-02,  1.3622e-02,  1.6564e-03,  5.6063e-02,\n",
    "        -9.6542e-03,  4.0601e-02,  4.2192e-02,  2.7282e-02, -4.0123e-02,\n",
    "         8.1882e-02, -2.6121e-02,  2.3286e-02,  4.9314e-02, -2.8606e-03,\n",
    "        -7.4387e-02, -5.1504e-03,  6.0365e-04, -5.9640e-02,  1.2083e-02,\n",
    "        -2.4037e-02, -4.7537e-02, -8.9067e-02,  1.5172e-02, -7.0192e-02,\n",
    "         1.1639e-02, -4.6833e-03, -3.5606e-02, -2.1418e-02, -8.3400e-03,\n",
    "        -1.1812e-02,  8.8803e-04,  1.5831e-03, -3.0439e-02,  5.4649e-03,\n",
    "         6.3953e-03, -1.4518e-02, -1.8259e-03, -5.2184e-02, -2.1835e-02,\n",
    "         5.8973e-03, -8.5032e-03,  2.1707e-02,  4.8014e-02,  4.7849e-02,\n",
    "         2.3958e-03,  1.2922e-02,  4.2764e-02, -2.6646e-02,  7.2119e-02,\n",
    "        -2.0728e-02,  5.8134e-02, -5.1347e-02,  5.2139e-02, -5.4106e-02,\n",
    "        -3.1487e-02, -4.2379e-03, -1.1163e-02, -3.4577e-02,  2.0586e-02,\n",
    "         1.2625e-04, -3.5940e-02, -6.4555e-02, -4.3353e-02, -5.8778e-02,\n",
    "         1.5517e-02, -1.0910e-03, -6.6798e-02, -2.7699e-02,  7.8522e-02,\n",
    "         8.6146e-02,  7.7630e-03,  5.0124e-02, -5.7992e-02, -2.2200e-02,\n",
    "        -4.3185e-03, -2.8709e-03,  3.9246e-02, -8.1500e-02, -1.0423e-02,\n",
    "        -5.7764e-02, -5.9561e-03, -5.1224e-03, -2.0711e-02, -3.7797e-03,\n",
    "         3.2807e-03,  2.6843e-02, -3.6737e-02, -3.0800e-02, -1.0503e-02,\n",
    "        -5.3457e-02, -4.1221e-02, -3.3799e-02, -8.6349e-03, -4.7439e-03,\n",
    "        -2.4835e-02, -5.2798e-02,  9.5092e-03,  3.1724e-02, -9.0052e-03,\n",
    "         7.5382e-03,  3.8200e-03, -7.7856e-03,  2.9869e-02, -4.8750e-03,\n",
    "        -7.8292e-03,  2.3758e-02,  1.4873e-02, -9.4694e-03, -4.6604e-02,\n",
    "        -6.2212e-03, -2.0669e-02,  1.2701e-02,  1.7831e-02,  3.5741e-02,\n",
    "        -3.8615e-03, -2.3966e-02,  2.1824e-03,  2.3840e-02,  4.9500e-02,\n",
    "         1.8837e-02, -5.2920e-05,  2.6114e-02,  1.0073e-01,  1.9235e-02,\n",
    "        -1.7241e-02,  3.3887e-02,  2.9006e-02, -9.1689e-02, -3.1656e-02,\n",
    "        -1.1117e-02,  4.8628e-02,  2.7492e-02,  3.9009e-02, -3.0926e-02,\n",
    "        -4.2895e-03,  3.5016e-02, -3.9638e-02,  5.9491e-02,  8.9960e-03,\n",
    "         3.0452e-02, -1.2911e-03,  1.0351e-02, -5.9746e-02, -1.9469e-03,\n",
    "        -2.3070e-02, -4.6439e-03,  1.4491e-02,  2.8515e-02, -2.3155e-02,\n",
    "        -2.7369e-02, -4.1841e-02,  4.8859e-02, -1.4076e-02,  5.4690e-02,\n",
    "        -4.4104e-02, -2.7567e-02, -8.5394e-03,  2.0238e-02,  3.1467e-02,\n",
    "         7.9279e-02, -3.3372e-03,  3.4812e-02, -3.9662e-02, -2.3040e-02,\n",
    "        -8.3274e-03,  3.8463e-03,  6.4554e-02,  1.8813e-02,  2.2556e-02,\n",
    "        -6.8544e-02,  2.1708e-02, -4.1814e-02, -3.0375e-02, -1.5941e-02,\n",
    "         1.7288e-04, -2.9641e-02, -3.3388e-02,  2.9383e-02, -7.4426e-03,\n",
    "         1.2448e-02,  3.7900e-02, -6.3293e-02, -2.8445e-02,  4.9981e-03,\n",
    "         5.2990e-02, -1.4166e-03,  3.9616e-02, -6.4212e-03, -2.0357e-02,\n",
    "        -1.9757e-03,  5.4820e-02,  5.0063e-02,  3.7313e-03,  1.3351e-02,\n",
    "        -8.3974e-03, -4.8734e-02,  3.1723e-02,  4.3567e-02,  3.4318e-02,\n",
    "         3.6629e-04, -3.9988e-02, -1.3103e-02, -1.5312e-02, -1.7292e-02,\n",
    "         3.1800e-03,  9.6552e-02,  2.7071e-02,  6.1272e-02,  3.2394e-02,\n",
    "         1.4733e-02, -7.6506e-02, -1.1652e-02, -3.4143e-02,  8.1766e-03,\n",
    "         4.8835e-02, -1.5100e-02,  1.8776e-02,  4.3010e-02, -2.3001e-02,\n",
    "         8.3834e-03,  2.6201e-02,  3.0167e-02, -5.6585e-03, -2.6402e-02,\n",
    "         2.6883e-02, -5.3909e-02, -4.0254e-03,  6.2743e-02, -7.9913e-04,\n",
    "        -3.7325e-02,  6.0301e-02, -2.7798e-02, -3.3212e-02,  1.3380e-02,\n",
    "         8.1164e-03,  2.3546e-02,  2.8834e-02, -1.3174e-02, -2.3931e-02,\n",
    "         4.8266e-02, -8.1734e-03,  7.8013e-03, -2.5567e-02,  2.6538e-02,\n",
    "         3.7964e-02, -9.0547e-02,  3.2511e-02,  4.3442e-02,  2.7266e-02,\n",
    "        -6.2117e-02, -7.3943e-02,  3.6751e-04,  4.8327e-02,  4.1736e-02,\n",
    "        -2.4926e-02,  2.3507e-03, -2.5331e-02,  2.5488e-02,  1.6013e-02,\n",
    "         2.1877e-02, -2.1653e-02,  2.1313e-02, -1.2351e-02, -1.2053e-02,\n",
    "        -5.0269e-02,  4.3835e-02, -2.0020e-02,  2.4839e-03,  2.5738e-02,\n",
    "        -5.5224e-03, -5.8843e-03,  5.7469e-02, -1.6048e-02, -3.1025e-02,\n",
    "         2.1149e-02, -4.8876e-02, -2.8336e-02,  2.5896e-02,  6.3234e-03,\n",
    "        -3.7355e-02, -8.4070e-03, -4.3917e-02,  2.3865e-02,  7.2084e-03,\n",
    "         5.4681e-02, -1.8412e-02,  4.4411e-02, -6.0587e-02, -3.7133e-02,\n",
    "        -2.3951e-02, -9.2139e-04, -4.0388e-02,  1.4934e-02,  5.0151e-02,\n",
    "         5.7851e-02, -1.0180e-02, -6.8106e-03,  5.6417e-04, -1.7309e-02,\n",
    "        -1.0165e-03, -2.9528e-02,  4.1398e-02,  1.6000e-02, -2.4035e-02,\n",
    "        -2.0178e-04, -2.7096e-02, -1.1785e-02, -5.6973e-02,  2.4596e-02,\n",
    "         2.5844e-02,  3.5915e-02,  2.1663e-02, -5.6594e-02,  7.3724e-02,\n",
    "         2.6818e-02, -2.1379e-02,  7.2263e-04,  3.2312e-02,  5.1703e-02,\n",
    "         6.0316e-03,  1.5219e-02, -5.1856e-02, -3.6642e-02,  2.4889e-02,\n",
    "        -3.0064e-02, -1.4434e-02,  1.3706e-02, -8.8471e-02,  1.5549e-02,\n",
    "         1.0244e-02,  1.1725e-02,  2.1740e-02, -4.2962e-02, -6.0788e-02,\n",
    "        -2.9457e-02,  1.5662e-02, -1.5633e-02,  8.1673e-02, -2.5786e-03,\n",
    "        -5.7903e-02,  1.7249e-03, -5.6719e-02, -3.6636e-02,  3.4668e-02,\n",
    "        -2.4915e-02, -3.2883e-02, -3.5396e-02,  2.5812e-02,  2.2283e-02,\n",
    "         2.9422e-02,  1.9445e-02,  2.8004e-02,  5.4516e-02, -1.4896e-02,\n",
    "         2.1365e-02, -4.3958e-03, -4.8943e-02,  4.0144e-02,  3.8219e-02,\n",
    "        -7.8286e-04, -4.1919e-02, -4.6140e-02, -1.1789e-02,  1.0216e-02,\n",
    "         2.0828e-03,  2.9919e-02,  8.3159e-03, -8.5125e-03, -8.1263e-03,\n",
    "        -2.0636e-02,  6.6536e-02, -6.5788e-03, -1.0293e-02,  2.8445e-02,\n",
    "         5.2921e-02, -5.5562e-02, -1.7273e-02,  7.5576e-02, -4.2413e-02,\n",
    "         2.7888e-02, -3.1687e-02,  3.3248e-02,  6.2094e-02, -4.7034e-03,\n",
    "        -1.7305e-02,  1.5567e-02, -3.3363e-03, -1.3980e-02,  3.3695e-02,\n",
    "         1.9907e-02, -6.3426e-02, -2.1559e-02,  5.6458e-03,  1.0381e-02,\n",
    "        -1.7983e-02, -6.4998e-02,  3.9473e-02,  1.6300e-02,  1.9143e-02,\n",
    "         1.0056e-02,  3.0917e-03, -1.3239e-02,  4.7693e-03, -6.5798e-02,\n",
    "         2.6512e-02, -1.3273e-02, -6.7242e-02,  2.7019e-02,  6.2950e-03,\n",
    "        -1.7237e-02, -3.8856e-02,  6.9744e-02, -2.7839e-02, -1.1612e-02,\n",
    "         2.5963e-04, -4.9502e-03, -2.7478e-02, -2.4752e-02,  2.2492e-02,\n",
    "        -4.3067e-02,  6.9937e-03,  5.5968e-02,  2.0326e-02, -1.6968e-02,\n",
    "        -1.9992e-02, -1.0691e-02,  2.8936e-02,  4.6730e-02, -2.8219e-03,\n",
    "         1.9446e-02,  2.0802e-02,  1.2312e-02,  1.4868e-02, -4.9224e-02,\n",
    "         7.9079e-03, -2.0139e-02, -6.3683e-03, -3.9114e-03,  2.8311e-02,\n",
    "         3.4844e-02,  3.7943e-02,  8.6546e-03, -1.7897e-02, -2.7561e-02,\n",
    "        -2.7213e-02,  8.1288e-02, -4.9007e-02,  2.4349e-02, -5.4233e-02,\n",
    "         3.3028e-02, -7.7637e-03, -3.3826e-02,  1.7515e-02, -7.4451e-02,\n",
    "        -5.0871e-03,  1.7512e-02, -6.9451e-02, -1.2467e-02, -4.8812e-02,\n",
    "         7.1086e-02,  2.0196e-02,  3.5404e-03]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_cs = torch.nn.functional.cosine_similarity(smaller_dict, original_feature, dim=1)\n",
    "all_cs.topk(10), all_cs.topk(10, largest=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_cs = torch.nn.functional.cosine_similarity(smaller_dict, torch.tensor(original_feature_2), dim=1)\n",
    "all_cs.topk(10), all_cs.topk(10, largest=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_direction.shape, original_feature.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 3527 affects the B word negatively\n",
    "diff_A_tensor[3527], diff_B_tensor[3527]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lll = [1734,  998, 6035, 4786, 1152, 3383, 4225, 6102, 1538, 5524]\n",
    "for iii in range(10):\n",
    "    print(lll[iii])\n",
    "    print(diff_A_tensor[lll[iii]], diff_B_tensor[lll[iii]])\n",
    "# maybe 6035?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_features = smaller_dict.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Text\n",
    "You can use the functions below to find interesting features to then add here to \"feature\""
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "# Feature Search\n",
    "Type in a sentence & see which features activate \n",
    "Note: Some features may be outliers, which will typically show up as high activations for the first word & first period or \\n (or high positive bias)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = \"いさんさん��に\"\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": "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": [
    "# l1 = 8e-4, first 8 features don't make sense\n",
    "# 1-3k mostly makes sense. Degrades after 3k (.6 MCS)\n",
    "# 1505, Japanese words\n",
    "\n",
    "# l1 = 4e-4. Loads of single letter features, unsure how legit they are\n",
    "# 127 - induction thing\n",
    "# 1001 - Japanese language, predicts common japanese endings & grammar (their comma char & quotations)\n",
    "# 1012 ? maybe poly, maybe context neuron\n",
    "# 1500: num after +/-\n",
    "# 1501: food-related\n",
    "# 1505 Poly\n",
    "# 2506 Poly\n",
    "# l1 = 2e-4. Definitely some real features, but many many polysemantic ones\n",
    "N = 3004\n",
    "best_feature = int(max_indices[N])\n",
    "# Feature_ids [1644, 3277, 1696, 1963, 1809, 2332, 3235, 93, 2140, 601]\n",
    "ind_for_feature=0\n",
    "indices=[2270, 1000, 4809, 2865, 4236,  612, 2769, 2188, 3806, 3499]\n",
    "# indices=[4786,  998, 1152, 3750, 1753, 3383, 6035, 4225, 5915, 1413]\n",
    "best_feature = 3383\n",
    "\n",
    "# 3383\n",
    "# 6035\n",
    "# 1413\n",
    "# 4466\n",
    "# 3806\n",
    "# 3667\n",
    "# 5455\n",
    "# 4290\n",
    "# 5164\n",
    "# 795\n",
    "# 5380\n",
    "# 3001\n",
    "print(\"bias:\", smaller_auto_encoder.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(full_text, best_feature, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_text(text_list, best_feature, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ex_sent_combined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ex_sent_combined = [\n",
    "#     '<|endoftext|>Then, Richard and Jennifer were working at the plateau. Jennifer decided to give a feather to Richard',\n",
    "#     '<|endoftext|>Then, Laura and Catherine were working at the cafe. Catherine decided to give a towel to Laura',\n",
    "#     '<|endoftext|>Then, Parker and Jake were working at the glacier. Jake decided to give fins to Parker',\n",
    "#     '<|endoftext|>Then, Jennifer and Richard were working at the plateau. Richard decided to give a feather to Jennifer',\n",
    "#     '<|endoftext|>Then, Catherine and Laura were working at the cafe. Laura decided to give a towel to Catherine',\n",
    "#     '<|endoftext|>Then, Jake and Parker were working at the glacier. Parker decided to give fins to Jake'\n",
    "# ]\n",
    "\n",
    "# 3383\n",
    "# 6035\n",
    "# 1413\n",
    "# 4466\n",
    "# 3806\n",
    "# 3667\n",
    "# 5455\n",
    "# 4290\n",
    "# 5164\n",
    "# 795\n",
    "# 5380\n",
    "# 3001\",\n",
    "# New ones\n",
    "features_ind_found = [\n",
    "    3527,\n",
    "    4211,\n",
    "    2675,\n",
    "    4006,\n",
    "    5189,\n",
    "    4478,\n",
    "    4977,\n",
    "    852,\n",
    "    2898,\n",
    "    5112,\n",
    "    4879,\n",
    "    3971,\n",
    "    4261,\n",
    "    1746,\n",
    "    4647,\n",
    "    2022,\n",
    "    3160,\n",
    "    32,\n",
    "    3657,\n",
    "    3173\n",
    "]\n",
    "ex_sent_combined = [ex_s + a_word for ex_s, a_word in zip(ex_sent, all_A)] \n",
    "wrong_sentences = [ex_s + b_word for ex_s, b_word in zip(ex_sent, all_B)] \n",
    "# feature_list\n",
    "# 1 is okay-ish\n",
    "# 3 is just after 'and'\n",
    "# 4 is a mixed bag, but activates for names\n",
    "# 5 mixed bag, but worse\n",
    "# 6 is mixed\n",
    "# 7 ???\n",
    "selector = 1\n",
    "best_feature = feature_list[selector]\n",
    "print(feature_list[selector])\n",
    "print(torch.nn.functional.cosine_similarity(smaller_dict[best_feature], torch.tensor(original_feature_2), dim=0))\n",
    "# best_feature = feature_list2[selector]\n",
    "# best_feature = 3657 # & selector 0\n",
    "# selector = 8\n",
    "# cs_features = [5051, 3149, 4725, 3657, 1827, 3944, 2945, 4863, 1056,  432]\n",
    "# cs_features =  [4811, 1266, 1491, 1486, 4124, 4502, 2847, 4923, 1940, 2096]\n",
    "# best_feature = cs_features[selector]\n",
    "visualize_text(ex_sent_combined, best_feature, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(ex_sent_combined, best_feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(wrong_sentences, best_feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(wrong_sentences, original_feature.to(device), entire_feature_direction=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(ex_sent_combined, original_feature.to(device), entire_feature_direction=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(wrong_sentences, torch.tensor(original_feature_2).to(device), entire_feature_direction=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(ex_sent_combined, torch.tensor(original_feature_2).to(device), entire_feature_direction=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(full_text, best_feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_feature.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_feature = torch.tensor([ 2.4080e-02, 2.4359e-02, 2.5915e-03, 4.9120e-02, 1.2658e-01,\n",
    "         3.2408e-02, -2.5182e-02, 4.3675e-02, -5.4410e-02, 1.6136e-02,\n",
    "        -2.0048e-02, -2.2420e-02, 3.3210e-02, 3.4793e-02, -2.8834e-02,\n",
    "         7.4147e-02, 3.7202e-02, 5.3582e-02, 1.1046e-03, 1.2104e-02,\n",
    "         5.3141e-03, 1.5514e-02, -3.4517e-02, -1.2846e-02, -1.4674e-02,\n",
    "         8.6960e-03, -1.5789e-02, 9.3915e-04, 5.9387e-04, -6.1752e-02,\n",
    "         4.1706e-02, -6.9348e-02, 6.0245e-02, -1.5278e-02, 4.6114e-03,\n",
    "        -6.0599e-02, 4.6384e-02, 8.7855e-03, 9.8222e-02, -1.4912e-02,\n",
    "        -2.0680e-02, 8.0620e-04, 2.1009e-02, -1.8937e-02, 1.9966e-02,\n",
    "         1.8603e-02, 9.0589e-02, -5.7756e-03, -4.1959e-02, 5.0267e-03,\n",
    "         2.3697e-02, -5.8151e-02, -7.7805e-02, -4.9728e-02, -4.8614e-03,\n",
    "        -3.9631e-02, -4.0052e-02, 3.9487e-02, 6.3542e-03, 2.7683e-02,\n",
    "        -3.4077e-02, -2.9967e-02, 3.2453e-03, -8.4002e-02, -3.5077e-02,\n",
    "        -9.1061e-03, -8.3716e-04, 7.6448e-03, -7.9082e-03, 3.9118e-02,\n",
    "        -3.9306e-02, -2.0184e-02, -1.7360e-02, -9.5160e-03, 1.3632e-02,\n",
    "        -1.1137e-01, -3.2181e-02, -5.6427e-02, -2.3615e-02, 9.5962e-03,\n",
    "         5.5335e-02, -1.8662e-02, 3.4204e-02, -5.0086e-02, -7.4954e-02,\n",
    "         8.4810e-03, 5.8262e-02, 1.7019e-02, 3.1432e-03, 1.1476e-02,\n",
    "         6.3471e-03, -4.1524e-02, 1.8465e-02, 7.5723e-02, 1.5079e-02,\n",
    "         1.9857e-02, 1.1886e-02, 9.8358e-03, -5.3220e-02, 5.9651e-02,\n",
    "         2.2428e-02, 3.4533e-02, -1.7907e-02, 1.0397e-02, 4.1672e-02,\n",
    "        -1.8936e-02, 2.1871e-02, 2.5337e-03, 2.3077e-03, 3.9737e-02,\n",
    "        -3.9214e-02, 9.2625e-03, -3.8041e-02, -1.2577e-02, -3.1696e-03,\n",
    "        -2.7372e-02, -5.1618e-02, 4.4042e-02, 2.4820e-02, 2.7953e-02,\n",
    "         1.5850e-02, -1.6006e-03, 5.3238e-02, 9.8043e-03, 3.8326e-02,\n",
    "        -4.1484e-02, 3.6432e-02, 2.2653e-02, 1.8522e-02, -3.7113e-02,\n",
    "         1.7174e-02, -8.3303e-03, 2.8099e-03, 1.0931e-01, -4.4335e-02,\n",
    "         1.1882e-02, -6.4499e-03, 3.7061e-02, -9.4151e-02, -5.1065e-02,\n",
    "         2.7060e-02, -6.5621e-02, -3.0455e-02, -1.6427e-02, 3.9223e-02,\n",
    "        -3.9643e-02, 6.4269e-04, -1.2907e-03, -5.0777e-03, -5.1587e-02,\n",
    "        -3.4985e-02, 3.7856e-02, -9.9967e-02, -6.5395e-03, -4.6955e-02,\n",
    "         2.4562e-02, 1.5419e-02, -2.1492e-02, -8.9394e-03, 4.2825e-02,\n",
    "        -4.5467e-02, -1.2693e-02, -4.9931e-02, 1.9798e-03, -3.1977e-02,\n",
    "         6.7696e-03, -3.7578e-02, -3.1094e-02, -3.4964e-02, 3.0003e-02,\n",
    "         3.0392e-02, 2.8398e-02, 7.5477e-04, 2.7859e-02, 8.7245e-03,\n",
    "         1.9231e-02, -3.1447e-02, -3.8160e-02, 3.6587e-02, -1.2608e-02,\n",
    "         6.7020e-02, 9.5656e-02, 2.0806e-02, -1.0801e-02, -1.6442e-02,\n",
    "         2.9674e-02, -4.6152e-02, -6.7173e-02, 5.4613e-02, -7.5615e-03,\n",
    "        -1.2474e-05, -9.5788e-03, 6.4168e-03, 1.7756e-02, -1.0295e-02,\n",
    "        -2.1090e-02, -3.4596e-02, -1.4731e-02, 1.0119e-02, 5.8010e-03,\n",
    "        -5.4254e-03, -4.6099e-02, 6.3387e-02, -1.7798e-02, -3.1989e-02,\n",
    "        -2.7206e-02, -7.4612e-02, 3.3296e-02, 9.3995e-03, 3.7320e-02,\n",
    "        -8.7339e-02, 2.0072e-02, -3.8689e-02, -9.5642e-03, -2.8851e-03,\n",
    "         3.4954e-03, -3.7127e-03, -3.8476e-03, 1.9904e-02, -4.4326e-03,\n",
    "         2.4047e-02, 2.1409e-02, -5.3637e-04, -1.3460e-02, -2.5007e-02,\n",
    "        -3.6522e-02, -6.4226e-02, 6.0937e-03, -4.9409e-03, -5.5151e-02,\n",
    "         1.1245e-02, -5.9511e-02, 4.5356e-02, 2.3629e-02, -3.4198e-03,\n",
    "         2.6451e-02, 5.8297e-02, -1.5115e-02, 3.8943e-02, -1.3371e-02,\n",
    "        -2.6044e-02, -5.6735e-02, 6.8021e-02, -3.7823e-02, 7.3421e-02,\n",
    "        -2.4523e-02, -6.3215e-03, -8.1932e-03, 5.6012e-02, 4.7597e-04,\n",
    "        -4.0288e-02, -1.3732e-03, -3.2098e-02, -9.2334e-03, 5.5161e-03,\n",
    "         2.5030e-02, 1.5095e-02, -2.9644e-02, 5.2908e-02, -2.3363e-02,\n",
    "         5.3523e-03, 3.9384e-03, 2.1839e-02, 1.2031e-02, 2.4092e-02,\n",
    "         2.3921e-02, -7.7738e-02, -1.1835e-01, -5.1954e-02, -2.1009e-02,\n",
    "         2.5704e-02, -2.4611e-02, 3.3695e-02, 2.6925e-02, 2.6932e-02,\n",
    "         1.2319e-02, -8.9244e-02, -8.6162e-03, 3.5292e-02, -3.8614e-02,\n",
    "        -3.9250e-02, 1.6301e-02, -3.1443e-02, 1.9028e-03, 5.7968e-02,\n",
    "         3.9234e-03, -2.0115e-02, -7.6918e-03, 4.2788e-02, 3.0034e-02,\n",
    "         9.1978e-02, -5.9110e-02, 2.3647e-02, 2.0807e-02, -1.2634e-02,\n",
    "        -7.6869e-03, 4.4716e-02, -1.8892e-02, 2.7425e-02, -9.9979e-03,\n",
    "         1.9796e-02, 1.0973e-02, -1.6451e-03, -4.3213e-02, -3.5061e-03,\n",
    "         2.9969e-03, -7.1300e-02, 6.3207e-03, 8.0630e-03, -3.5077e-02,\n",
    "        -9.7324e-03, -5.2640e-02, -5.5317e-02, -2.2824e-02, 2.7998e-02,\n",
    "        -7.8537e-02, 4.3458e-02, -1.4984e-02, -7.7363e-02, 3.7290e-02,\n",
    "         7.7163e-02, -3.3954e-03, 1.3120e-02, 5.7814e-02, -1.7183e-02,\n",
    "         2.4303e-03, 3.1352e-02, 7.0760e-02, -5.9879e-03, 6.1438e-02,\n",
    "        -9.8398e-03, 1.6338e-02, 2.8400e-02, 5.7705e-02, -2.2048e-03,\n",
    "        -4.5825e-03, -1.2012e-02, -1.3619e-03, 3.8692e-02, -1.8441e-02,\n",
    "        -1.1080e-02, 3.3503e-03, 1.1516e-02, 5.0336e-02, 2.0678e-03,\n",
    "        -2.6509e-02, 4.7976e-03, -4.7704e-03, -3.4007e-02, -5.1153e-02,\n",
    "        -1.9541e-02, -1.5282e-02, -3.3280e-02, 2.1600e-02, -6.8638e-02,\n",
    "         7.2270e-03, -4.8205e-02, 6.4255e-02, -3.0913e-02, 3.3795e-02,\n",
    "         3.1021e-02, 1.7689e-02, 7.8526e-03, 4.4164e-02, -2.6782e-03,\n",
    "         2.9056e-02, 3.4432e-02, 3.8272e-02, 7.4780e-03, 5.0119e-02,\n",
    "        -1.8974e-02, -2.9729e-02, 4.4389e-02, -2.5922e-02, -6.9515e-02,\n",
    "        -6.1480e-02, -1.2153e-02, -5.7087e-02, 8.2086e-02, 3.9201e-02,\n",
    "         2.0275e-02, 8.8203e-03, -4.9722e-02, 6.9982e-02, -7.2797e-03,\n",
    "         4.0989e-02, 4.1517e-03, 8.9593e-03, 2.5157e-02, -4.8101e-03,\n",
    "        -9.6375e-03, -2.2509e-02, 5.1393e-02, 7.2553e-03, 1.3507e-02,\n",
    "         4.3849e-02, 6.2116e-02, 3.8052e-02, 1.5074e-02, 1.4775e-02,\n",
    "         2.6600e-02, 2.7580e-02, -6.2622e-03, -2.5771e-02, 1.7867e-02,\n",
    "        -1.4621e-02, 1.3103e-02, 1.5070e-02, -4.5592e-02, 6.9849e-06,\n",
    "         5.1731e-03, -2.9146e-02, -1.2707e-02, 2.8848e-02, 8.5701e-02,\n",
    "        -1.1278e-02, 2.7468e-02, -4.3419e-03, -2.8725e-02, -3.2829e-02,\n",
    "        -1.2745e-02, 3.3975e-02, -1.1140e-02, -3.2737e-02, -3.7541e-02,\n",
    "        -1.2454e-02, -1.5200e-02, -2.8173e-02, -1.0102e-01, -1.6749e-02,\n",
    "         2.8634e-02, -1.3546e-02, -1.8982e-02, 9.3449e-02, 4.1330e-02,\n",
    "        -1.3391e-02, -2.9870e-02, -1.0827e-02, -4.5005e-02, 1.5958e-02,\n",
    "        -2.7183e-03, -3.2368e-02, 5.6117e-02, -4.7420e-02, -2.1395e-02,\n",
    "        -4.4957e-02, -5.9798e-03, 2.2991e-02, 7.0553e-02, -9.5782e-03,\n",
    "         9.5198e-03, 5.8040e-03, -3.0258e-02, -5.3331e-02, 8.5738e-03,\n",
    "         5.9829e-02, 4.9407e-02, -3.9949e-02, -1.5012e-02, -5.7299e-02,\n",
    "         3.3791e-02, 2.6755e-02, 1.6138e-02, 1.5216e-02, -1.4221e-02,\n",
    "        -5.5148e-02, -6.4813e-03, -4.5133e-02, 2.5805e-02, 1.9024e-02,\n",
    "         2.3914e-03, 5.3716e-03, -5.0496e-02, -3.6670e-04, -2.6395e-02,\n",
    "         5.3805e-02, -1.9281e-02, 2.9166e-02, 1.8902e-02, 6.9283e-03,\n",
    "        -8.6916e-03, 1.3865e-01, 4.3540e-02, 7.9599e-03, -1.1489e-02,\n",
    "        -1.9561e-02, -3.3223e-02, 4.2447e-02, 2.4892e-02, 7.6215e-03,\n",
    "        -5.5037e-02, -4.0666e-03, -4.7167e-03, 3.1368e-02, 2.2559e-02,\n",
    "        -9.4491e-03, -3.7170e-02, -4.0265e-02, -2.2772e-04, -6.9743e-03,\n",
    "        -7.5886e-03, 3.6595e-02, -1.0402e-02, -3.0940e-02, -1.4825e-02,\n",
    "         1.4237e-02, 4.5008e-02, 2.5084e-02, 1.0162e-02, 2.5703e-02,\n",
    "        -1.3544e-02, -1.1389e-01, -4.0357e-02, -4.7644e-02, -3.7380e-02,\n",
    "        -1.4030e-02, 5.7597e-02, 1.7919e-04, 1.8142e-02, -1.6253e-02,\n",
    "        -7.6142e-02, 3.8851e-02, -2.1893e-02, -1.5224e-02, 2.1425e-02,\n",
    "         1.6085e-02, -2.6741e-02, -2.4140e-02, -1.4206e-02, 4.1798e-02,\n",
    "        -3.2894e-02, 4.9439e-02, -9.9219e-03, -7.2214e-02, -4.4329e-03,\n",
    "         2.9387e-02, -3.9025e-02, 1.9442e-02, 3.4890e-02, 1.3546e-02,\n",
    "         6.8964e-04, -2.7797e-02, -6.8082e-02, 9.9170e-03, 2.3642e-02,\n",
    "        -6.8385e-02, 1.0770e-02, -5.6083e-03, 1.7623e-02, -3.4506e-02,\n",
    "        -4.2243e-02, 7.4690e-02, -2.4971e-02, -5.2503e-02, -1.4595e-02,\n",
    "         5.1033e-02, 6.1160e-02, 3.7497e-03, -3.8879e-02, -4.1647e-03,\n",
    "         2.6918e-02, -1.4865e-02, 9.0000e-03, -1.4608e-02, -2.1336e-02,\n",
    "        -5.9530e-03, -1.6575e-03, -4.7094e-02, 2.3869e-02, 1.4982e-02,\n",
    "         4.9038e-02, -4.5014e-02, 4.8011e-02, -1.5510e-02, -1.6457e-02,\n",
    "         1.2145e-02, 4.9655e-03, -4.4219e-02, 2.6135e-02, 1.6504e-02,\n",
    "         3.7103e-03, 3.4081e-02, -8.6893e-03, -9.9807e-03, 1.2927e-02,\n",
    "         3.5594e-02, 2.4592e-02, 3.7850e-02, -1.1348e-02, -1.1096e-03,\n",
    "        -4.4439e-02, 2.7691e-02, -6.9026e-02, 5.7734e-02, 2.9289e-02,\n",
    "         3.4821e-02, -5.2970e-03, 3.7610e-02, -1.7178e-02, -5.0140e-02,\n",
    "        -8.3218e-02, 4.2783e-02, -1.2670e-02, 1.1850e-02, 3.1992e-02,\n",
    "        -4.4622e-04, 3.7421e-02, -2.3873e-02, -1.9229e-02, 1.4420e-02,\n",
    "         1.5459e-02, 2.4360e-02, -1.4355e-03, 3.8322e-02, -4.1234e-02,\n",
    "        -4.1331e-02, -2.9026e-02, -1.7661e-02, 5.0412e-02, -5.7450e-02,\n",
    "        -2.2180e-02, 1.4419e-02, -1.6419e-02, -2.6536e-02, -6.5096e-02,\n",
    "        -8.5083e-03, -5.8017e-03, 4.6651e-02, 3.4101e-02, -1.8907e-02,\n",
    "         2.7087e-02, 2.5500e-02, 1.1684e-03, 6.8219e-02, -3.2993e-02,\n",
    "        -2.3292e-02, 1.5120e-03, -2.1995e-02, 6.2556e-02, 4.9651e-02,\n",
    "         1.6080e-02, -1.3513e-02, 3.5774e-02, -4.4417e-02, 4.9697e-03,\n",
    "         4.9093e-02, 1.3976e-02, 2.6455e-02, 3.7071e-02, -4.3469e-02,\n",
    "         1.3540e-02, 5.2918e-02, 3.4559e-02, -3.1488e-02, -1.9801e-02,\n",
    "        -3.4259e-02, -1.0953e-02, -3.2256e-02, -3.9682e-02, 9.6360e-03,\n",
    "        -9.6176e-03, 1.2003e-03, 4.2134e-02, -5.7197e-02, -4.2081e-02,\n",
    "        -4.6472e-03, 2.6133e-02, 3.7587e-02, 9.8670e-03, -3.4046e-02,\n",
    "        -3.2811e-03, -1.5316e-02, -4.8647e-03, 1.9632e-02, 1.1838e-02,\n",
    "         1.5452e-02, -6.3881e-02, 1.2259e-02, 1.2232e-03, -4.3000e-02,\n",
    "        -6.8827e-02, 4.1810e-02, 7.1543e-03, -5.1488e-02, 2.4518e-02,\n",
    "        -1.0627e-03, 1.4221e-02, -5.3517e-02, -4.3952e-02, -3.0068e-02,\n",
    "        -1.9294e-03, -1.0837e-02, -1.0782e-02, 1.1658e-02, -7.7354e-03,\n",
    "        -9.4467e-03, 5.8047e-02, 6.2752e-03, 1.7554e-02, 5.2714e-04,\n",
    "         1.2364e-02, 4.7268e-02, -3.3143e-02, -2.9131e-02, -7.5221e-03,\n",
    "        -1.3250e-02, -1.0116e-02, 1.5731e-02, 1.2194e-03, 5.5533e-02,\n",
    "        -2.4715e-02, 3.5104e-04, 7.7462e-02, -4.2483e-02, 9.3196e-03,\n",
    "         3.5436e-04, 3.0903e-02, -6.8819e-02, 1.9523e-02, 2.2681e-04,\n",
    "        -2.1906e-02, -3.0358e-03, 1.9632e-02, 1.0104e-02, -4.5632e-03,\n",
    "         2.3078e-02, -2.9149e-02, -4.9732e-02, 1.2652e-05, 4.3399e-03,\n",
    "         3.1472e-02, 1.0199e-02, -3.4592e-04, -3.0710e-02, 8.5242e-03,\n",
    "        -6.5155e-03, -3.7067e-02, -6.3603e-03, -1.0648e-02, 4.0477e-02,\n",
    "        -1.5365e-02, 2.0180e-02, 8.3195e-03, 1.8911e-02, -3.3727e-02,\n",
    "        -3.8646e-02, -3.2142e-02, -1.6740e-02, 2.0304e-02, 3.7467e-02,\n",
    "         2.5757e-02, -9.1715e-02, 6.6637e-02, -2.6891e-02, 3.7948e-02,\n",
    "        -2.8087e-02, -1.2200e-02, 2.7008e-02, -6.5081e-03, 7.5711e-02,\n",
    "         1.4658e-02, -1.6642e-02, 4.5490e-02, 1.1131e-02, 4.4595e-02,\n",
    "        -5.6698e-02, 1.6837e-02, -7.4032e-03])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logit_lens(model,best_feature, smaller_dict, layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_text = [\n",
    "    \"e-tikaのサンプル\",\n",
    "    \n",
    "    \"每个人都有他的作战策略，直到脸上中了一拳。\",\n",
    "    \"这是我年轻时候住的房子\",\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": [
    "# 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*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, \" \", 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
}