minimal_feature_interp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformer_lens import HookedTransformer\n",
    "import numpy as np \n",
    "# autoencoder_filename = \"/home/mchorse/logan/sparse_coding/output_sweep_tied_mlpout_l1_r8/_9/learned_dicts.pt\"\n",
    "# autoencoder_filename = \"/home/mchorse/logan/sparse_coding/output_sweep_tied_mlpout_l2_r4/_19/learned_dicts.pt\"\n",
    "autoencoder_filename = \"/mnt/ssd-cluster/longrun2408/tied_residual_l4_r6/_31/learned_dicts.pt\"\n",
    "auto_num = 5 # Selects which specific autoencoder to use\n",
    "all_autoencoders = torch.load(autoencoder_filename)\n",
    "num_dictionaries = len(all_autoencoders)\n",
    "autoencoder, hyperparams = all_autoencoders[auto_num]\n",
    "l1_alpha = hyperparams['l1_alpha']\n",
    "autoencoder2, hyperparams2 = all_autoencoders[auto_num+1]\n",
    "smaller_dict = autoencoder.get_learned_dict()\n",
    "larger_dict = autoencoder2.get_learned_dict()\n",
    "\n",
    "#Change these settings to load the correct autoencoder\n",
    "layer = 4\n",
    "setting = \"residual\"\n",
    "# setting = \"attention\"\n",
    "# setting = \"mlp\"\n",
    "# setting = \"mlp_out\"\n",
    "model_name = \"EleutherAI/pythia-70m-deduped\"\n",
    "# model_name = \"EleutherAI/pythia-160m\"\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",
    "elif setting == \"attention\":\n",
    "    cache_name = f\"blocks.{layer}.hook_attn_out\"\n",
    "    neurons = model.cfg.d_model\n",
    "elif setting == \"mlp_out\":\n",
    "    cache_name = f\"blocks.{layer}.hook_mlp_out\"\n",
    "    neurons = model.cfg.d_model\n",
    "else:\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparams, hyperparams2"
   ]
  },
  {
   "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]\n",
    "max_cosine_similarities[max_indices][:20]\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 = 32\n",
    "neuron_activations = torch.zeros((datapoints*token_amount, neurons))\n",
    "dictionary_activations = torch.zeros((datapoints*token_amount, smaller_dict_features))\n",
    "smaller_auto_encoder = autoencoder\n",
    "smaller_auto_encoder.to_device(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",
    "        batched_dictionary_activations = smaller_auto_encoder.encode(batched_neuron_activations)\n",
    "        dictionary_activations[i*batch_size*token_amount:(i+1)*batch_size*token_amount,:] = batched_dictionary_activations.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dictionary_activations[:100000].count_nonzero(dim=1).float().mean()"
   ]
  },
  {
   "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",
    "            neuron_act_batch = rearrange(neuron_act_batch, \"b s n -> (b s) n\" )\n",
    "            act = smaller_auto_encoder.encode(neuron_act_batch)\n",
    "            return act[:, feature].tolist()\n",
    "        else: # neuron/residual basis\n",
    "            return neuron_act_batch[0, :, feature].tolist()\n",
    "\n",
    "def ablate_text(text, feature, model, setting=\"dictionary_basis\", return_activations=False):\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, setting)[-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 return_activations:\n",
    "        return activation_list, display_text_list\n",
    "    else:\n",
    "        return text_neuron_activations(tokens=display_text_list, activations=activation_list)\n",
    "\n",
    "def visualize_text(text, feature, model, setting=\"dictionary_basis\", max_activation = None, return_activations=False):\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",
    "    if return_activations:\n",
    "        return act_list, display_text_list\n",
    "    else:\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):\n",
    "    def mlp_ablation_hook(value, hook):\n",
    "        # Rearrange to fit autoencoder\n",
    "        int_val = rearrange(value, 'b s h -> (b s) h')\n",
    "\n",
    "        # Run through the autoencoder\n",
    "        act = autoencoder.encode(int_val)\n",
    "        feature_to_ablate = feature # TODO: bring this out of the function\n",
    "\n",
    "        # Subtract value with feature direction*act_of_feature\n",
    "        dictionary_for_this_autoencoder = autoencoder.get_learned_dict()\n",
    "        feature_direction = torch.outer(act[:, feature_to_ablate].squeeze(), dictionary_for_this_autoencoder[feature_to_ablate].squeeze())\n",
    "        batch, seq_len, hidden_size = value.shape\n",
    "        feature_direction = rearrange(feature_direction, '(b s) h -> b s h', b=batch, s=seq_len)\n",
    "        value -= feature_direction\n",
    "        return value\n",
    "\n",
    "    return model.run_with_hooks(tokens, \n",
    "        fwd_hooks=[(\n",
    "            cache_name, \n",
    "            mlp_ablation_hook\n",
    "            )]\n",
    "        )\n",
    "def add_feature_direction(tokens, feature, model, autoencoder, scalar=1.0):\n",
    "    def residual_add_hook(value, hook):\n",
    "        feature_direction = autoencoder.decoder.weight[:, feature].squeeze()\n",
    "        value += scalar*feature_direction\n",
    "        return value\n",
    "\n",
    "    return model.run_with_hooks(tokens, \n",
    "        fwd_hooks=[(\n",
    "            cache_name,\n",
    "            residual_add_hook\n",
    "            )]\n",
    "        )\n",
    "def ablate_feature_direction_display(text, features=None, setting=\"true_tokens\", verbose=False):\n",
    "\n",
    "    if features==None:\n",
    "        features = torch.tensor([best_feature])\n",
    "    if isinstance(features, int):\n",
    "        features = torch.tensor([features])\n",
    "    if isinstance(features, list):\n",
    "        features = torch.tensor(features)\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    text_list = []\n",
    "    logit_list = []\n",
    "    for t in text:\n",
    "        tokens = model.to_tokens(t, prepend_bos=False)\n",
    "        with torch.no_grad():\n",
    "            original_logits = model(tokens).log_softmax(-1).cpu()\n",
    "            ablated_logits = ablate_feature_direction(tokens, features, model, smaller_auto_encoder).log_softmax(-1).cpu()\n",
    "        diff_logits = ablated_logits  - original_logits# ablated > original -> negative diff\n",
    "        tokens = tokens.cpu()\n",
    "        if setting == \"true_tokens\":\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            gather_tokens = rearrange(tokens[:,1:], \"b s -> b s 1\") # TODO: verify this is correct\n",
    "            # Gather the logits for the true tokens\n",
    "            diff = rearrange(diff_logits[:, :-1].gather(-1,gather_tokens), \"b s n -> (b s n)\")\n",
    "        elif setting == \"max\":\n",
    "            # Negate the diff_logits to see which tokens have the largest effect on the neuron\n",
    "            val, ind = (-1*diff_logits).max(-1)\n",
    "            diff = rearrange(val[:, :-1], \"b s -> (b s)\")\n",
    "            diff*= -1 # Negate the values gathered\n",
    "            split_text = model.to_str_tokens(ind, prepend_bos=False)\n",
    "            gather_tokens = rearrange(ind[:,1:], \"1 s -> 1 s 1\")\n",
    "        split_text = split_text[1:] # Remove the first token since we're not predicting it\n",
    "        if(verbose):\n",
    "            text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            orig = rearrange(original_logits[:, :-1].gather(-1, gather_tokens), \"b s n -> (b s n)\")\n",
    "            ablated = rearrange(ablated_logits[:, :-1].gather(-1, gather_tokens), \"b s n -> (b s n)\")\n",
    "            logit_list += orig.tolist() + [0.0]\n",
    "            logit_list += ablated.tolist() + [0.0]\n",
    "        text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "        logit_list += diff.tolist() + [0.0]\n",
    "    logit_list = torch.tensor(logit_list).reshape(-1,1,1)\n",
    "    if verbose:\n",
    "        print(f\"Max & Min logit-diff: {logit_list.max().item():.2f} & {logit_list.min().item():.2f}\")\n",
    "    return text_neuron_activations(tokens=text_list, activations=logit_list)\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": [
    "# Better feature search for token level features\n",
    "# Find top-k activating features index for each feature\n",
    "num_data, num_features = dictionary_activations.shape\n",
    "topk = 20\n",
    "search_first = 10000\n",
    "threshold = 0.1\n",
    "topk_features = torch.arange(topk).repeat(num_features, 1)\n",
    "for i in range(num_features):\n",
    "    # topk_features[i] = torch.topk(dictionary_activations[:search_first, i], topk)[1]\n",
    "    # Convert feature index into (datapoint, seq) index\n",
    "    activations = dictionary_activations[:search_first, i]\n",
    "    sorted_indices = torch.argsort(activations)  # Sort activations\n",
    "    non_zero_indices = sorted_indices[activations[sorted_indices] > threshold]  # Select only non-zero activations\n",
    "    \n",
    "    if len(non_zero_indices) < 10:  # No non-zero activations for this feature\n",
    "        continue\n",
    "\n",
    "    # # Take 'topk' evenly spaced samples from non_zero_indices\n",
    "    selected_indices = torch.linspace(0, len(non_zero_indices) - 1, steps=topk).long()\n",
    "    topk_features[i] = non_zero_indices[selected_indices]\n",
    "    # topk_features[i] = torch.cat((non_zero_indices[selected_indices][:topk//2], non_zero_indices[selected_indices][-(topk//2):]))\n",
    "\n",
    "topk_features = topk_features.long()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "activations[non_zero_indices[selected_indices]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_dim_indices = np.unravel_index(topk_features.flatten(), (datapoints, token_amount))\n",
    "\n",
    "tokens = np.zeros((num_features*topk), dtype=int)\n",
    "for i, (md, s_ind) in enumerate(tqdm(zip(*multi_dim_indices))):\n",
    "    tokens[i] = dataset[int(md)][\"input_ids\"][int(s_ind)]\n",
    "\n",
    "tokens = tokens.reshape(num_features, topk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "most_common_tokens = []\n",
    "most_common_tokens_count = []\n",
    "for feature_tokens in tokens:\n",
    "    token_counter = Counter(feature_tokens)\n",
    "    most_common_token, count = token_counter.most_common(1)[0]\n",
    "    \n",
    "    most_common_tokens.append(most_common_token)\n",
    "    most_common_tokens_count.append(count)\n",
    "\n",
    "# Convert to NumPy arrays for further manipulation\n",
    "most_common_tokens = np.array(most_common_tokens)\n",
    "most_common_tokens_count = np.array(most_common_tokens_count)\n",
    "\n",
    "# Argsort the array\n",
    "sorted_indices = np.argsort(most_common_tokens_count)[::-1]\n",
    "\n",
    "print(\"Most common tokens:\", most_common_tokens)\n",
    "print(\"Count of most common tokens:\", most_common_tokens_count[sorted_indices])\n",
    "print(\"Argsorted indices:\", sorted_indices)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "most_com_tok = most_common_tokens[sorted_indices][:200]\n",
    "# Count the number of times each token appears\n",
    "c =  Counter(most_com_tok).most_common(20)\n",
    "[(model.tokenizer.decode([x[0]]), x[1]) for x in c], [model.tokenizer.decode([x[0]]) for x in c]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = [x[0] for x in c]\n",
    "tokens_to_check = [187, 13, 14, 15, 253, 285, 50, 281, 247, 273, 313, 275, 27, 457, 16, 326, 352, 310, 50276, 452]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_indices[:10], most_common_tokens_count[sorted_indices][:10], most_common_tokens[sorted_indices][:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# t = \" I do like a\"\n",
    "t = \" He had a first one (1), and then a second (2\"\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.encode(neuron_act_batch.squeeze())\n",
    "neg = \" He had a first one (1), and then a second 2\"\n",
    "split_text = model.to_str_tokens(neg, prepend_bos=False)\n",
    "token = model.to_tokens(neg, prepend_bos=False)\n",
    "_, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "neuron_act_batch = cache[cache_name]\n",
    "act[-1, :] -= smaller_auto_encoder.encode(neuron_act_batch.squeeze())[-1,:]\n",
    "v, i = act[-1, :].topk(10)\n",
    "\n",
    "print(\"Activations:\",[round(val,2) for val in v.tolist()])\n",
    "print(\"Feature_ids\", i.tolist())"
   ]
  },
  {
   "cell_type": "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": [
    "# N = 106 # 106\n",
    "# N = 116 # 106\n",
    "# best_feature = int(max_indices[N])\n",
    "best_feature = 1814  #array([ 333, 2126, 2933,  156,  820, 2652,  911,  307, 2004,  169])\n",
    "best_feature = 69  #(array([ 333, 2004,  553, 1770, 2126, 1878, 1296,   97, 1094,   69]),\n",
    "# best_feature = 11 # Change this one for global index (N is sorted by MCS)\n",
    "# best_feature =  699 # 1841, 1034, 1359, 1987,  795, 578, 1501, 1603, 1448, 699\n",
    "# best_feature =  776 # 1841, 1034, 1987, 1448, 776, 1319, 1217, 1603, 1594, 1500\n",
    "N_ind = 6\n",
    "best_feature = int(sorted_indices[N_ind])\n",
    "print(most_common_tokens_count[sorted_indices][N_ind])\n",
    "print(f\"Feature index: {best_feature}\")\n",
    "print(f\"MCS: {max_cosine_similarities[best_feature]}\")\n",
    "text_list, full_text, token_list, full_token_list = get_feature_datapoints(best_feature, dictionary_activations, dataset, setting=\"uniform\")\n",
    "# text_list, full_text, token_list, full_token_list = get_feature_datapoints(best_feature, dictionary_activations, dataset, setting=\"max\")\n",
    "visualize_text(text_list, best_feature, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_indices[N_ind]"
   ]
  },
  {
   "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": [
    "custom_text = [\n",
    "    \" 3.2/1.6\",\n",
    "    \" 5.6.1.53.2\",\n",
    "    \" This is KZ 96.5\"\n",
    "]\n",
    "# visualize_text(custom_text, best_feature, model)\n",
    "act_list, text_display_list = visualize_text(custom_text, best_feature, model, return_activations=True)\n",
    "ablated_act_list, _ = ablate_text(custom_text, best_feature, model, return_activations=True)\n",
    "act_list = act_list.squeeze().squeeze()\n",
    "ablated_act_list = ablated_act_list.squeeze().squeeze()\n",
    "# Find indices of \"\\n\" in text_list\n",
    "split_indices = [i for i, text in enumerate(text_display_list) if text == \"\\n\"]\n",
    "\n",
    "# Split text_list based on indices of \"\\n\"\n",
    "split_text_lists = [text_display_list[i:j] for i, j in zip([0]+[idx+1 for idx in split_indices], split_indices+[None])][:-1]\n",
    "\n",
    "# Similarly, split act_list. Since there's a one-to-one correspondence between text_list and act_list\n",
    "split_act_lists = [act_list[i:j] for i, j in zip([0]+[idx+1 for idx in split_indices], split_indices+[None])][:-1]\n",
    "\n",
    "split_ablated_act_lists = [ablated_act_list[i:j] for i, j in zip([0]+[idx+1 for idx in split_indices], split_indices+[None])][:-1]\n",
    "\n",
    "#Save split_text/act/ablated\n",
    "import pickle\n",
    "with open(f\"images/feature_activation_and_ablated_{best_feature}.pkl\", \"wb\") as f:\n",
    "    pickle.dump((split_text_lists, split_act_lists, split_ablated_act_lists), f)\n",
    "import matplotlib.pyplot as plt\n",
    "invisible_char = '\\u200B'\n",
    "\n",
    "nrows = len(split_text_lists)\n",
    "fig, axes = plt.subplots(nrows=nrows, figsize=(3, 2*nrows))\n",
    "# Iterate through both split lists and plot them\n",
    "for text_ind in range(nrows):\n",
    "# for text_seq, act_seq in zip(split_text_lists, split_act_lists):\n",
    "    ax = axes[text_ind]\n",
    "    text_seq = split_text_lists[text_ind]\n",
    "    act_seq = split_act_lists[text_ind]\n",
    "    # ablated_act_seq = split_ablated_act_lists[text_ind]\n",
    "    # Modify the text_seq to add the invisible character\n",
    "    modified_text_seq = [text + invisible_char*i for i, text in enumerate(text_seq)]\n",
    "    # plt.figure(figsize=(3, 3))  # Create a new figure for each text sequence\n",
    "    # Color is blue and labeled \"activation\"\n",
    "    ax.plot(modified_text_seq, act_seq, color='b', label='Feature Activation')\n",
    "    # ax.plot(modified_text_seq, ablated_act_seq, color='r', label='Ablated Activation-diff of last token')\n",
    "    ax.legend()  # Show the legend\n",
    "    ax.set_xlabel(\"Text Tokens\")\n",
    "    ax.set_ylabel(\"Activations\")\n",
    "\n",
    "# Adjust layout for better display\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(right=0.85) \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrows = len(split_text_lists)\n",
    "fig, axes = plt.subplots(nrows=nrows, figsize=(3, 2*nrows))\n",
    "last_token_activation = [act_list[-1] for act_list in split_act_lists]\n",
    "for text_ind in range(nrows):\n",
    "# for text_seq, act_seq in zip(split_text_lists, split_act_lists):\n",
    "    ax = axes[text_ind]\n",
    "    text_seq = split_text_lists[text_ind]\n",
    "    ablated_act_seq = split_ablated_act_lists[text_ind]\n",
    "    # Modify the text_seq to add the invisible character\n",
    "    modified_text_seq = [text + invisible_char*i for i, text in enumerate(text_seq)]\n",
    "    last_token_act = last_token_activation[text_ind] \n",
    "    ablated_act_seq += last_token_act\n",
    "    ax.plot(modified_text_seq, ablated_act_seq, color='r', label='Ablated Activation-diff of last token')\n",
    "\n",
    "    # Plot a blue, dotted line to show the activation of the last token\n",
    "    # ax.plot(modified_text_seq, [0, last_token_act], color='b', linestyle='--', label='Activation of last token')\n",
    "    ax.plot(modified_text_seq, [last_token_act]*len(modified_text_seq), color='b', linestyle='--', label='Activation of last token')\n",
    "    ax.legend()  # Show the legend\n",
    "    ax.set_xlabel(\"Text Tokens\")\n",
    "    ax.set_ylabel(\"Activations\")\n",
    "\n",
    "# Adjust layout for better display\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(right=0.85) \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.show()"
   ]
  },
  {
   "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": [
    "logit_lens(model,best_feature, smaller_dict, layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_text = [\n",
    "    \" 3.2/1.6\",\n",
    "    \" Section 5.6.1\",\n",
    "    \" This is KZ 96.5\"\n",
    "]\n",
    "visualize_text(custom_text, best_feature, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Centric Viewpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Go through datapoints & see if the features that activate on them make sense.\n",
    "d_point = 0\n",
    "# text = tokens_dataset[d_point]\n",
    "data_ind, sequence_pos = np.unravel_index(d_point, (datapoints, token_amount))\n",
    "feature_val, feature_ind = dictionary_activations[d_point].topk(10)\n",
    "data_ind = int(data_ind)\n",
    "sequence_pos = int(sequence_pos)\n",
    "full_tok = torch.tensor(dataset[data_ind][\"input_ids\"])\n",
    "full_text = []\n",
    "full_text.append(model.tokenizer.decode(full_tok))\n",
    "visualize_text(full_text, feature_ind, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check the neuron/residual basis\n",
    "When we look at the weights of a feature, we are seeing the literal dimensions from the residual stream/neurons being read from the feature. \n",
    "\n",
    "Here I'm visualizing the weight values for the residual stream. If there are outliers, then it's mainly reading from that dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(weights*max_activation).topk(20), (weights*max_activation).topk(20, largest=False).values, (weights*max_activation > 0.2).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepend/Append tokens\n",
    "We can iterate over all tokens to check which ones activate a feature a lot to more rigorously test a hypothesis on what a feature means.\n",
    "\n",
    "Note: I'm literately running the model through all 50k tokens prepended to the text here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepend_all_tokens_and_get_feature_activation(model, minimal_activating_example, feature, setting=\"prepend\"):\n",
    "    tokens = model.to_tokens(minimal_activating_example, prepend_bos=False)\n",
    "\n",
    "    # Run through every number up to vocab size\n",
    "    vocab_size = model.cfg.d_vocab\n",
    "    batch_size = 256*2 # Define your desired batch size\n",
    "\n",
    "    dollar_feature_activations = torch.zeros(vocab_size)\n",
    "    for start in range(0, vocab_size, batch_size):\n",
    "        end = min(start + batch_size, vocab_size)\n",
    "\n",
    "        token_prep = torch.arange(start, end).to(device)\n",
    "        token_prep = token_prep.unsqueeze(1)  # Add a dimension for concatenation\n",
    "\n",
    "        # 1. Prepend to the tokens\n",
    "        if setting == \"prepend\":\n",
    "            tokens_catted = torch.cat((token_prep, tokens.repeat(end - start, 1)), dim=1).long()\n",
    "        elif setting == \"append\":\n",
    "            tokens_catted = torch.cat((tokens.repeat(end - start, 1), token_prep), dim=1).long()\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown setting: {setting}\")\n",
    "\n",
    "        # 2. Run through the model\n",
    "        with torch.no_grad():\n",
    "            _, cache = model.run_with_cache(tokens_catted.to(device))\n",
    "            neuron_act_batch = cache[cache_name]\n",
    "            _, act = smaller_auto_encoder(neuron_act_batch)\n",
    "\n",
    "        # 3. Get the feature\n",
    "        dollar_feature_activations[start:end] = act[:, -1, feature].cpu().squeeze()\n",
    "\n",
    "    k = 20\n",
    "    k_increasing_val, k_increasing_ind = dollar_feature_activations.topk(k)\n",
    "    k_decreasing_val, k_decreasing_ind = dollar_feature_activations.topk(k, largest=False)\n",
    "    if(setting == \"prepend\"):\n",
    "        print(f\"[token]{minimal_activating_example}\")\n",
    "    elif(setting == \"append\"):\n",
    "        print(f\"{minimal_activating_example}[token]\")\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown setting: {setting}\")\n",
    "    # Print indices converted to tokens\n",
    "    print(f\"Top-{k} increasing: {model.to_str_tokens(k_increasing_ind)}\")\n",
    "    # Print values\n",
    "    print(f\"Top-{k} increasing: {[f'{val:.2f}' for val in k_increasing_val]}\")\n",
    "    print(f\"Top-{k} decreasing: {model.to_str_tokens(k_decreasing_ind)}\")\n",
    "    print(f\"Top-{k} decreasing: {[f'{val:.2f}' for val in k_decreasing_val]}\")\n",
    "    print(f\"Number of 0 activations: {torch.sum(dollar_feature_activations == 0)}\")\n",
    "    if(setting == \"prepend\"):\n",
    "        best_text = \"\".join(model.to_str_tokens(dollar_feature_activations.argmax()) + [minimal_activating_example])\n",
    "    else:\n",
    "        best_text = \"\".join([minimal_activating_example] + model.to_str_tokens(dollar_feature_activations.argmax()))\n",
    "    return best_text\n",
    "\n",
    "best_text = \"\"\n",
    "for x in range(3):\n",
    "    # best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"prepend\")\n",
    "    best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"append\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" for all $\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"append\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_text = \"\"\n",
    "for x in range(3):\n",
    "    best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"prepend\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepend_all_tokens_and_get_feature_activation(model, \" for all $\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \"The\", best_feature, setting=\"append\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"append\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}