Upload folder using huggingface_hub

.github/workflows/update_space.yml

@@ -0,0 +1,28 @@
+name: Run Python script
+
+on:
+  push:
+    branches:
+      - y
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Install Gradio
+      run: python -m pip install gradio
+
+    - name: Log in to Hugging Face
+      run: python -c 'import huggingface_hub; huggingface_hub.login(token="${{ secrets.hf_token }}")'
+
+    - name: Deploy to Spaces
+      run: gradio deploy

.gitignore

@@ -0,0 +1,4 @@
+venv
+*.safetensors
+*.pyc
+.DS_Store

README.md

@@ -1,12 +1,6 @@
 ---
-title: Steer Hexter
-emoji: 🦀
-colorFrom: red
-colorTo: purple
+title: steer-hexter
+app_file: app.py
 sdk: gradio
 sdk_version: 4.38.1
-app_file: app.py
-pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,200 @@
+from transformer_lens import HookedTransformer
+from sae_lens import SAE
+import torch
+
+if torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+class Inference:
+
+    def __init__(self, model, pretrained_sae, layer):
+        self.layer = layer
+        if model == "gemma-2b":
+            self.sae_id = f"blocks.{layer}.hook_resid_post"
+        elif model == "gpt2-small":
+            print(f"using {model}")
+            self.sae_id = f"blocks.{0}.hook_resid_pre"
+        self.sampling_kwargs = dict(temperature=1.0, top_p=0.1, freq_penalty=1.0)
+        self.set_coeff(1)
+        self.set_model(model)
+        self.set_SAE(pretrained_sae)
+ 
+    
+    def set_model(self, model):
+        self.model = HookedTransformer.from_pretrained(model, device = device)
+
+    def set_coeff(self, coeff):
+        self.coeff = coeff
+    
+    def set_temperature(self, temperature):
+        self.sampling_kwargs['temperature'] = temperature
+
+    def set_steering_vector_prompt(self, prompt: str):
+        self.steering_vector_prompt = prompt
+
+    def set_SAE(self, sae_name):
+        sae, cfg_dict, _ = SAE.from_pretrained(
+            release = sae_name,
+            sae_id = self.sae_id,
+            device = device
+        )
+        self.sae = sae
+        self.cfg_dict = cfg_dict
+
+    def _get_sae_out_and_feature_activations(self):
+        # given the words in steering_vectore_prompt, the SAE predicts that the neurons(aka features) in activateCache will be activated
+        sv_logits, activationCache = self.model.run_with_cache(self.steering_vector_prompt, prepend_bos=True)
+        sv_feature_acts = self.sae.encode(activationCache[self.sae.cfg.hook_name])
+        # get top_k of 1
+        # self.sae_out = sae.decode(sv_feature_acts) 
+        return self.sae.decode(sv_feature_acts), sv_feature_acts
+
+    def _hooked_generate(self, prompt_batch, fwd_hooks, seed=None, **kwargs):
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        with self.model.hooks(fwd_hooks=fwd_hooks):
+            tokenized = self.model.to_tokens(prompt_batch)
+            result = self.model.generate(
+                stop_at_eos=False,  # avoids a bug on MPS
+                input=tokenized,
+                max_new_tokens=50,
+                do_sample=True,
+                **kwargs)
+        return result
+
+    def _get_features(self, sv_feature_activations):
+        # return torch.topk(sv_feature_acts, 1).indices.tolist()
+        features = torch.topk(sv_feature_activations, 1).indices
+        print(f'features that align with the text prompt: {features}')
+        print("pump the features into the tool that gives you the words associated with each feature")
+        return features
+
+    
+    def _get_steering_hook(self, feature, sae_out):
+        coeff = self.coeff
+        steering_vector = self.sae.W_dec[feature]
+        steering_vector = steering_vector[0]
+        def steering_hook(resid_pre, hook):
+            if resid_pre.shape[1] == 1:
+                return
+
+            position = sae_out.shape[1]
+            # using our steering vector and applying the coefficient
+            resid_pre[:, :position - 1, :] += coeff * steering_vector
+        
+        return steering_hook
+
+    def _get_steering_hooks(self):
+        # TODO: refactor this. It works because sae_out.shape[1] = sv_feature_acts.shape[1] = len(features[0])
+        # you can manipulate views to retrieve hooks more cleanly
+        # and not use the seperate function _get_steering_hook()
+        sae_out, sv_feature_acts = self._get_sae_out_and_feature_activations()
+        features = self._get_features(sv_feature_acts)
+        steering_hooks = [self._get_steering_hook(feature, sae_out) for feature in features[0]]
+
+        return steering_hooks
+
+
+    def _run_generate(self, example_prompt, steering_on: bool):
+        
+        self.model.reset_hooks()
+        steer_hooks = self._get_steering_hooks()
+        editing_hooks = [ (self.sae_id, steer_hook) for steer_hook in steer_hooks]
+        # editing_hooks = [(self.sae_id, steer_hook)]
+        # ^^change this to support steer_hooks being a list of steer_hooks
+        print(f"steering by {len(editing_hooks)} hooks")
+        if steering_on: 
+            res = self._hooked_generate([example_prompt] * 3, editing_hooks, seed=None, **self.sampling_kwargs)
+        else:
+            tokenized = self.model.to_tokens([example_prompt])
+            res = self.model.generate(
+                stop_at_eos=False,  # avoids a bug on MPS
+                input=tokenized,
+                max_new_tokens=50,
+                do_sample=True,
+                **self.sampling_kwargs)
+            
+        # Print results, removing the ugly beginning of sequence token
+        res_str = self.model.to_string(res[:, 1:])
+        response = ("\n\n" + "-" * 80 + "\n\n").join(res_str)
+        print(response)
+        return response
+
+
+    def generate(self, message: str, steering_on: bool):
+        return self._run_generate(message, steering_on)
+
+
+
+MODEL = "gemma-2b"
+PRETRAINED_SAE = "gemma-2b-res-jb"
+MODEL = "gpt2-small"
+PRETRAINED_SAE = "gpt2-small-res-jb"
+LAYER = 10
+chatbot_model = Inference(MODEL,PRETRAINED_SAE, LAYER)
+
+
+import time
+import gradio as gr
+
+default_image = "Hexter-Hackathon.png"
+
+def slow_echo(message, history):
+    result = chatbot_model.generate(message, False)
+    for i in range(len(result)):
+        time.sleep(0.01)
+        yield result[: i + 1]
+def slow_echo_steering(message, history):
+    result = chatbot_model.generate(message, True)
+    for i in range(len(result)):
+        time.sleep(0.01)
+        yield result[: i + 1]
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        gr.Markdown("*STANDARD HEXTER BOT*")
+    with gr.Row():
+        chatbot = gr.ChatInterface(
+            slow_echo,
+            chatbot=gr.Chatbot(min_width=1000),
+            textbox=gr.Textbox(placeholder="Ask Hexter anything!", min_width=1000),
+            theme="soft",
+            cache_examples=False,
+            retry_btn=None,
+            clear_btn=None,
+            undo_btn=None,
+        )
+    with gr.Row():
+        gr.Markdown("*STEERED HEXTER BOT*")
+    with gr.Row():
+        chatbot_steered = gr.ChatInterface(
+            slow_echo_steering,
+            chatbot=gr.Chatbot(min_width=1000),
+            textbox=gr.Textbox(placeholder="Ask Hexter anything!", min_width=1000),
+            theme="soft",
+            cache_examples=False,
+            retry_btn=None,
+            clear_btn=None,
+            undo_btn=None,
+        )
+    with gr.Row():
+        steering_prompt = gr.Textbox(label="Steering prompt", value="Golden Gate Bridge")
+    with gr.Row():
+        coeff = gr.Slider(1, 1000, 300, label="Coefficient", info="Coefficient is..", interactive=True)
+    with gr.Row():
+        temp = gr.Slider(0, 5, 1, label="Temperature", info="Temperature is..", interactive=True)
+
+    # Set up an action when the sliders change
+    temp.change(chatbot_model.set_temperature, inputs=[temp], outputs=[])
+    coeff.change(chatbot_model.set_coeff, inputs=[coeff], outputs=[])
+    chatbot_model.set_steering_vector_prompt(steering_prompt)
+    steering_prompt.change(chatbot_model.set_steering_vector_prompt, inputs=[steering_prompt], outputs=[])
+
+demo.queue()
+demo.launch(debug=True)
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=["/"])

loading_analyzing.py

@@ -0,0 +1,266 @@
+
+COLAB = False
+# from IPython import get_ipython # type: ignore
+# ipython = get_ipython(); assert ipython is not None
+# ipython.run_line_magic("load_ext", "autoreload")
+# ipython.run_line_magic("autoreload", "2")
+
+
+# Standard imports
+import os
+import torch
+from tqdm import tqdm
+import plotly.express as px
+
+# Imports for displaying vis in Colab / notebook
+import webbrowser
+import http.server
+import socketserver
+import threading
+PORT = 8000
+
+torch.set_grad_enabled(False);
+
+# For the most part I'll try to import functions and classes near where they are used
+# to make it clear where they come from.
+
+if torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(f"Device: {device}")
+
+
+def display_vis_inline(filename: str, height: int = 850):
+    '''
+    Displays the HTML files in Colab. Uses global `PORT` variable defined in prev cell, so that each
+    vis has a unique port without having to define a port within the function.
+    '''
+    if not(COLAB):
+        webbrowser.open(filename);
+
+    else:
+        global PORT
+
+        def serve(directory):
+            os.chdir(directory)
+
+            # Create a handler for serving files
+            handler = http.server.SimpleHTTPRequestHandler
+
+            # Create a socket server with the handler
+            with socketserver.TCPServer(("", PORT), handler) as httpd:
+                print(f"Serving files from {directory} on port {PORT}")
+                httpd.serve_forever()
+
+        thread = threading.Thread(target=serve, args=("/content",))
+        thread.start()
+
+        # output.serve_kernel_port_as_iframe(PORT, path=f"/{filename}", height=height, cache_in_notebook=True)
+
+        PORT += 1
+
+
+    
+
+from datasets import load_dataset
+from transformer_lens import HookedTransformer
+from sae_lens import SAE
+
+model = HookedTransformer.from_pretrained("gpt2-small", device = device)
+
+# the cfg dict is returned alongside the SAE since it may contain useful information for analysing the SAE (eg: instantiating an activation store)
+# Note that this is not the same as the SAEs config dict, rather it is whatever was in the HF repo, from which we can extract the SAE config dict
+# We also return the feature sparsities which are stored in HF for convenience.
+sae, cfg_dict, sparsity = SAE.from_pretrained(
+    release = "gpt2-small-res-jb", # see other options in sae_lens/pretrained_saes.yaml
+    sae_id = "blocks.8.hook_resid_pre", # won't always be a hook point
+    device = device
+)
+
+
+from transformer_lens.utils import tokenize_and_concatenate
+
+dataset = load_dataset(
+    path = "NeelNanda/pile-10k",
+    split="train",
+    streaming=False,
+)
+
+token_dataset = tokenize_and_concatenate(
+    dataset= dataset,# type: ignore
+    tokenizer = model.tokenizer, # type: ignore
+    streaming=True,
+    max_length=sae.cfg.context_size,
+    add_bos_token=sae.cfg.prepend_bos,
+)
+
+
+sae.eval()  # prevents error if we're expecting a dead neuron mask for who grads
+
+with torch.no_grad():
+    # activation store can give us tokens.
+    batch_tokens = token_dataset[:32]["tokens"]
+    _, cache = model.run_with_cache(batch_tokens, prepend_bos=True)
+
+    # Use the SAE
+    feature_acts = sae.encode(cache[sae.cfg.hook_name])
+    sae_out = sae.decode(feature_acts)
+
+    # save some room
+    del cache
+
+    # ignore the bos token, get the number of features that activated in each token, averaged accross batch and position
+    l0 = (feature_acts[:, 1:] > 0).float().sum(-1).detach()
+    print("average l0", l0.mean().item())
+    px.histogram(l0.flatten().cpu().numpy()).show()
+
+
+
+
+from transformer_lens import utils
+from functools import partial
+
+# next we want to do a reconstruction test.
+def reconstr_hook(activation, hook, sae_out):
+    return sae_out
+
+
+def zero_abl_hook(activation, hook):
+    return torch.zeros_like(activation)
+
+
+print("Orig", model(batch_tokens, return_type="loss").item())
+print(
+    "reconstr",
+    model.run_with_hooks(
+        batch_tokens,
+        fwd_hooks=[
+            (
+                sae.cfg.hook_name,
+                partial(reconstr_hook, sae_out=sae_out),
+            )
+        ],
+        return_type="loss",
+    ).item(),
+)
+print(
+    "Zero",
+    model.run_with_hooks(
+        batch_tokens,
+        return_type="loss",
+        fwd_hooks=[(sae.cfg.hook_name, zero_abl_hook)],
+    ).item(),
+)
+
+
+
+
+example_prompt = "When John and Mary went to the shops, John gave the bag to"
+example_answer = " Mary"
+utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)
+
+logits, cache = model.run_with_cache(example_prompt, prepend_bos=True)
+tokens = model.to_tokens(example_prompt)
+sae_out = sae(cache[sae.cfg.hook_name])
+
+
+def reconstr_hook(activations, hook, sae_out):
+    return sae_out
+
+
+def zero_abl_hook(mlp_out, hook):
+    return torch.zeros_like(mlp_out)
+
+
+hook_name = sae.cfg.hook_name
+
+print("Orig", model(tokens, return_type="loss").item())
+print(
+    "reconstr",
+    model.run_with_hooks(
+        tokens,
+        fwd_hooks=[
+            (
+                hook_name,
+                partial(reconstr_hook, sae_out=sae_out),
+            )
+        ],
+        return_type="loss",
+    ).item(),
+)
+print(
+    "Zero",
+    model.run_with_hooks(
+        tokens,
+        return_type="loss",
+        fwd_hooks=[(hook_name, zero_abl_hook)],
+    ).item(),
+)
+
+
+with model.hooks(
+    fwd_hooks=[
+        (
+            hook_name,
+            partial(reconstr_hook, sae_out=sae_out),
+        )
+    ]
+):
+    utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)
+
+
+
+
+from sae_dashboard.sae_vis_data import SaeVisConfig
+from sae_dashboard.sae_vis_runner import SaeVisRunner
+
+test_feature_idx_gpt = list(range(10)) + [14057]
+
+feature_vis_config_gpt = SaeVisConfig(
+    hook_point=hook_name,
+    features=test_feature_idx_gpt,
+    minibatch_size_features=64,
+    minibatch_size_tokens=256,
+    verbose=True,
+    device=device,
+)
+
+visualization_data_gpt = SaeVisRunner(feature_vis_config_gpt).run(
+    encoder=sae, # type: ignore
+    model=model,
+    tokens=token_dataset[:10000]["tokens"],  # type: ignore
+)
+# SaeVisData.create(
+#     encoder=sae,
+#     model=model, # type: ignore
+#     tokens=token_dataset[:10000]["tokens"],  # type: ignore
+#     cfg=feature_vis_config_gpt,
+# )
+
+
+
+
+
+from sae_dashboard.data_writing_fns import save_feature_centric_vis
+
+filename = f"demo_feature_dashboards.html"
+save_feature_centric_vis(sae_vis_data=visualization_data_gpt, filename=filename)
+
+
+
+from sae_lens.analysis.neuronpedia_integration import get_neuronpedia_quick_list
+
+# this function should open
+neuronpedia_quick_list = get_neuronpedia_quick_list(
+    test_feature_idx_gpt,
+    layer=sae.cfg.hook_layer,
+    model="gpt2-small",
+    dataset="res-jb",
+    name="A quick list we made",
+)
+
+if COLAB:
+  # If you're on colab, click the link below
+  print(neuronpedia_quick_list)

sae_tiny-stories-1L-21M_blocks.0.hook_mlp_out_16384/cfg.json

@@ -0,0 +1 @@
+{"model_name": "tiny-stories-1L-21M", "model_class_name": "HookedTransformer", "hook_name": "blocks.0.hook_mlp_out", "hook_eval": "NOT_IN_USE", "hook_layer": 0, "hook_head_index": null, "dataset_path": "apollo-research/roneneldan-TinyStories-tokenizer-gpt2", "dataset_trust_remote_code": true, "streaming": true, "is_dataset_tokenized": true, "context_size": 512, "use_cached_activations": false, "cached_activations_path": null, "architecture": "standard", "d_in": 1024, "d_sae": 16384, "b_dec_init_method": "zeros", "expansion_factor": 16, "activation_fn": "relu", "activation_fn_kwargs": {}, "normalize_sae_decoder": false, "noise_scale": 0.0, "from_pretrained_path": null, "apply_b_dec_to_input": false, "decoder_orthogonal_init": false, "decoder_heuristic_init": true, "init_encoder_as_decoder_transpose": true, "n_batches_in_buffer": 64, "training_tokens": 122880000, "finetuning_tokens": 0, "store_batch_size_prompts": 16, "train_batch_size_tokens": 4096, "normalize_activations": "expected_average_only_in", "device": "cuda", "act_store_device": "cuda", "seed": 42, "dtype": "float32", "prepend_bos": true, "autocast": false, "autocast_lm": false, "compile_llm": false, "llm_compilation_mode": null, "compile_sae": false, "sae_compilation_mode": null, "adam_beta1": 0.9, "adam_beta2": 0.999, "mse_loss_normalization": null, "l1_coefficient": 5, "lp_norm": 1.0, "scale_sparsity_penalty_by_decoder_norm": true, "l1_warm_up_steps": 1500, "lr": 5e-05, "lr_scheduler_name": "constant", "lr_warm_up_steps": 0, "lr_end": 5e-06, "lr_decay_steps": 6000, "n_restart_cycles": 1, "finetuning_method": null, "use_ghost_grads": false, "feature_sampling_window": 1000, "dead_feature_window": 1000, "dead_feature_threshold": 0.0001, "n_eval_batches": 10, "eval_batch_size_prompts": null, "log_to_wandb": true, "log_activations_store_to_wandb": false, "log_optimizer_state_to_wandb": false, "wandb_project": "sae_lens_tutorial", "wandb_id": null, "run_name": "16384-L1-5-LR-5e-05-Tokens-1.229e+08", "wandb_entity": null, "wandb_log_frequency": 30, "eval_every_n_wandb_logs": 20, "resume": false, "n_checkpoints": 0, "checkpoint_path": "checkpoints/q5ut4uqw", "verbose": true, "model_kwargs": {}, "model_from_pretrained_kwargs": {}, "sae_lens_version": "3.13.0", "sae_lens_training_version": "3.13.0", "tokens_per_buffer": 134217728}

steer_stories.py

@@ -0,0 +1,150 @@
+# general imports
+import os
+import torch
+from tqdm import tqdm
+import plotly.express as px
+
+torch.set_grad_enabled(False);
+
+
+# package import
+from torch import Tensor
+from transformer_lens import utils
+from functools import partial
+from jaxtyping import Int, Float
+
+# device setup
+if torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(f"Device: {device}")
+
+
+
+
+
+from transformer_lens import HookedTransformer
+from sae_lens import SAE
+
+# Choose a layer you want to focus on
+# For this tutorial, we're going to use layer ????
+layer = 0
+
+# get model
+model = HookedTransformer.from_pretrained("tiny-stories-1L-21M", device = device)
+
+# get the SAE for this layer
+sae = SAE.load_from_pretrained("sae_tiny-stories-1L-21M_blocks.0.hook_mlp_out_16384", device = device)
+
+# get hook point
+hook_point = sae.cfg.hook_name
+print(hook_point)
+
+
+
+sv_prompt = " Lily"
+sv_logits, activationCache = model.run_with_cache(sv_prompt, prepend_bos=True)
+sv_feature_acts = sae.encode(activationCache[hook_point])
+print(torch.topk(sv_feature_acts, 3).indices.tolist())
+
+# Generate
+
+
+sv_prompt = " Lily"
+sv_logits, activationCache = model.run_with_cache(sv_prompt, prepend_bos=True)
+tokens = model.to_tokens(sv_prompt)
+print(tokens)
+
+# get the feature activations from our SAE
+sv_feature_acts = sae.encode(activationCache[hook_point])
+
+# get sae_out
+sae_out = sae.decode(sv_feature_acts)
+
+# print out the top activations, focus on the indices
+print(torch.topk(sv_feature_acts, 3))
+
+
+
+# get the neurons to use;
+print(torch.topk(sv_feature_acts, 3).indices.tolist())
+
+
+# choose the vector -- find this from the above section
+# 
+steering_vector = sae.W_dec[10284]
+
+example_prompt = "Once upon a time"
+coeff = 1000
+sampling_kwargs = dict(temperature=1.0, top_p=0.1, freq_penalty=1.0)
+
+
+
+# apply steering vector when the model generates
+
+def steering_hook(resid_pre, hook):
+    if resid_pre.shape[1] == 1:
+        return
+
+    position = sae_out.shape[1]
+    if steering_on:
+      breakpoint()
+      # using our steering vector and applying the coefficient
+      resid_pre[:, :position - 1, :] += coeff * steering_vector
+
+
+def hooked_generate(prompt_batch, fwd_hooks=[], seed=None, **kwargs):
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    with model.hooks(fwd_hooks=fwd_hooks):
+        tokenized = model.to_tokens(prompt_batch)
+        result = model.generate(
+            stop_at_eos=False,  # avoids a bug on MPS
+            input=tokenized,
+            max_new_tokens=50,
+            do_sample=True,
+            **kwargs)
+    return result
+
+
+def run_generate(example_prompt):
+  model.reset_hooks()
+  editing_hooks = [(f"blocks.{layer}.hook_resid_post", steering_hook)]
+  res = hooked_generate([example_prompt] * 3, editing_hooks, seed=None, **sampling_kwargs)
+
+  # Print results, removing the ugly beginning of sequence token
+  res_str = model.to_string(res[:, 1:])
+  print(("\n\n" + "-" * 80 + "\n\n").join(res_str))
+
+
+steering_on = True
+run_generate(example_prompt)
+
+
+
+# evaluate features
+
+import pandas as pd
+
+# Let's start by getting the top 10 logits for each feature
+projection_onto_unembed = sae.W_dec @ model.W_U
+
+
+# get the top 10 logits.
+vals, inds = torch.topk(projection_onto_unembed, 10, dim=1)
+
+# get 10 random features
+random_indices = torch.randint(0, projection_onto_unembed.shape[0], (10,))
+
+# Show the top 10 logits promoted by those features
+top_10_logits_df = pd.DataFrame(
+    [model.to_str_tokens(i) for i in inds[random_indices]],
+    index=random_indices.tolist(),
+).T
+top_10_logits_df
+# [7195, 5910, 2041]
+top_10_associated_words_logits_df = model.to_str_tokens(inds[5910])
+# See the words associated with feature 7195 (Should be "Golden")

using_an_sae_as_a_steering_vector.ipynb

@@ -0,0 +1,2171 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GoXn14ltnGh3"
+      },
+      "source": [
+        "# Using an SAE as a steering vector\n",
+        "\n",
+        "This notebook demonstrates how to use SAE lens to identify a feature on a pretrained model, and then construct a steering vector to affect the models output to various prompts. This notebook will also make use of Neuronpedia for identifying features of interest.\n",
+        "\n",
+        "The steps below include:\n",
+        "\n",
+        "\n",
+        "\n",
+        "*   Installing relevant packages (Colab or locally)\n",
+        "*   Load your SAE and the model it used\n",
+        "*   Determining your feature of interest and its index\n",
+        "*   Implementing your steering vector\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gf3lJYPEXh0v"
+      },
+      "source": [
+        "## Setting up packages and notebook"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "l9k5iGyOXtuN"
+      },
+      "source": [
+        "### Import and installs"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fapxk8MDrs6R"
+      },
+      "source": [
+        "#### Environment Setup\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "collapsed": true,
+        "id": "0TwNmRkRUgR7",
+        "outputId": "ffeb827a-9af2-4b09-b8dd-78e0d594ddf6"
+      },
+      "outputs": [],
+      "source": [
+        "try:\n",
+        "  # for google colab users\n",
+        "    import google.colab # type: ignore\n",
+        "    from google.colab import output\n",
+        "    COLAB = True\n",
+        "    %pip install sae-lens transformer-lens\n",
+        "except:\n",
+        "  # for local setup\n",
+        "    COLAB = False\n",
+        "    from IPython import get_ipython # type: ignore\n",
+        "    ipython = get_ipython(); assert ipython is not None\n",
+        "    ipython.run_line_magic(\"load_ext\", \"autoreload\")\n",
+        "    ipython.run_line_magic(\"autoreload\", \"2\")\n",
+        "\n",
+        "# Imports for displaying vis in Colab / notebook\n",
+        "import webbrowser\n",
+        "import http.server\n",
+        "import socketserver\n",
+        "import threading\n",
+        "PORT = 8000\n",
+        "\n",
+        "# general imports\n",
+        "import os\n",
+        "import torch\n",
+        "from tqdm import tqdm\n",
+        "import plotly.express as px\n",
+        "\n",
+        "torch.set_grad_enabled(False);"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "NGgIu1ZVYDub"
+      },
+      "outputs": [],
+      "source": [
+        "def display_vis_inline(filename: str, height: int = 850):\n",
+        "    '''\n",
+        "    Displays the HTML files in Colab. Uses global `PORT` variable defined in prev cell, so that each\n",
+        "    vis has a unique port without having to define a port within the function.\n",
+        "    '''\n",
+        "    if not(COLAB):\n",
+        "        webbrowser.open(filename);\n",
+        "\n",
+        "    else:\n",
+        "        global PORT\n",
+        "\n",
+        "        def serve(directory):\n",
+        "            os.chdir(directory)\n",
+        "\n",
+        "            # Create a handler for serving files\n",
+        "            handler = http.server.SimpleHTTPRequestHandler\n",
+        "\n",
+        "            # Create a socket server with the handler\n",
+        "            with socketserver.TCPServer((\"\", PORT), handler) as httpd:\n",
+        "                print(f\"Serving files from {directory} on port {PORT}\")\n",
+        "                httpd.serve_forever()\n",
+        "\n",
+        "        thread = threading.Thread(target=serve, args=(\"/content\",))\n",
+        "        thread.start()\n",
+        "\n",
+        "        output.serve_kernel_port_as_iframe(PORT, path=f\"/{filename}\", height=height, cache_in_notebook=True)\n",
+        "\n",
+        "        PORT += 1"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CmaPYLpGrxbo"
+      },
+      "source": [
+        "#### General Installs and device setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "tdUm9rZKr1Qb",
+        "outputId": "9b73b762-1356-437b-8925-91c514093b43"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Device: mps\n"
+          ]
+        }
+      ],
+      "source": [
+        "# package import\n",
+        "from torch import Tensor\n",
+        "from transformer_lens import utils\n",
+        "from functools import partial\n",
+        "from jaxtyping import Int, Float\n",
+        "\n",
+        "# device setup\n",
+        "if torch.backends.mps.is_available():\n",
+        "    device = \"mps\"\n",
+        "else:\n",
+        "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+        "\n",
+        "print(f\"Device: {device}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "lsB0qORUaXiK"
+      },
+      "source": [
+        "### Load your model and SAE\n",
+        "\n",
+        "We're going to work with a pretrained GPT2-small model, and the RES-JB SAE set which is for the residual stream."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 40,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "collapsed": true,
+        "id": "bCvNtm1OOhlR",
+        "outputId": "e6fd27ab-ee94-46ec-a07e-ee48c8f30da3"
+      },
+      "outputs": [
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "8607cfc3f17548078c7b3ff7ebcca055",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "WARNING:root:You are not using LayerNorm, so the writing weights can't be centered! Skipping\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Loaded pretrained model gemma-2b into HookedTransformer\n",
+            "blocks.6.hook_resid_post\n"
+          ]
+        }
+      ],
+      "source": [
+        "from transformer_lens import HookedTransformer\n",
+        "from sae_lens import SAE\n",
+        "from sae_lens.toolkit.pretrained_saes import get_gpt2_res_jb_saes\n",
+        "\n",
+        "# Choose a layer you want to focus on\n",
+        "# For this tutorial, we're going to use layer 2\n",
+        "layer = 6\n",
+        "\n",
+        "# get model\n",
+        "model = HookedTransformer.from_pretrained(\"gemma-2b\", device = device)\n",
+        "\n",
+        "# get the SAE for this layer\n",
+        "sae, cfg_dict, _ = SAE.from_pretrained(\n",
+        "    release = \"gemma-2b-res-jb\",\n",
+        "    sae_id = f\"blocks.{layer}.hook_resid_post\",\n",
+        "    device = device\n",
+        ")\n",
+        "\n",
+        "# get hook point\n",
+        "hook_point = sae.cfg.hook_name\n",
+        "print(hook_point)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NkAAoyFbu5a5"
+      },
+      "source": [
+        "## Determine your feature of interest and its index"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DkQNvdd54q4S"
+      },
+      "source": [
+        "### Find your feature"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wzeY2D13xRjY"
+      },
+      "source": [
+        "#### Explore through code by using the feature activations for a prompt\n",
+        "\n",
+        "For the purpose of the tutorial, we are selecting a simple token prompt.\n",
+        "\n",
+        "In this example we will look trying to find and steer a \"Jedi\" feature.\n",
+        "\n",
+        "We run our prompt on our model and get the cache, which we then use with our sae to get our feature activations.\n",
+        "\n",
+        "Now we'll look at the top feature activations and look them up on Neuronpedia to determine what they have been intepreted as."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 41,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "IIrdJ36mlXgB",
+        "outputId": "c4014b87-3af6-4c27-8f79-3b5a3c2c03dc"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "tensor([[    2,   714, 17489, 22352, 16125]], device='mps:0')\n",
+            "torch.return_types.topk(\n",
+            "values=tensor([[[72.5067, 70.9109, 68.8217],\n",
+            "         [37.8954, 31.1813, 15.6114],\n",
+            "         [65.9133, 14.2098, 13.3081],\n",
+            "         [22.8078, 21.7340, 17.3972],\n",
+            "         [43.6480, 12.6739, 10.7545]]], device='mps:0'),\n",
+            "indices=tensor([[[ 3390, 15881,  5347],\n",
+            "         [ 6518, 13743,  1959],\n",
+            "         [ 1571, 12529, 15173],\n",
+            "         [12773, 10200, 15173],\n",
+            "         [ 5192, 15173, 12030]]], device='mps:0'))\n"
+          ]
+        }
+      ],
+      "source": [
+        "sv_prompt = \" The Golden Gate Bridge\"\n",
+        "sv_logits, cache = model.run_with_cache(sv_prompt, prepend_bos=True)\n",
+        "tokens = model.to_tokens(sv_prompt)\n",
+        "print(tokens)\n",
+        "\n",
+        "# get the feature activations from our SAE\n",
+        "sv_feature_acts = sae.encode(cache[hook_point])\n",
+        "\n",
+        "# get sae_out\n",
+        "sae_out = sae.decode(sv_feature_acts)\n",
+        "\n",
+        "# print out the top activations, focus on the indices\n",
+        "print(torch.topk(sv_feature_acts, 3))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 16,
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
+            "To disable this warning, you can either:\n",
+            "\t- Avoid using `tokenizers` before the fork if possible\n",
+            "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
+          ]
+        },
+        {
+          "data": {
+            "text/plain": [
+              "'https://neuronpedia.org/quick-list/?name=temporary_list&features=%5B%7B%22modelId%22%3A%20%22gemma-2b%22%2C%20%22layer%22%3A%20%220-res-jb%22%2C%20%22index%22%3A%20%22%5B%5B9036%2C%20347%2C%205775%5D%2C%20%5B12244%2C%208086%2C%2015895%5D%5D%22%7D%5D'"
+            ]
+          },
+          "execution_count": 16,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "from sae_lens.analysis.neuronpedia_integration import get_neuronpedia_quick_list\n",
+        "get_neuronpedia_quick_list(torch.topk(sv_feature_acts, 3).indices.tolist(), layer = layer, model = \"gemma-2b\", dataset=\"res-jb\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7hy8RbbyTb8n"
+      },
+      "source": [
+        "As we can see from our print out of tokens, the prompt is made of three tokens in total - \"<endoftext>\", \"J\", and \"edi\".\n",
+        "\n",
+        "Our feature activation indexes at sv_feature_acts[2] - for \"edi\" - are of most interest to us.\n",
+        "\n",
+        "Because we are using pretrained saes that have published feature maps, you can search on Neuronpedia for a feature of interest."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gFv4iBHFcOmE"
+      },
+      "source": [
+        "### Steps for Neuronpedia use\n",
+        "\n",
+        "Use the interface to search for a specific concept or item and determine which layer and at what index it is.\n",
+        "\n",
+        "1.   Open the [Neuronpedia](https://www.neuronpedia.org/) homepage.\n",
+        "2.   Using the \"Models\" dropdown, select your model. Here we are using GPT2-SM (GPT2-small).\n",
+        "3.   The next page will have a search bar, which allows you to enter your index of interest. We're interested in the \"RES-JB\" SAE set, make sure to select it.\n",
+        "4.   We found these indices in the previous step: [ 7650,   718, 22372]. Select them in the search to see the feature dashboard for each.\n",
+        "5.   As we'll see, some of the indices may relate to features you don't care about.\n",
+        "\n",
+        "From using Neuronpedia, I have determined that my feature of interest is in layer 2, at index 7650: [here](https://www.neuronpedia.org/gpt2-small/2-res-jb/7650) is the feature."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KX0rXziniH9O"
+      },
+      "source": [
+        "### Note: 2nd Option - Starting with Neuronpedia\n",
+        "\n",
+        "Another option here is that you can start with Neuronpedia to identify features of interest. By using your prompt in the interface you can explore which features were involved and search across all the layers. This allows you to first determine your layer and index of interest in Neuronpedia before focusing them in your code. Start [here](https://www.neuronpedia.org/search) if you want to begin with search."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "YACtNFzGcNua"
+      },
+      "source": [
+        "## Implement your steering vector and affect the output"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "pO8hjg8j5bb-"
+      },
+      "source": [
+        "### Define values for your steering vector\n",
+        "To create our steering vector, we now need to get the decoder weights from our sparse autoencoder found at our index of interest.\n",
+        "\n",
+        "Then to use our steering vector, we want a prompt for text generation, as well as a scaling factor coefficent to apply with the steering vector\n",
+        "\n",
+        "We also set common sampling kwargs - temperature, top_p and freq_penalty"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 46,
+      "metadata": {
+        "id": "rgYEWGV0t0L2"
+      },
+      "outputs": [],
+      "source": [
+        "steering_vector = sae.W_dec[10200]\n",
+        "\n",
+        "example_prompt = \"What is the most iconic structure known to man?\"\n",
+        "coeff = 300\n",
+        "sampling_kwargs = dict(temperature=1.0, top_p=0.1, freq_penalty=1.0)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cexaoBR65lIa"
+      },
+      "source": [
+        "### Set up hook functions\n",
+        "\n",
+        "Finally, we need to create a hook that allows us to apply the steering vector when our model runs generate() on our defined prompt. We have also added a boolean value 'steering_on' that allows us to easily toggle the steering vector on and off for each prompt\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 47,
+      "metadata": {
+        "collapsed": true,
+        "id": "3kcVWeJoIAlC"
+      },
+      "outputs": [],
+      "source": [
+        "def steering_hook(resid_pre, hook):\n",
+        "    if resid_pre.shape[1] == 1:\n",
+        "        return\n",
+        "\n",
+        "    position = sae_out.shape[1]\n",
+        "    if steering_on:\n",
+        "      # using our steering vector and applying the coefficient\n",
+        "      resid_pre[:, :position - 1, :] += coeff * steering_vector\n",
+        "\n",
+        "\n",
+        "def hooked_generate(prompt_batch, fwd_hooks=[], seed=None, **kwargs):\n",
+        "    if seed is not None:\n",
+        "        torch.manual_seed(seed)\n",
+        "\n",
+        "    with model.hooks(fwd_hooks=fwd_hooks):\n",
+        "        tokenized = model.to_tokens(prompt_batch)\n",
+        "        result = model.generate(\n",
+        "            stop_at_eos=False,  # avoids a bug on MPS\n",
+        "            input=tokenized,\n",
+        "            max_new_tokens=50,\n",
+        "            do_sample=True,\n",
+        "            **kwargs)\n",
+        "    return result\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 48,
+      "metadata": {
+        "id": "VcuRkX0yA2WH"
+      },
+      "outputs": [],
+      "source": [
+        "def run_generate(example_prompt):\n",
+        "  model.reset_hooks()\n",
+        "  editing_hooks = [(f\"blocks.{layer}.hook_resid_post\", steering_hook)]\n",
+        "  res = hooked_generate([example_prompt] * 3, editing_hooks, seed=None, **sampling_kwargs)\n",
+        "\n",
+        "  # Print results, removing the ugly beginning of sequence token\n",
+        "  res_str = model.to_string(res[:, 1:])\n",
+        "  print((\"\\n\\n\" + \"-\" * 80 + \"\\n\\n\").join(res_str))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XYx--hIn61VQ"
+      },
+      "source": [
+        "### Generate text influenced by steering vector\n",
+        "\n",
+        "You may want to experiment with the scaling factor coefficient value that you set and see how it affects the generated output."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 49,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 337,
+          "referenced_widgets": [
+            "9f555c5ada38495eb4281cbb49169abe",
+            "79b59cbde9444bf892931d31afec7f2a",
+            "a157870318114d459a33d795850967ef",
+            "635162e10abc441797d4e5b74713bf44",
+            "720b4d010c364e3fbf72a53b267e8db9",
+            "d9c33fbfb3164cbbb7b9a4cd172d20ae",
+            "df53331cce124bd1ada5aa9e9a977015",
+            "229dad8e29f04c279c5603286e2c0643",
+            "83d947fc3338491ab4155b87c443884c",
+            "5e9700580d6b4ad0bfac34bf3b3919fc",
+            "a2c30462ef8d41fd9158f194a746d5a7"
+          ]
+        },
+        "id": "hN_YOzBE6lz8",
+        "outputId": "e263b8ff-86ce-439e-81e5-bbecb0d7e187"
+      },
+      "outputs": [
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "634ddfad68cb49208e63733402859842",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "  0%|          | 0/50 [00:00<?, ?it/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "What is the most iconic structure known to man? The Golden Gate Bridge. The Golden Gate Bridge is a cable- suspension bridge that connects San Francisco and the Presidio of San Francisco, connecting the northern part of San Francisco with the northern part of Marin County. It’s one of the most famous\n",
+            "\n",
+            "--------------------------------------------------------------------------------\n",
+            "\n",
+            "What is the most iconic structure known to man? The Golden Gate Bridge. The Golden Gate Bridge is a cable- suspension bridge that connects San Francisco and the northern part of San Francisco’s famous “Golden” Coast. It’s one of the most iconic landmarks in the whole world, and it\n",
+            "\n",
+            "--------------------------------------------------------------------------------\n",
+            "\n",
+            "What is the most iconic structure known to man? The Golden Gate Bridge. The Golden Gate Bridge is a cable- suspension bridge that connects San Francisco and the Presidio of San Francisco, connecting the northern part of San Francisco with the city’s northernmost part. It’s one of the most\n"
+          ]
+        }
+      ],
+      "source": [
+        "steering_on = True\n",
+        "run_generate(example_prompt)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ltZEm1VW7Tsr"
+      },
+      "source": [
+        "### Generate text with no steering"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 50,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 337,
+          "referenced_widgets": [
+            "1a5dd5f7c9d340b6ab00ecaf43525ae9",
+            "8211cc6c973a43fcaf18e14f6d7f08a2",
+            "3d3584d1feec459287ffa24c4ef790c3",
+            "5f03835168e64ec588c50ee21fedd198",
+            "b833db18729f422cb86deed4be6f1900",
+            "66d406d6eb1f49699ee09c9a2fd4ffa9",
+            "38341454dd6b4e9ca2fe5b85d2e371e1",
+            "a30c82833f55441995744300c2ef538d",
+            "4932983d4f1a4199b3d24c730c765a24",
+            "c20e9e14100d45f3bdff1b6df943940f",
+            "5c53f97287d54c03a378fc44ab791cd7"
+          ]
+        },
+        "collapsed": true,
+        "id": "nA9cs1BY7XhS",
+        "outputId": "22a03d47-1afb-4217-d77a-979c94392f2a"
+      },
+      "outputs": [
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "f581d3a3b4d44a5e92dae704116e4445",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "  0%|          | 0/50 [00:00<?, ?it/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "What is the most iconic structure known to man? The Eiffel Tower, of course!\n",
+            "\n",
+            "The Eiffel Tower is a symbol of Paris and France. It was built in 1889 for the World’s Fair and has been a popular tourist attraction ever since.\n",
+            "\n",
+            "The tower stands at \n",
+            "\n",
+            "--------------------------------------------------------------------------------\n",
+            "\n",
+            "What is the most iconic structure known to man? The Eiffel Tower, of course!\n",
+            "\n",
+            "The Eiffel Tower is a wrought iron lattice tower located in Paris, France. It was built for the 1889 World’s Fair and has become one of the most recognizable symbols of Paris.\n",
+            "\n",
+            "\n",
+            "\n",
+            "--------------------------------------------------------------------------------\n",
+            "\n",
+            "What is the most iconic structure known to man? The Eiffel Tower, of course!\n",
+            "\n",
+            "The Eiffel Tower is a symbol of Paris and France. It was built in 1889 for the World’s Fair and has been a popular tourist attraction ever since.\n",
+            "\n",
+            "The tower stands at \n"
+          ]
+        }
+      ],
+      "source": [
+        "steering_on = False\n",
+        "run_generate(example_prompt)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Q_duIXtnAcj9"
+      },
+      "source": [
+        "### General Question test\n",
+        "We'll also attempt a more general prompt which is a better indication of whether our steering vector is having an effect or not"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "UmqQEAM3Ab0i"
+      },
+      "outputs": [],
+      "source": [
+        "question_prompt = \"What is on your mind?\"\n",
+        "coeff = 100\n",
+        "sampling_kwargs = dict(temperature=1.0, top_p=0.1, freq_penalty=1.0)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 337,
+          "referenced_widgets": [
+            "106650a69f4c4bd0a340d58c4bd4f1bb",
+            "06d77984a1a64d39938bfe68e114539b",
+            "6571f57262c447ce9177223fb583e707",
+            "a2179cafb63f475db0162cd990a17ff7",
+            "c0bb81765e93420796cd5f959e9d3534",
+            "fe6cae73e861414eaff54680113676bc",
+            "3f5f9cad86e24dd489146215c3a208c9",
+            "70006fb01d6a49fb909e4a3bfc5b940a",
+            "7980b120d41247548f49667cea6156a5",
+            "359ef2b8a4ac4a9c9a91edc4a2dd1326",
+            "c66dc6c14a4c4274900abe8fc993266a"
+          ]
+        },
+        "id": "HUanDPQeAss3",
+        "outputId": "ecb100a3-d855-4c3e-a758-bd7a3cfebd23"
+      },
+      "outputs": [],
+      "source": [
+        "steering_on = True\n",
+        "run_generate(question_prompt)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 337,
+          "referenced_widgets": [
+            "a8bdc4ecce4f48e0ba6483ea9e679336",
+            "60604227dac34e37a0a9f3bfb3984317",
+            "4024c181581c485abd3181586afc2574",
+            "7761a50a602f41f1a21aa826c491eb9d",
+            "25ebd285de2e49c483c3b22b5c8364c0",
+            "3b74befc8d70471697ce6686ab4ac5c3",
+            "b2ff537e768b43ef98c412e633ab9e49",
+            "3fdf0c5e62f24f30b02bcdc37b17c2e7",
+            "07c0dd1a8de149408b981a8892f6e46d",
+            "b272384164504fa5b81d5502c12f8800",
+            "f525b9f19c334fe6b2305ad6bcfa20bf"
+          ]
+        },
+        "id": "W07bAiWqBlXh",
+        "outputId": "a6b074e6-8183-41ec-c390-2d6430eefdc7"
+      },
+      "outputs": [],
+      "source": [
+        "steering_on = False\n",
+        "run_generate(question_prompt)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JVTbMgMzCLB9"
+      },
+      "source": [
+        "## Next Steps\n",
+        "\n",
+        "Ideas you could take for further exploration:\n",
+        "\n",
+        "*   Try ablating the feature\n",
+        "*   Try and get a response where just the feature token prints over and over\n",
+        "*   Investigate other features with more complex usage\n",
+        "\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "fapxk8MDrs6R",
+        "CmaPYLpGrxbo"
+      ],
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "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.11.7"
+    },
+    "vscode": {
+      "interpreter": {
+        "hash": "088c6e4e32c1710b3b346fe2c9e3084abd3190c888871e6e5b66f23c765b3959"
+      }
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "06d77984a1a64d39938bfe68e114539b": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_fe6cae73e861414eaff54680113676bc",
+            "placeholder": "​",
+            "style": "IPY_MODEL_3f5f9cad86e24dd489146215c3a208c9",
+            "value": "100%"
+          }
+        },
+        "07c0dd1a8de149408b981a8892f6e46d": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "106650a69f4c4bd0a340d58c4bd4f1bb": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_06d77984a1a64d39938bfe68e114539b",
+              "IPY_MODEL_6571f57262c447ce9177223fb583e707",
+              "IPY_MODEL_a2179cafb63f475db0162cd990a17ff7"
+            ],
+            "layout": "IPY_MODEL_c0bb81765e93420796cd5f959e9d3534"
+          }
+        },
+        "1a5dd5f7c9d340b6ab00ecaf43525ae9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_8211cc6c973a43fcaf18e14f6d7f08a2",
+              "IPY_MODEL_3d3584d1feec459287ffa24c4ef790c3",
+              "IPY_MODEL_5f03835168e64ec588c50ee21fedd198"
+            ],
+            "layout": "IPY_MODEL_b833db18729f422cb86deed4be6f1900"
+          }
+        },
+        "229dad8e29f04c279c5603286e2c0643": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "25ebd285de2e49c483c3b22b5c8364c0": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "359ef2b8a4ac4a9c9a91edc4a2dd1326": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "38341454dd6b4e9ca2fe5b85d2e371e1": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "3b74befc8d70471697ce6686ab4ac5c3": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "3d3584d1feec459287ffa24c4ef790c3": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_a30c82833f55441995744300c2ef538d",
+            "max": 50,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_4932983d4f1a4199b3d24c730c765a24",
+            "value": 50
+          }
+        },
+        "3f5f9cad86e24dd489146215c3a208c9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "3fdf0c5e62f24f30b02bcdc37b17c2e7": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "4024c181581c485abd3181586afc2574": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_3fdf0c5e62f24f30b02bcdc37b17c2e7",
+            "max": 50,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_07c0dd1a8de149408b981a8892f6e46d",
+            "value": 50
+          }
+        },
+        "4932983d4f1a4199b3d24c730c765a24": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "5c53f97287d54c03a378fc44ab791cd7": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "5e9700580d6b4ad0bfac34bf3b3919fc": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "5f03835168e64ec588c50ee21fedd198": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_c20e9e14100d45f3bdff1b6df943940f",
+            "placeholder": "​",
+            "style": "IPY_MODEL_5c53f97287d54c03a378fc44ab791cd7",
+            "value": " 50/50 [00:01&lt;00:00, 29.69it/s]"
+          }
+        },
+        "60604227dac34e37a0a9f3bfb3984317": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_3b74befc8d70471697ce6686ab4ac5c3",
+            "placeholder": "​",
+            "style": "IPY_MODEL_b2ff537e768b43ef98c412e633ab9e49",
+            "value": "100%"
+          }
+        },
+        "635162e10abc441797d4e5b74713bf44": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_5e9700580d6b4ad0bfac34bf3b3919fc",
+            "placeholder": "​",
+            "style": "IPY_MODEL_a2c30462ef8d41fd9158f194a746d5a7",
+            "value": " 50/50 [00:02&lt;00:00, 29.94it/s]"
+          }
+        },
+        "6571f57262c447ce9177223fb583e707": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_70006fb01d6a49fb909e4a3bfc5b940a",
+            "max": 50,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_7980b120d41247548f49667cea6156a5",
+            "value": 50
+          }
+        },
+        "66d406d6eb1f49699ee09c9a2fd4ffa9": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "70006fb01d6a49fb909e4a3bfc5b940a": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "720b4d010c364e3fbf72a53b267e8db9": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "7761a50a602f41f1a21aa826c491eb9d": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_b272384164504fa5b81d5502c12f8800",
+            "placeholder": "​",
+            "style": "IPY_MODEL_f525b9f19c334fe6b2305ad6bcfa20bf",
+            "value": " 50/50 [00:01&lt;00:00, 28.55it/s]"
+          }
+        },
+        "7980b120d41247548f49667cea6156a5": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "79b59cbde9444bf892931d31afec7f2a": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_d9c33fbfb3164cbbb7b9a4cd172d20ae",
+            "placeholder": "​",
+            "style": "IPY_MODEL_df53331cce124bd1ada5aa9e9a977015",
+            "value": "100%"
+          }
+        },
+        "8211cc6c973a43fcaf18e14f6d7f08a2": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_66d406d6eb1f49699ee09c9a2fd4ffa9",
+            "placeholder": "​",
+            "style": "IPY_MODEL_38341454dd6b4e9ca2fe5b85d2e371e1",
+            "value": "100%"
+          }
+        },
+        "83d947fc3338491ab4155b87c443884c": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "9f555c5ada38495eb4281cbb49169abe": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_79b59cbde9444bf892931d31afec7f2a",
+              "IPY_MODEL_a157870318114d459a33d795850967ef",
+              "IPY_MODEL_635162e10abc441797d4e5b74713bf44"
+            ],
+            "layout": "IPY_MODEL_720b4d010c364e3fbf72a53b267e8db9"
+          }
+        },
+        "a157870318114d459a33d795850967ef": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_229dad8e29f04c279c5603286e2c0643",
+            "max": 50,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_83d947fc3338491ab4155b87c443884c",
+            "value": 50
+          }
+        },
+        "a2179cafb63f475db0162cd990a17ff7": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_359ef2b8a4ac4a9c9a91edc4a2dd1326",
+            "placeholder": "​",
+            "style": "IPY_MODEL_c66dc6c14a4c4274900abe8fc993266a",
+            "value": " 50/50 [00:01&lt;00:00, 28.62it/s]"
+          }
+        },
+        "a2c30462ef8d41fd9158f194a746d5a7": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "a30c82833f55441995744300c2ef538d": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "a8bdc4ecce4f48e0ba6483ea9e679336": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_60604227dac34e37a0a9f3bfb3984317",
+              "IPY_MODEL_4024c181581c485abd3181586afc2574",
+              "IPY_MODEL_7761a50a602f41f1a21aa826c491eb9d"
+            ],
+            "layout": "IPY_MODEL_25ebd285de2e49c483c3b22b5c8364c0"
+          }
+        },
+        "b272384164504fa5b81d5502c12f8800": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "b2ff537e768b43ef98c412e633ab9e49": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "b833db18729f422cb86deed4be6f1900": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "c0bb81765e93420796cd5f959e9d3534": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "c20e9e14100d45f3bdff1b6df943940f": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "c66dc6c14a4c4274900abe8fc993266a": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "d9c33fbfb3164cbbb7b9a4cd172d20ae": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "df53331cce124bd1ada5aa9e9a977015": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "f525b9f19c334fe6b2305ad6bcfa20bf": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "fe6cae73e861414eaff54680113676bc": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        }
+      }
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}