sae_dashboard_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo Notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Steps:\n",
    "1. Download SAE with SAE Lens.\n",
    "2. Create a dataset consistent with that SAE. \n",
    "3. Fold the SAE decoder norm weights so that feature activations are \"correct\".\n",
    "4. Estimate the activation normalization constant if needed, and fold it into the SAE weights.\n",
    "5. Run the SAE generator for the features you want."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from sae_lens import SAE\n",
    "from transformer_lens import HookedTransformer\n",
    "from sae_dashboard.sae_vis_data import SaeVisConfig\n",
    "from sae_dashboard.sae_vis_runner import SaeVisRunner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1. Download / Initialize SAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For the most part I'll try to import functions and classes near where they are used\n",
    "# to make it clear where they come from.\n",
    "\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}\")\n",
    "\n",
    "model = HookedTransformer.from_pretrained(\"gemma-2b\", device=device, dtype=\"bfloat16\")\n",
    "\n",
    "# the cfg dict is returned alongside the SAE since it may contain useful information for analysing the SAE (eg: instantiating an activation store)\n",
    "# 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\n",
    "# We also return the feature sparsities which are stored in HF for convenience.\n",
    "sae, cfg_dict, sparsity = SAE.from_pretrained(\n",
    "    release=\"gemma-2b-res-jb\",  # see other options in sae_lens/pretrained_saes.yaml\n",
    "    sae_id=\"blocks.6.hook_resid_post\",  # won't always be a hook point\n",
    "    device=device,\n",
    ")\n",
    "# fold w_dec norm so feature activations are accurate\n",
    "sae.fold_W_dec_norm()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Get token dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sae_lens import ActivationsStore\n",
    "\n",
    "activations_store = ActivationsStore.from_sae(\n",
    "    model=model,\n",
    "    sae=sae,\n",
    "    streaming=True,\n",
    "    store_batch_size_prompts=16,\n",
    "    n_batches_in_buffer=8,\n",
    "    device=device,\n",
    ")\n",
    "\n",
    "# Some SAEs will require we estimate the activation norm and fold it into the weights. This is easy with SAE Lens.\n",
    "if sae.cfg.normalize_activations == \"expected_average_only_in\":\n",
    "    norm_scaling_factor = activations_store.estimate_norm_scaling_factor(\n",
    "        n_batches_for_norm_estimate=30\n",
    "    )\n",
    "    sae.fold_activation_norm_scaling_factor(norm_scaling_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def get_tokens(\n",
    "    activations_store: ActivationsStore,\n",
    "    n_prompts: int,\n",
    "):\n",
    "    all_tokens_list = []\n",
    "    pbar = tqdm(range(n_prompts))\n",
    "    for _ in pbar:\n",
    "        batch_tokens = activations_store.get_batch_tokens()\n",
    "        batch_tokens = batch_tokens[torch.randperm(batch_tokens.shape[0])][\n",
    "            : batch_tokens.shape[0]\n",
    "        ]\n",
    "        all_tokens_list.append(batch_tokens)\n",
    "\n",
    "    all_tokens = torch.cat(all_tokens_list, dim=0)\n",
    "    all_tokens = all_tokens[torch.randperm(all_tokens.shape[0])]\n",
    "    return all_tokens\n",
    "\n",
    "\n",
    "# 1000 prompts is plenty for a demo.\n",
    "token_dataset = get_tokens(activations_store, 128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3 Evaluate the SAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sae_lens import run_evals\n",
    "\n",
    "eval_metrics = run_evals(\n",
    "    sae=sae,\n",
    "    activation_store=activations_store,\n",
    "    model=model,\n",
    "    n_eval_batches=3,\n",
    "    eval_batch_size_prompts=8,\n",
    ")\n",
    "\n",
    "# CE Loss score should be high for residual stream SAEs\n",
    "print(eval_metrics[\"metrics/CE_loss_score\"])\n",
    "\n",
    "# ce loss without SAE should be fairly low < 3.5 suggesting the Model is being run correctly\n",
    "print(eval_metrics[\"metrics/ce_loss_without_sae\"])\n",
    "\n",
    "# ce loss with SAE shouldn't be massively higher\n",
    "print(eval_metrics[\"metrics/ce_loss_with_sae\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Generate Feature Dashboards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "test_feature_idx_gpt = list(range(256))\n",
    "\n",
    "feature_vis_config_gpt = SaeVisConfig(\n",
    "    hook_point=sae.cfg.hook_name,\n",
    "    features=test_feature_idx_gpt,\n",
    "    minibatch_size_features=64,\n",
    "    minibatch_size_tokens=256,  # this is number of prompts at a time.\n",
    "    verbose=True,\n",
    "    device=\"cuda\",\n",
    "    cache_dir=Path(\n",
    "        \"demo_activations_cache\"\n",
    "    ),  # this will enable us to skip running the model for subsequent features.\n",
    "    dtype=\"bfloat16\",\n",
    ")\n",
    "\n",
    "data = SaeVisRunner(feature_vis_config_gpt).run(\n",
    "    encoder=sae,  # type: ignore\n",
    "    model=model,\n",
    "    tokens=token_dataset,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sae_dashboard.data_writing_fns import save_feature_centric_vis\n",
    "\n",
    "filename = f\"demo_feature_dashboards.html\"\n",
    "save_feature_centric_vis(sae_vis_data=data, filename=filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}