case_studies_loop.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",
    "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, autoencoder, 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 = autoencoder.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, autoencoder, setting=\"dictionary_basis\"):\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    display_text_list = []\n",
    "    activation_list = []\n",
    "    for t in text:\n",
    "        # Convert text into tokens\n",
    "        if isinstance(t, str): # If the text is a list of tokens\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            tokens = model.to_tokens(t, prepend_bos=False)\n",
    "        else: # t equals tokens\n",
    "            tokens = t\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "        seq_size = tokens.shape[1]\n",
    "        if(seq_size == 1): # If the text is a single token, we can't ablate it\n",
    "            continue\n",
    "        original = get_neuron_activation(tokens, feature, model, autoencoder)[-1]\n",
    "        changed_activations = torch.zeros(seq_size, device=device).cpu()\n",
    "        for i in range(seq_size):\n",
    "            # Remove the i'th token from the input\n",
    "            ablated_tokens = torch.cat((tokens[:,:i], tokens[:,i+1:]), dim=1)\n",
    "            changed_activations[i] += get_neuron_activation(ablated_tokens, feature, model,autoencoder, 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",
    "    return text_neuron_activations(tokens=display_text_list, activations=activation_list)\n",
    "\n",
    "def visualize_text(text, feature, model, autoencoder, setting=\"dictionary_basis\", max_activation = None):\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    if isinstance(feature, int):\n",
    "        feature = [feature]\n",
    "    display_text_list = []\n",
    "    act_list = []\n",
    "    for t in text:\n",
    "        if isinstance(t, str): # If the text is a list of tokens\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            token = model.to_tokens(t, prepend_bos=False)\n",
    "        else: # t are tokens\n",
    "            token = t\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "        for f in feature:\n",
    "            display_text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            act_list += get_neuron_activation(token, f, model, autoencoder, 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):\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, autoencoder, setting=\"true_tokens\", verbose=False):\n",
    "\n",
    "    if features==None:\n",
    "        features = torch.tensor([best_feature])\n",
    "    if isinstance(features, int):\n",
    "        features = torch.tensor([features])\n",
    "    if isinstance(features, list):\n",
    "        features = torch.tensor(features)\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    text_list = []\n",
    "    logit_list = []\n",
    "    for t in text:\n",
    "        tokens = model.to_tokens(t, prepend_bos=False)\n",
    "        with torch.no_grad():\n",
    "            original_logits = model(tokens).log_softmax(-1).cpu()\n",
    "            ablated_logits = ablate_feature_direction(tokens, features, model, autoencoder).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",
    "\n",
    "import re\n",
    "import numpy as np\n",
    "\n",
    "# Assuming you have a tokenizer like this (for example, from HuggingFace's tokenizers)\n",
    "  # Replace with your tokenizer\n",
    "\n",
    "def char_to_token_position(tokenized_output, char_idx):\n",
    "    \"\"\"\n",
    "    Given a char-level index in the original string, find the corresponding token-level index.\n",
    "    \"\"\"\n",
    "    token_positions = tokenized_output.char_to_token(char_idx)\n",
    "    return token_positions\n",
    "\n",
    "def find_text_in_regex(tokens, regex_pattern):\n",
    "    \"\"\"\n",
    "    Find tokens inside parentheses and return token-level start and end indices.\n",
    "    \"\"\"\n",
    "    detokenized_text = model.tokenizer.decode(tokens.input_ids)\n",
    "    \n",
    "    # Find matches in the detokenized text\n",
    "    pattern = re.compile(regex_pattern)\n",
    "    matches = pattern.finditer(detokenized_text)\n",
    "    \n",
    "    token_positions = []\n",
    "    for match in matches:\n",
    "        start_char, end_char = match.span()\n",
    "        start_token = char_to_token_position(tokens, start_char)\n",
    "        end_token = char_to_token_position(tokens, end_char)\n",
    "        token_positions.append((start_token, end_token))\n",
    "    return token_positions\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from collections import Counter\n",
    "\n",
    "def prepare_data(dictionary_activations, feature_index, num_bins=10):\n",
    "    best_feature_activations = dictionary_activations[:, feature_index]\n",
    "    min_value = torch.min(best_feature_activations)\n",
    "    max_value = torch.max(best_feature_activations)\n",
    "    bin_boundaries = torch.linspace(min_value, max_value, num_bins+1)\n",
    "    bins = torch.bucketize(best_feature_activations, bin_boundaries)\n",
    "    return best_feature_activations, bin_boundaries, bins\n",
    "\n",
    "def compute_bar_positions(bin_boundaries):\n",
    "    barWidth = (bin_boundaries[1] - bin_boundaries[0]) * 0.8\n",
    "    r = (bin_boundaries[:-1] + bin_boundaries[1:]) / 2\n",
    "    return r, barWidth\n",
    "\n",
    "def process_data_type1(bins, dataset, num_bins, num_unique_tokens, model):\n",
    "    nonzero = bins.nonzero()[:,0]\n",
    "    datapoint_indices = [np.unravel_index(i, (datapoints, token_amount)) for i in nonzero]\n",
    "    max_token = [dataset[int(md)][\"input_ids\"][int(s_ind)] for md, s_ind in datapoint_indices]\n",
    "    \n",
    "    ac = Counter(max_token)\n",
    "    max_tokens = [token_ind for token_ind, count in ac.most_common(num_unique_tokens)]\n",
    "    token_dict = {token: np.zeros(num_bins) for token in max_tokens}\n",
    "    misc_count = np.zeros(num_bins)\n",
    "\n",
    "\n",
    "    for i in range(1, num_bins+1):\n",
    "        bin_tokens_ind = (bins == i).nonzero()[:,0]\n",
    "        datapoint_indices = [np.unravel_index(ind, (datapoints, token_amount)) for ind in bin_tokens_ind]\n",
    "        max_token = [dataset[int(md)][\"input_ids\"][int(s_ind)] for md, s_ind in datapoint_indices]\n",
    "        \n",
    "        total_tokens = len(max_token)\n",
    "        running_count = 0\n",
    "        ac = Counter(max_token)\n",
    "        for token_ind, count in ac.most_common(10):\n",
    "            if token_ind in max_tokens:\n",
    "                token_dict[token_ind][i-1] = count\n",
    "                running_count += count\n",
    "        misc_count[i-1] = total_tokens - running_count\n",
    "    # Covert max_tokens to text\n",
    "    max_text = [model.tokenizer.decode(token) for token in max_tokens]\n",
    "    return token_dict, misc_count, max_tokens, max_text\n",
    "\n",
    "def process_data_type2(bins, dataset, num_bins, masks):\n",
    "    regex_matches_vs_nonmatches = np.zeros((num_bins, 2))\n",
    "    all_non_matched_tokens = []\n",
    "    non_matched_all_tokens = []\n",
    "    for i in range(1,num_bins+1):\n",
    "        bin_tokens_ind = (bins == i).nonzero()[:,0]\n",
    "        datapoint_indices = [np.unravel_index(i, (datapoints, token_amount)) for i in bin_tokens_ind]\n",
    "        for md, s_ind in datapoint_indices:\n",
    "            md = int(md)\n",
    "            s_ind = int(s_ind)\n",
    "            regex_matches_vs_nonmatches[i-1, masks[md][s_ind]] += 1\n",
    "            if not masks[md][s_ind]:\n",
    "                all_non_matched_tokens.append(dataset[md][\"input_ids\"][s_ind])\n",
    "                non_matched_all_tokens.append(dataset[md][\"input_ids\"][:s_ind+1])\n",
    "    return regex_matches_vs_nonmatches, all_non_matched_tokens\n",
    "\n",
    "def plot_data_type1(r, barWidth, bin_boundaries, data, title):\n",
    "    token_dict, misc_count, max_tokens, max_text = data\n",
    "    colors = sns.color_palette('colorblind', len(max_tokens) + 1)\n",
    "    running_sum = np.zeros(len(bin_boundaries) - 1)\n",
    "    for i in range(len(max_tokens)):\n",
    "        token_id_name = max_tokens[i]\n",
    "        token_count_array = token_dict[token_id_name]\n",
    "        text_label = max_text[i]\n",
    "        plt.bar(r, token_count_array, bottom=running_sum, label=text_label, color=colors[i], width=barWidth)\n",
    "        running_sum += token_count_array\n",
    "    plt.bar(r, misc_count, bottom=running_sum, label='[All Other Tokens]', color=colors[-1], width=barWidth)\n",
    "    plt.title(title)\n",
    "    plt.xlabel('Activation')\n",
    "    plt.xticks(bin_boundaries, [f\"{val:.2f}\" for val in bin_boundaries])\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    return plt.gcf()\n",
    "\n",
    "def plot_data_type2(r, barWidth, bin_boundaries, data, regex_pattern):\n",
    "    regex_matches_vs_nonmatches, all_non_matched_tokens = data\n",
    "    colors = sns.color_palette('deep', regex_matches_vs_nonmatches.shape[1] + 1)\n",
    "    running_sum = np.zeros(len(bin_boundaries) - 1)\n",
    "    for i in range(regex_matches_vs_nonmatches.shape[1]):\n",
    "        text_label = regex_pattern if i == 0 else \"non-matches\"\n",
    "        token_count_array = regex_matches_vs_nonmatches[:, (i+1)%2]\n",
    "        plt.bar(r, token_count_array, bottom=running_sum, label=text_label, color=colors[i], width=barWidth)\n",
    "        running_sum += token_count_array\n",
    "    plt.title(f'Token Count for Feature {feature_index} Across Activation Ranges')\n",
    "    plt.xlabel('Activation')\n",
    "    plt.xticks(bin_boundaries, [f\"{val:.2f}\" for val in bin_boundaries])\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "def generate_masks(dataset, model, feature_index, token_amount, regex_pattern, extra_token_amount=2):\n",
    "    \"\"\"\n",
    "    Generates masks based on a regex pattern.\n",
    "\n",
    "    Parameters:\n",
    "    - dataset: dataset containing the texts to be searched\n",
    "    - model: the model containing the tokenizer\n",
    "    - feature_index: index of the feature (not used in the function based on given code, but kept for consistency)\n",
    "    - token_amount: limit to number of tokens to be considered from the tokenized text\n",
    "    - regex_pattern: the regular expression pattern to be searched for\n",
    "    - extra_token_amount: additional tokens to consider\n",
    "\n",
    "    Returns:\n",
    "    - masks: a list of masks generated based on the regex_pattern\n",
    "    \"\"\"\n",
    "    texts = dataset[\"text\"]\n",
    "    masks = []\n",
    "\n",
    "    for text_ind, text in enumerate(texts):\n",
    "        tokens = model.tokenizer(text)\n",
    "        tokens['input_ids'] = tokens['input_ids'][:token_amount + extra_token_amount]\n",
    "        token_positions = find_text_in_regex(tokens, regex_pattern=regex_pattern)\n",
    "        mask = np.zeros(token_amount, dtype=int)  # Initial mask\n",
    "        for start, end in token_positions:\n",
    "            if end is None:\n",
    "                print(\"No match found for text: \", model.tokenizer.decode(tokens.input_ids))\n",
    "                print(start, end)\n",
    "                continue\n",
    "            if end > token_amount:\n",
    "                end = token_amount\n",
    "            mask[start:end] = 1\n",
    "        masks.append(mask)\n",
    "\n",
    "    return masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Downnload dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "from einops import rearrange\n",
    "from datasets import Dataset, load_dataset\n",
    "model_name = \"EleutherAI/pythia-70m-deduped\"\n",
    "device = torch.device(\"cuda:3\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = HookedTransformer.from_pretrained(model_name, device=device)\n",
    "\n",
    "\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": [
    "\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "import matplotlib.pyplot as plt\n",
    "# autoencoder_filename = \"/home/mchorse/logan/sparse_coding/output_sweep_tied_mlpout_l1_r8/_9/learned_dicts.pt\"\n",
    "autoencoder_filename = \"/home/mchorse/sparse_coding_hoagy/tiedlong_tied_residual_l4_r4/_80/learned_dicts.pt\"\n",
    "all_autoencoders = torch.load(autoencoder_filename)\n",
    "num_dictionaries = len(all_autoencoders)\n",
    "print(f\"Loaded {num_dictionaries} autoencoders\")\n",
    "for smaller_autoencoder, hyperparams in all_autoencoders[3:]:\n",
    "    # auto_num = 4 # Selects which specific autoencoder to use\n",
    "    # smaller_autoencoder, hyperparams = all_autoencoders[auto_num]\n",
    "    l1_alpha = hyperparams['l1_alpha']\n",
    "    smaller_dict = smaller_autoencoder.get_learned_dict()\n",
    "    smaller_dict_features, _ = smaller_dict.shape\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",
    "\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\n",
    "\n",
    "    #Dictionary Comparison\n",
    "    # Now we can use the model to get the activations\n",
    "\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_autoencoder.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",
    "            # normal dictionary\n",
    "            batched_dictionary_activations = smaller_autoencoder.encode(batched_neuron_activations)\n",
    "            dictionary_activations[i*batch_size*token_amount:(i+1)*batch_size*token_amount,:] = batched_dictionary_activations.cpu()\n",
    "    # Feature Search\n",
    "    t = \" I don't know about that. It is now up to Dave'\"\n",
    "    #no gradient\n",
    "    with torch.no_grad():\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_autoencoder.encode(neuron_act_batch.squeeze())\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())\n",
    "    dataset_count = 10000\n",
    "    percentage_activated = dictionary_activations[:dataset_count, i.cpu()].count_nonzero(dim=0) / dataset_count\n",
    "    # Grab up to 3 indices, as long as it's not > 30% activated\n",
    "    indices = torch.where(percentage_activated < 0.1)[0][:3]\n",
    "    print(\"Features < 10\\%:\", i[indices].tolist())\n",
    "    # Token level\n",
    "    for feature_index in i[indices]:\n",
    "        num_bins=10\n",
    "        activations, bin_boundaries, bins = prepare_data(dictionary_activations, feature_index)\n",
    "        r, barWidth = compute_bar_positions(bin_boundaries)\n",
    "        data1 = process_data_type1(bins, dataset, num_bins, num_unique_tokens=5, model=model)\n",
    "        title = f\" Token Count for Feature {feature_index} at l1_alpha: {hyperparams['l1_alpha']:.1E}\"\n",
    "        plt_figure = plot_data_type1(r, barWidth, bin_boundaries, data1, title = title)\n",
    "\n",
    "        # Save all data for plots e.g. r, barWidth, bin_boundaries, data1,title\n",
    "        all_variables = (r, barWidth, bin_boundaries, data1,title)\n",
    "        with open(f\"images/{hyperparams['l1_alpha']:.1E}_{feature_index}.pkl\", 'wb') as f:\n",
    "            pickle.dump(all_variables, f)\n",
    "        plt_figure.savefig(f\"images/{hyperparams['l1_alpha']:.1E}_{feature_index}.png\", bbox_inches='tight')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r, barWidth, bin_boundaries, data1,title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_figure = plot_data_type1(r, barWidth, bin_boundaries, data1, model, title = title)\n",
    "plt_figure.savefig(f\"{hyperparams['l1_alpha']:.1E}_{feature_index}.png\", bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_index = 2808\n",
    "# regex_pattern =  r'(?<=\\().*?(?=\\))'\n",
    "regex_pattern = r'\\((?=[A-Z]{1,})'\n",
    "masks = generate_masks(dataset, model, feature_index, token_amount, regex_pattern)\n",
    "\n",
    "activations, bin_boundaries, bins = prepare_data(dictionary_activations, feature_index)\n",
    "r, barWidth = compute_bar_positions(bin_boundaries)\n",
    "data2 = process_data_type2(bins, dataset, num_bins, masks)\n",
    "plot_data_type2(r, barWidth, bin_boundaries, data2, regex_pattern)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regex_nonzero, feature_nonzero = np.array(masks).flatten().nonzero()[0], dictionary_activations[:, feature_index].numpy().nonzero()[0]\n",
    "index_matches = np.intersect1d(regex_nonzero, feature_nonzero)\n",
    "false_positives = np.setdiff1d(regex_nonzero, index_matches)\n",
    "false_negatives = np.setdiff1d(feature_nonzero, index_matches)\n",
    "print(index_matches.shape, false_positives.shape, false_negatives.shape)\n",
    "print(f\"Matches: {index_matches.shape[0]}, False Positives: {false_positives.shape[0]}, False Negatives: {false_negatives.shape[0]}\")\n",
    "found_indices = false_positives[:20]\n",
    "\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",
    "# text_list, full_text, token_list, full_token_list\n",
    "visualize_text(text_list, feature_index, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 104\n",
    "best_feature = int(max_indices[N])\n",
    "# best_feature = 21 # Change this one for global index (N is sorted by MCS)\n",
    "best_feature = 1764 # Feature_ids [1359, 1764, 475, 1815, 513, 1184, 1100, 995, 1001, 365]\n",
    "\n",
    "\n",
    "feature_setting = \"dictionary\"\n",
    "# feature_setting = \"pca\"\n",
    "# feature_setting = \"ica\"\n",
    "# feature_setting = \"neuron\"\n",
    "setting=\"dictionary_basis\"\n",
    "if(feature_setting == \"dictionary\"):\n",
    "    autoencoder = smaller_autoencoder\n",
    "    activations = dictionary_activations\n",
    "elif(feature_setting == \"pca\"):\n",
    "    autoencoder = pca_topk\n",
    "    activations = pca_dictionary_activations\n",
    "elif(feature_setting == \"ica\"):\n",
    "    autoencoder = ica_topk\n",
    "    activations = ica_dictionary_activations\n",
    "elif(feature_setting == \"neuron\"):\n",
    "    autoencoder = smaller_autoencoder # Doesn't matter actually\n",
    "    activations = neuron_activations\n",
    "    setting = \"neuron_basis\"\n",
    "\n",
    "# autoencoder = pca_topk\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, activations, dataset, setting=\"uniform\")\n",
    "# text_list, full_text, token_list, full_token_list = get_feature_datapoints(best_feature, activations, dataset, setting=\"max\")\n",
    "# visualize_text(text_list, best_feature, model, autoencoder, setting=setting)\n",
    "visualize_text(full_text, best_feature, model, autoencoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_text(text_list, best_feature, model, autoencoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ablate_feature_direction_display(full_text, best_feature, autoencoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logit_lens(model,best_feature, smaller_dict, layer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "logan",
   "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
}