', 'When', ' John', ' and', ' Mary', ' went', ' to', ' the', ' shops', ',', ' John', ' gave', ' the', ' bag', ' to']\n",
+ "Tokenized answer: [' Mary']\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "Performance on answer token:\n",
+ "Rank: 0 Logit: 17.90 Prob: 75.62% Token: | Mary|\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "Performance on answer token:\n",
+ "\u001b[1mRank: \u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m Logit: \u001b[0m\u001b[1;36m17.90\u001b[0m\u001b[1m Prob: \u001b[0m\u001b[1;36m75.62\u001b[0m\u001b[1m% Token: | Mary|\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Top 0th token. Logit: 17.90 Prob: 75.62% Token: | Mary|\n",
+ "Top 1th token. Logit: 16.47 Prob: 18.18% Token: | the|\n",
+ "Top 2th token. Logit: 14.11 Prob: 1.71% Token: | his|\n",
+ "Top 3th token. Logit: 13.87 Prob: 1.34% Token: | a|\n",
+ "Top 4th token. Logit: 12.59 Prob: 0.38% Token: | their|\n",
+ "Top 5th token. Logit: 11.97 Prob: 0.20% Token: | John|\n",
+ "Top 6th token. Logit: 11.68 Prob: 0.15% Token: | Sarah|\n",
+ "Top 7th token. Logit: 11.45 Prob: 0.12% Token: | Mrs|\n",
+ "Top 8th token. Logit: 11.17 Prob: 0.09% Token: | someone|\n",
+ "Top 9th token. Logit: 11.10 Prob: 0.08% Token: | her|\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "Ranks of the answer tokens: [(' Mary', 0)]\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' Mary'\u001b[0m, \u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Orig 3.7073116302490234\n",
+ "reconstr 3.7859463691711426\n",
+ "Zero 12.452932357788086\n",
+ "Tokenized prompt: ['', 'When', ' John', ' and', ' Mary', ' went', ' to', ' the', ' shops', ',', ' John', ' gave', ' the', ' bag', ' to']\n",
+ "Tokenized answer: [' Mary']\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "Performance on answer token:\n",
+ "Rank: 0 Logit: 17.26 Prob: 67.92% Token: | Mary|\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "Performance on answer token:\n",
+ "\u001b[1mRank: \u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m Logit: \u001b[0m\u001b[1;36m17.26\u001b[0m\u001b[1m Prob: \u001b[0m\u001b[1;36m67.92\u001b[0m\u001b[1m% Token: | Mary|\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Top 0th token. Logit: 17.26 Prob: 67.92% Token: | Mary|\n",
+ "Top 1th token. Logit: 16.21 Prob: 23.79% Token: | the|\n",
+ "Top 2th token. Logit: 13.92 Prob: 2.42% Token: | a|\n",
+ "Top 3th token. Logit: 12.88 Prob: 0.86% Token: | his|\n",
+ "Top 4th token. Logit: 11.99 Prob: 0.35% Token: | her|\n",
+ "Top 5th token. Logit: 11.77 Prob: 0.28% Token: | John|\n",
+ "Top 6th token. Logit: 11.71 Prob: 0.27% Token: | shop|\n",
+ "Top 7th token. Logit: 11.62 Prob: 0.24% Token: | their|\n",
+ "Top 8th token. Logit: 11.44 Prob: 0.20% Token: | him|\n",
+ "Top 9th token. Logit: 11.31 Prob: 0.18% Token: | Mrs|\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "Ranks of the answer tokens: [(' Mary', 0)]\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' Mary'\u001b[0m, \u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "example_prompt = \"When John and Mary went to the shops, John gave the bag to\"\n",
+ "example_answer = \" Mary\"\n",
+ "utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)\n",
+ "\n",
+ "logits, cache = model.run_with_cache(example_prompt, prepend_bos=True)\n",
+ "tokens = model.to_tokens(example_prompt)\n",
+ "sae_out = sae(cache[sae.cfg.hook_name])\n",
+ "\n",
+ "\n",
+ "def reconstr_hook(activations, hook, sae_out):\n",
+ " return sae_out\n",
+ "\n",
+ "\n",
+ "def zero_abl_hook(mlp_out, hook):\n",
+ " return torch.zeros_like(mlp_out)\n",
+ "\n",
+ "\n",
+ "hook_name = sae.cfg.hook_name\n",
+ "\n",
+ "print(\"Orig\", model(tokens, return_type=\"loss\").item())\n",
+ "print(\n",
+ " \"reconstr\",\n",
+ " model.run_with_hooks(\n",
+ " tokens,\n",
+ " fwd_hooks=[\n",
+ " (\n",
+ " hook_name,\n",
+ " partial(reconstr_hook, sae_out=sae_out),\n",
+ " )\n",
+ " ],\n",
+ " return_type=\"loss\",\n",
+ " ).item(),\n",
+ ")\n",
+ "print(\n",
+ " \"Zero\",\n",
+ " model.run_with_hooks(\n",
+ " tokens,\n",
+ " return_type=\"loss\",\n",
+ " fwd_hooks=[(hook_name, zero_abl_hook)],\n",
+ " ).item(),\n",
+ ")\n",
+ "\n",
+ "\n",
+ "with model.hooks(\n",
+ " fwd_hooks=[\n",
+ " (\n",
+ " hook_name,\n",
+ " partial(reconstr_hook, sae_out=sae_out),\n",
+ " )\n",
+ " ]\n",
+ "):\n",
+ " utils.test_prompt(example_prompt, example_answer, model, prepend_bos=True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "1swj9KA7v612"
+ },
+ "source": [
+ "# Generating Feature Interfaces\n",
+ "\n",
+ "Feature dashboards are an important part of SAE Evaluation. They work by:\n",
+ "- 1. Collecting feature activations over a larger number of examples.\n",
+ "- 2. Aggregating feature specific statistics (such as max activating examples).\n",
+ "- 3. Representing that information in a standardized way\n",
+ "\n",
+ "For our feature visualizations, we will use a separate library called SAEDashboard."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 26,
+ "metadata": {
+ "id": "edt8ag4fv612"
+ },
+ "outputs": [
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "e6c9cf17032c45de85bd0b6b18adaed5",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Forward passes to cache data for vis: 0%| | 0/2 [00:00┏━━━━━━┳━━━━━━┳━━━━━━━┓\n",
+ "┃ Task ┃ Time ┃ Pct % ┃\n",
+ "┡━━━━━━╇━━━━━━╇━━━━━━━┩\n",
+ "└──────┴──────┴───────┘\n",
+ "\n"
+ ],
+ "text/plain": [
+ "┏━━━━━━┳━━━━━━┳━━━━━━━┓\n",
+ "┃\u001b[1m \u001b[0m\u001b[1mTask\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mTime\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mPct %\u001b[0m\u001b[1m \u001b[0m┃\n",
+ "┡━━━━━━╇━━━━━━╇━━━━━━━┩\n",
+ "└──────┴──────┴───────┘\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "# %pip install sae_dashboard\n",
+ "from sae_dashboard.sae_vis_data import SaeVisConfig\n",
+ "from sae_dashboard.sae_vis_runner import SaeVisRunner\n",
+ "\n",
+ "test_feature_idx_gpt = list(range(100)) + [14057] # TODO: test all the features\n",
+ "\n",
+ "feature_vis_config_gpt = SaeVisConfig(\n",
+ " hook_point=hook_name,\n",
+ " features=test_feature_idx_gpt,\n",
+ " minibatch_size_features=64,\n",
+ " minibatch_size_tokens=256,\n",
+ " verbose=True,\n",
+ " device=device,\n",
+ ")\n",
+ "\n",
+ "visualization_data_gpt = SaeVisRunner(\n",
+ " feature_vis_config_gpt\n",
+ ").run(\n",
+ " encoder=sae, # type: ignore\n",
+ " model=model,\n",
+ " tokens=token_dataset_tensors[:1], # type: ignore \n",
+ ")\n",
+ "# SaeVisData.create(\n",
+ "# encoder=sae,\n",
+ "# model=model, # type: ignore\n",
+ "# tokens=token_dataset[:10000][\"tokens\"], # type: ignore\n",
+ "# cfg=feature_vis_config_gpt,\n",
+ "# )"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 27,
+ "metadata": {
+ "id": "yQ94Frzbv612"
+ },
+ "outputs": [
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "d55a7d33346e47559a5f60f961a543c6",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Saving feature-centric vis: 0%| | 0/101 [00:00