tutorials/training_a_sparse_autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5O8tQblzOVHu"
   },
   "source": [
    "# Training a basic SAE with SAELens\n",
    "\n",
    "This tutorial demonstrates training a simple, relatively small Sparse Autoencoder, specifically on the tiny-stories-1L-21M model. \n",
    "\n",
    "As the SAELens library is under active development, please open an issue if this tutorial is stale [here](https://github.com/jbloomAus/SAELens)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "shAFb9-lOVHu"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "LeRi_tw2dhae"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: sae-lens in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (3.3.0)\n",
      "Requirement already satisfied: transformer-lens in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (1.19.0)\n",
      "Requirement already satisfied: circuitsvis in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (1.43.2)\n",
      "Requirement already satisfied: automated-interpretability<0.0.4,>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.0.3)\n",
      "Requirement already satisfied: babe<0.0.8,>=0.0.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.0.7)\n",
      "Requirement already satisfied: datasets<3.0.0,>=2.17.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (2.19.2)\n",
      "Requirement already satisfied: matplotlib<4.0.0,>=3.8.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (3.9.0)\n",
      "Requirement already satisfied: matplotlib-inline<0.2.0,>=0.1.6 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.1.7)\n",
      "Requirement already satisfied: nltk<4.0.0,>=3.8.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (3.8.1)\n",
      "Requirement already satisfied: plotly<6.0.0,>=5.19.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (5.22.0)\n",
      "Requirement already satisfied: plotly-express<0.5.0,>=0.4.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.4.1)\n",
      "Requirement already satisfied: pytest-profiling<2.0.0,>=1.7.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (1.7.0)\n",
      "Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (1.0.1)\n",
      "Requirement already satisfied: pyyaml<7.0.0,>=6.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (6.0.1)\n",
      "Requirement already satisfied: pyzmq==26.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (26.0.0)\n",
      "Requirement already satisfied: sae-vis<0.3.0,>=0.2.18 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.2.18)\n",
      "Requirement already satisfied: safetensors<0.5.0,>=0.4.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.4.3)\n",
      "Requirement already satisfied: transformers<5.0.0,>=4.38.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (4.41.2)\n",
      "Requirement already satisfied: typer<0.13.0,>=0.12.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.12.3)\n",
      "Requirement already satisfied: accelerate>=0.23.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.31.0)\n",
      "Requirement already satisfied: beartype<0.15.0,>=0.14.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.14.1)\n",
      "Requirement already satisfied: better-abc<0.0.4,>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.0.3)\n",
      "Requirement already satisfied: einops>=0.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.7.0)\n",
      "Requirement already satisfied: fancy-einsum>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.0.3)\n",
      "Requirement already satisfied: jaxtyping>=0.2.11 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.2.29)\n",
      "Requirement already satisfied: numpy>=1.24 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (1.26.4)\n",
      "Requirement already satisfied: pandas>=1.1.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (2.2.2)\n",
      "Requirement already satisfied: rich>=12.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (13.7.1)\n",
      "Requirement already satisfied: sentencepiece in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.2.0)\n",
      "Requirement already satisfied: torch>=1.10 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (2.1.2)\n",
      "Requirement already satisfied: tqdm>=4.64.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (4.66.4)\n",
      "Requirement already satisfied: typing-extensions in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (4.12.2)\n",
      "Requirement already satisfied: wandb>=0.13.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.17.1)\n",
      "Requirement already satisfied: importlib-metadata>=5.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (7.1.0)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.18.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (2.18.1)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
      "Requirement already satisfied: triton==2.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (2.1.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->circuitsvis) (12.5.40)\n",
      "Requirement already satisfied: filelock in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from triton==2.1.0->circuitsvis) (3.14.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (24.0)\n",
      "Requirement already satisfied: psutil in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (5.9.0)\n",
      "Requirement already satisfied: huggingface-hub in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (0.23.3)\n",
      "Requirement already satisfied: blobfile<3.0.0,>=2.1.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.1.1)\n",
      "Requirement already satisfied: boostedblob<0.16.0,>=0.15.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.15.3)\n",
      "Requirement already satisfied: httpx<0.28.0,>=0.27.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.27.0)\n",
      "Requirement already satisfied: orjson<4.0.0,>=3.10.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.10.4)\n",
      "Requirement already satisfied: pytest<9.0.0,>=8.1.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (8.2.2)\n",
      "Requirement already satisfied: scikit-learn<2.0.0,>=1.4.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.5.0)\n",
      "Requirement already satisfied: tiktoken<0.7.0,>=0.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.6.0)\n",
      "Requirement already satisfied: py2store in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from babe<0.0.8,>=0.0.7->sae-lens) (0.1.20)\n",
      "Requirement already satisfied: graze in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from babe<0.0.8,>=0.0.7->sae-lens) (0.1.17)\n",
      "Requirement already satisfied: pyarrow>=12.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (16.1.0)\n",
      "Requirement already satisfied: pyarrow-hotfix in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.6)\n",
      "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.3.8)\n",
      "Requirement already satisfied: requests>=2.32.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (2.32.3)\n",
      "Requirement already satisfied: xxhash in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (3.4.1)\n",
      "Requirement already satisfied: multiprocess in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.70.16)\n",
      "Requirement already satisfied: fsspec<=2024.3.1,>=2023.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from fsspec[http]<=2024.3.1,>=2023.1.0->datasets<3.0.0,>=2.17.1->sae-lens) (2024.3.1)\n",
      "Requirement already satisfied: aiohttp in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (3.9.5)\n",
      "Requirement already satisfied: zipp>=0.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from importlib-metadata>=5.1.0->circuitsvis) (3.19.2)\n",
      "Requirement already satisfied: typeguard==2.13.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from jaxtyping>=0.2.11->transformer-lens) (2.13.3)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (1.2.1)\n",
      "Requirement already satisfied: cycler>=0.10 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (4.53.0)\n",
      "Requirement already satisfied: kiwisolver>=1.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (1.4.5)\n",
      "Requirement already satisfied: pillow>=8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (10.3.0)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (3.1.2)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (2.9.0)\n",
      "Requirement already satisfied: traitlets in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib-inline<0.2.0,>=0.1.6->sae-lens) (5.14.3)\n",
      "Requirement already satisfied: click in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (8.1.7)\n",
      "Requirement already satisfied: joblib in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (1.4.2)\n",
      "Requirement already satisfied: regex>=2021.8.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (2024.5.15)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pandas>=1.1.5->transformer-lens) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pandas>=1.1.5->transformer-lens) (2024.1)\n",
      "Requirement already satisfied: tenacity>=6.2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly<6.0.0,>=5.19.0->sae-lens) (8.3.0)\n",
      "Requirement already satisfied: statsmodels>=0.9.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (0.14.2)\n",
      "Requirement already satisfied: scipy>=0.18 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (1.13.1)\n",
      "Requirement already satisfied: patsy>=0.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (0.5.6)\n",
      "Requirement already satisfied: six in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens) (1.16.0)\n",
      "Requirement already satisfied: gprof2dot in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens) (2024.6.6)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from rich>=12.6.0->transformer-lens) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from rich>=12.6.0->transformer-lens) (2.18.0)\n",
      "Requirement already satisfied: dataclasses-json<0.7.0,>=0.6.4 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens) (0.6.7)\n",
      "Requirement already satisfied: eindex-callum<0.2.0,>=0.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens) (0.1.1)\n",
      "Requirement already satisfied: sympy in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (1.12.1)\n",
      "Requirement already satisfied: networkx in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (3.3)\n",
      "Requirement already satisfied: jinja2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (3.1.4)\n",
      "Requirement already satisfied: tokenizers<0.20,>=0.19 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformers<5.0.0,>=4.38.1->sae-lens) (0.19.1)\n",
      "Requirement already satisfied: shellingham>=1.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from typer<0.13.0,>=0.12.3->sae-lens) (1.5.4)\n",
      "Requirement already satisfied: docker-pycreds>=0.4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (0.4.0)\n",
      "Requirement already satisfied: gitpython!=3.1.29,>=1.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (3.1.43)\n",
      "Requirement already satisfied: platformdirs in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (4.2.2)\n",
      "Requirement already satisfied: protobuf!=4.21.0,<6,>=3.19.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (5.27.1)\n",
      "Requirement already satisfied: sentry-sdk>=1.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (2.5.1)\n",
      "Requirement already satisfied: setproctitle in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (1.3.3)\n",
      "Requirement already satisfied: setuptools in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (70.0.0)\n",
      "Requirement already satisfied: pycryptodomex~=3.8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.20.0)\n",
      "Requirement already satisfied: urllib3<3,>=1.25.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.2.1)\n",
      "Requirement already satisfied: lxml~=4.9 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (4.9.4)\n",
      "Requirement already satisfied: uvloop>=0.16.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from boostedblob<0.16.0,>=0.15.3->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.19.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (4.0.3)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (3.21.3)\n",
      "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (0.9.0)\n",
      "Requirement already satisfied: gitdb<5,>=4.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer-lens) (4.0.11)\n",
      "Requirement already satisfied: anyio in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (4.4.0)\n",
      "Requirement already satisfied: certifi in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2024.6.2)\n",
      "Requirement already satisfied: httpcore==1.* in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.0.5)\n",
      "Requirement already satisfied: idna in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.7)\n",
      "Requirement already satisfied: sniffio in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.3.1)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpcore==1.*->httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.14.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich>=12.6.0->transformer-lens) (0.1.2)\n",
      "Requirement already satisfied: iniconfig in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.0.0)\n",
      "Requirement already satisfied: pluggy<2.0,>=1.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.5.0)\n",
      "Requirement already satisfied: exceptiongroup>=1.0.0rc8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.2.0)\n",
      "Requirement already satisfied: tomli>=1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.0.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from requests>=2.32.1->datasets<3.0.0,>=2.17.1->sae-lens) (3.3.2)\n",
      "Requirement already satisfied: threadpoolctl>=3.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from scikit-learn<2.0.0,>=1.4.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.5.0)\n",
      "Requirement already satisfied: dol in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from graze->babe<0.0.8,>=0.0.7->sae-lens) (0.2.47)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from jinja2->torch>=1.10->transformer-lens) (2.1.5)\n",
      "Requirement already satisfied: config2py in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens) (0.1.33)\n",
      "Requirement already satisfied: importlib-resources in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens) (6.4.0)\n",
      "Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sympy->torch>=1.10->transformer-lens) (1.3.0)\n",
      "Requirement already satisfied: smmap<6,>=3.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer-lens) (5.0.1)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (1.0.0)\n",
      "Requirement already satisfied: i2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from config2py->py2store->babe<0.0.8,>=0.0.7->sae-lens) (0.1.17)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    # import google.colab # type: ignore\n",
    "    # from google.colab import output\n",
    "    %pip install sae-lens transformer-lens circuitsvis\n",
    "except:\n",
    "    from IPython import get_ipython  # type: ignore\n",
    "\n",
    "    ipython = get_ipython()\n",
    "    assert ipython is not None\n",
    "    ipython.run_line_magic(\"load_ext\", \"autoreload\")\n",
    "    ipython.run_line_magic(\"autoreload\", \"2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uy-b3CcSOVHu",
    "outputId": "58ce28d0-f91f-436d-cf87-76bb26e2ecaf"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import os\n",
    "\n",
    "from sae_lens import LanguageModelSAERunnerConfig, SAETrainingRunner\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = \"mps\"\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "\n",
    "print(\"Using device:\", device)\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oe2nlqf-OVHv"
   },
   "source": [
    "# Model Selection and Evaluation (Feel Free to Skip)\n",
    "\n",
    "We'll use the runner to train an SAE on a TinyStories Model. This is a very small model so we can train an SAE on it quite quickly. Before we get started, let's load in the model with `transformer_lens` and see what it can do.\n",
    "\n",
    "TransformerLens gives us 2 functions that are useful here (and circuits viz provides a third):\n",
    "1. `transformer_lens.utils.test_prompt` will help us see when the model can infer one token.\n",
    "2. `HookedTransformer.generate` will help us see what happens when we sample from the model.\n",
    "3. `circuitsvis.logits.token_log_probs` will help us visualize the log probs of tokens at several positions in a prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "hFz6JUMuOVHv"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n",
      "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained model tiny-stories-1L-21M into HookedTransformer\n"
     ]
    }
   ],
   "source": [
    "from transformer_lens import HookedTransformer\n",
    "\n",
    "model = HookedTransformer.from_pretrained(\n",
    "    \"tiny-stories-1L-21M\"\n",
    ")  # This will wrap huggingface models and has lots of nice utilities."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aUiXrjdUOVHv"
   },
   "source": [
    "### Getting a vibe for a model using `model.generate`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZZfKT5aDOVHv"
   },
   "source": [
    "Let's start by generating some stories using the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "G4ad4Zz1OVHv"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Once upon a time, Bobby was hungry and needed something to do. He went to the subway but was far away.\\n\\nThe man wanted to get the hat, so the people wanted it. He found it hard to be a big, powerful bird. It wanted'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Once upon a time, there was a trunk. The trunk was very rich, and it was a very special trunk. All the animals came across the trunk, and was very colorful. They took turns to fill it up when they bumped into a dragon. \\n\\n'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Once upon a time, there was a young man. He was three years old. He wanted to learn how to keep the match safe. So he kept checking it every day.\\n\\nOne day a 3 year old girl wanted to learn about the fun trunk of the'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Once upon a time, there was a little girl named Sally. She liked to play with her toys. One sunny day, Sally found a butterfly. She was so happy! She wanted to play with something new. So, she called her friend Tom.\\n\\nTom'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'Once upon a time, there was a little girl named Lola. She really loved playing with her pet cat, Tom and show them his appreciation.\\n\\nOne day, Lola licked the couch closer and soon found herself in a magical land! As soon as'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# here we use generate to get 10 completeions with temperature 1. Feel free to play with the prompt to make it more interesting.\n",
    "for i in range(5):\n",
    "    display(\n",
    "        model.generate(\n",
    "            \"Once upon a time\",\n",
    "            stop_at_eos=False,  # avoids a bug on MPS\n",
    "            temperature=1,\n",
    "            verbose=False,\n",
    "            max_new_tokens=50,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RDKr8o1xOVHv"
   },
   "source": [
    "One thing we notice is that the model seems to be able to repeat the name of the main character very consistently. It can output a pronoun intead but in some stories will repeat the protagonists name. This seems like an interesting capability to analyse with SAEs. To better understand the models ability to remember the protagonists name, let's extract a prompt where the next character is determined and use the \"test_prompt\" utility from TransformerLens to check the ranking of the token for that name."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KsfJX-YpOVHv"
   },
   "source": [
    "### Spot checking model abilities with `transformer_lens.utils.test_prompt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "TpmPoj7uOVHv"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokenized prompt: ['<|endoftext|>', 'Once', ' upon', ' a', ' time', ',', ' there', ' was', ' a', ' little', ' girl', ' named', ' Lily', '.', ' She', ' lived', ' in', ' a', ' big', ',', ' happy', ' little', ' girl', '.', ' On', ' her', ' big', ' adventure', ',']\n",
      "Tokenized answer: [' Lily']\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Performance on answer token:\n",
       "<span style=\"font-weight: bold\">Rank: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">        Logit: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">18.81</span><span style=\"font-weight: bold\"> Prob: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">13.46</span><span style=\"font-weight: bold\">% Token: | Lily|</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Performance on answer token:\n",
       "\u001b[1mRank: \u001b[0m\u001b[1;36m1\u001b[0m\u001b[1m        Logit: \u001b[0m\u001b[1;36m18.81\u001b[0m\u001b[1m Prob: \u001b[0m\u001b[1;36m13.46\u001b[0m\u001b[1m% Token: | Lily|\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top 0th token. Logit: 20.48 Prob: 71.06% Token: | she|\n",
      "Top 1th token. Logit: 18.81 Prob: 13.46% Token: | Lily|\n",
      "Top 2th token. Logit: 17.35 Prob:  3.11% Token: | the|\n",
      "Top 3th token. Logit: 17.26 Prob:  2.86% Token: | her|\n",
      "Top 4th token. Logit: 16.74 Prob:  1.70% Token: | there|\n",
      "Top 5th token. Logit: 16.43 Prob:  1.25% Token: | they|\n",
      "Top 6th token. Logit: 15.80 Prob:  0.66% Token: | all|\n",
      "Top 7th token. Logit: 15.64 Prob:  0.56% Token: | things|\n",
      "Top 8th token. Logit: 15.28 Prob:  0.39% Token: | one|\n",
      "Top 9th token. Logit: 15.24 Prob:  0.38% Token: | lived|\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Ranks of the answer tokens:</span> <span style=\"font-weight: bold\">[(</span><span style=\"color: #008000; text-decoration-color: #008000\">' Lily'</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">)]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' Lily'\u001b[0m, \u001b[1;36m1\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformer_lens.utils import test_prompt\n",
    "\n",
    "# Test the model with a prompt\n",
    "test_prompt(\n",
    "    \"Once upon a time, there was a little girl named Lily. She lived in a big, happy little town. On her big adventure,\",\n",
    "    \" Lily\",\n",
    "    model,\n",
    "    prepend_space_to_answer=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jGzOvReDOVHv"
   },
   "source": [
    "In the output above, we see that the model assigns ~ 70% probability to \"she\" being the next token, and a 13% chance to \" Lily\" being the next token. Other names like Lucy or Anna are not highly ranked."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QH8YOZOzOVHv"
   },
   "source": [
    "### Exploring Model Capabilities with Log Probs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "50mqTBihOVHw"
   },
   "source": [
    "Looking at token ranking for a single prompt is interesting, but a much higher through way to understand models is to look at token log probs for all tokens in text. We can use the `circuits_vis` package to get a nice visualization where we can see tokenization, and hover to get the top5 tokens by log probability. Darker tokens are tokens where the model assigned a higher probability to the actual next token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "Tic0RCUpOVHw"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-118354ce-b30f\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TokenLogProbs } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-118354ce-b30f\",\n",
       "      TokenLogProbs,\n",
       "      {\"prompt\": [\"<|endoftext|>\", \"Hi\", \",\", \" how\", \" are\", \" you\", \" doing\", \" this\", \"?\", \" I\", \"'m\", \" really\", \" enjoying\", \" your\", \" posts\"], \"topKLogProbs\": [[-0.020237978547811508, -6.104498863220215, -7.070311546325684, -7.073490142822266, -7.178879737854004, -7.389684677124023, -7.525315284729004, -7.588573455810547, -7.763055801391602, -8.224309921264648], [-1.1883893013000488, -3.4683852195739746, -3.4808287620544434, -3.4896111488342285, -3.500030994415283, -3.554358959197998, -4.017881870269775, -4.030544757843018, -4.298159122467041, -4.418168544769287], [-2.9201290607452393, -3.379486322402954, -3.4266788959503174, -3.531764268875122, -3.956162691116333, -4.256065368652344, -4.443942070007324, -4.4448394775390625, -4.559991836547852, -4.583307266235352], [-0.9917522668838501, -2.928676128387451, -3.2575316429138184, -3.317981243133545, -3.4740567207336426, -3.6419081687927246, -3.8674941062927246, -3.868630886077881, -3.8917040824890137, -4.005284786224365], [-0.037123147398233414, -4.683361530303955, -4.9722676277160645, -5.854288578033447, -5.9421305656433105, -6.388034343719482, -6.3940815925598145, -6.498584270477295, -6.593432903289795, -7.8692545890808105], [-0.23486259579658508, -2.3354179859161377, -2.947497606277466, -5.052804470062256, -5.078576564788818, -5.102746486663818, -5.116313457489014, -5.737276554107666, -5.768456935882568, -5.893773555755615], [-0.44884783029556274, -2.076716423034668, -3.433588981628418, -3.645073890686035, -3.7033796310424805, -4.046278953552246, -4.097006797790527, -4.209780693054199, -4.518967628479004, -4.961886405944824], [-0.4609467387199402, -2.7377679347991943, -3.158902883529663, -3.686795949935913, -4.416429042816162, -4.448527812957764, -4.571945667266846, -4.669140338897705, -4.695371150970459, -4.748204708099365], [-0.9866101741790771, -1.8349759578704834, -2.460031270980835, -2.511583089828491, -3.1156270503997803, -3.419478178024292, -3.5203635692596436, -3.8049819469451904, -4.003089904785156, -4.258264541625977], [-1.3091926574707031, -1.633901596069336, -1.78680419921875, -2.631612777709961, -3.4602813720703125, -3.615842819213867, -4.020046234130859, -4.093286514282227, -4.177457809448242, -4.309240341186523], [-1.9952350854873657, -2.55653715133667, -2.645555019378662, -2.8948235511779785, -3.11043119430542, -3.3945651054382324, -3.7311148643493652, -3.8203349113464355, -3.967912197113037, -4.026367664337158], [-1.0431095361709595, -1.959466576576233, -3.0507731437683105, -3.4215006828308105, -3.48012113571167, -3.4898486137390137, -3.6020045280456543, -4.332496166229248, -4.492417812347412, -4.517611026763916], [-1.5123934745788574, -2.020272731781006, -2.22328519821167, -2.2493205070495605, -3.376523494720459, -3.92171049118042, -3.9798388481140137, -4.199243068695068, -4.265368938446045, -4.6615891456604], [-3.188736915588379, -3.192061424255371, -3.4252281188964844, -3.494192123413086, -3.7325897216796875, -3.7353057861328125, -3.735713005065918, -3.833963394165039, -4.233694076538086, -4.265708923339844]], \"topKTokens\": [[\"\\n\", \",\", \"Words\", \"Summary\", \"\\n\\n\", \"<|endoftext|>\", \" \", \"Features\", \"Random\", \" the\"], [\",\", \" Tim\", \" bird\", \"!\", \" little\", \" Lily\", \" Tom\", \"!\\\"\", \" Max\", \" tree\"], [\" cat\", \" a\", \" bird\", \" little\", \" I\", \"\\n\", \" dog\", \" but\", \" frog\", \" sw\"], [\" are\", \" I\", \" was\", \" do\", \" fast\", \" big\", \" did\", \" to\", \" he\", \" old\"], [\" you\", \" we\", \" they\", \" the\", \" Lily\", \" your\", \" I\", \" Tim\", \" friends\", \" those\"], [\"?\\\"\", \" today\", \"?\", \" and\", \".\", \",\", \" doing\", \" going\", \" up\", \"?\\\".\"], [\"?\\\"\", \" this\", \" today\", \"?\", \" here\", \" a\", \" it\", \" these\", \".\", \" in\"], [\"?\\\"\", \"?\", \" great\", \" so\", \" job\", \" puzzle\", \" task\", \" very\", \".\", \" amazing\"], [\"\\ufffd\", \" You\", \" I\", \" It\", \" We\", \" Can\", \"\\n\", \" This\", \" Do\", \" Are\"], [\"'m\", \" am\", \" want\", \" like\", \" know\", \" have\", \" love\", \" can\", \" just\", \" hope\"], [\" a\", \" so\", \" very\", \" trying\", \" going\", \" an\", \" looking\", \" just\", \" making\", \" playing\"], [\" good\", \" hungry\", \" enjoying\", \" a\", \" tired\", \" proud\", \" happy\", \" excited\", \" busy\", \" glad\"], [\" the\", \" it\", \" this\", \" my\", \" a\", \" your\", \" playing\", \" some\", \" these\", \" myself\"], [\" meal\", \" food\", \" day\", \" ice\", \" sandwich\", \" delicious\", \" new\", \" dance\", \" cake\", \" time\"]], \"correctTokenRank\": [1718, 0, 675, 0, 0, 6, 1, 1, 2, 0, 23, 2, 5, 10036], \"correctTokenLogProb\": [-13.814948081970215, -1.1883893013000488, -8.445575714111328, -0.9917522668838501, -0.037123147398233414, -5.116313457489014, -2.076716423034668, -2.7377679347991943, -2.460031270980835, -1.3091926574707031, -4.860561847686768, -3.0507731437683105, -3.92171049118042, -17.039505004882812]}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x77eaf81e21d0>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import circuitsvis as cv  # optional dep, install with pip install circuitsvis\n",
    "\n",
    "# Let's make a longer prompt and see the log probabilities of the tokens\n",
    "example_prompt = \"\"\"Hi, how are you doing this? I'm really enjoying your posts\"\"\"\n",
    "logits, cache = model.run_with_cache(example_prompt)\n",
    "cv.logits.token_log_probs(\n",
    "    model.to_tokens(example_prompt),\n",
    "    model(example_prompt)[0].log_softmax(dim=-1),\n",
    "    model.to_string,\n",
    ")\n",
    "# hover on the output to see the result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lhGIl3YbOVHw"
   },
   "source": [
    "Let's combine `model.generate` and the token log probs visualization to see the log probs on text generated by the model. Note that we can play with the temperature and this should sample less likely trajectories according to the model. I've increased the maximum number of tokens in order to get a full story.\n",
    "\n",
    "Some things to explore:\n",
    "- Which tokens does the model assign high probability to? Can you see how the model should know which word comes next?\n",
    "- What happens if you increase / decrease the temperature?\n",
    "- Do the rankings of tokens seem sensible to you? What about where the model doesn't assign a high probability to the token which came next?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "Nikp2ASlOVHw"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:01<00:00, 103.20it/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-a7b13498-82dd\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TokenLogProbs } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-a7b13498-82dd\",\n",
       "      TokenLogProbs,\n",
       "      {\"prompt\": [\"<|endoftext|>\", \"Once\", \" upon\", \" a\", \" time\", \",\", \" there\", \" was\", \" a\", \" mom\", \"my\", \" and\", \" a\", \" big\", \",\", \" modern\", \" airport\", \".\", \" They\", \" saw\", \" many\", \" people\", \" walking\", \" around\", \" and\", \" waved\", \" at\", \" all\", \" the\", \" planes\", \" too\", \".\", \"\\n\", \"\\n\", \"The\", \" family\", \" asked\", \" the\", \" little\", \" girl\", \" and\", \" daddy\", \" lots\", \" of\", \" questions\", \",\", \" so\", \" the\", \" mom\", \"my\", \" replied\", \",\", \" \\\"\", \"Are\", \" we\", \" just\", \" going\", \" to\", \" explore\", \"?\\\"\", \"\\n\", \"\\n\", \"The\", \" little\", \" girl\", \" smiled\", \" and\", \" nodded\", \" her\", \" head\", \".\", \" She\", \" was\", \" excited\", \" to\", \" see\", \" all\", \" the\", \" planes\", \",\", \" and\", \" she\", \" got\", \" to\", \" the\", \" airport\", \".\", \"\\n\", \"\\n\", \"The\", \" little\", \" girl\", \" climbed\", \" up\", \" the\", \" stairs\", \" and\", \" said\", \",\", \" \\\"\", \"Well\", \" done\", \",\", \" everyone\", \" is\", \" so\", \" proud\", \" of\", \" you\", \"!\", \" Are\", \" you\", \" ready\", \" to\", \" show\", \" your\", \" family\", \"?\\\"\", \"\\n\", \"\\n\", \"The\", \" family\", \" looked\", \",\", \" and\", \" they\", \" all\", \" hugged\", \" each\", \" other\", \",\", \" and\", \" they\", \" agreed\", \"!\", \" The\", \" little\", \" girl\", \" was\", \" so\", \" happy\", \",\", \" she\", \" wanted\", \" to\", \" show\", \" her\", \" new\", \" toy\", \".\", \"\\n\", \"<|endoftext|>\", \".\", \"The\", \" family\", \" went\", \" on\", \" to\", \" the\", \" modern\", \" airport\", \" to\", \" come\", \" on\", \" these\", \" trips\", \".\", \" Every\", \" time\", \" one\", \" of\", \" the\", \" little\", \" girl\", \",\", \" her\", \" family\", \",\", \" would\", \" march\", \" with\", \" enthusiasm\", \" and\", \" hope\", \" for\", \" the\", \" family\", \".\", \"\\n\", \"\\n\", \"Eventually\", \",\", \" the\", \" little\", \" girl\", \" began\", \" to\", \" get\", \" thirsty\", \" again\", \".\", \" She\", \" was\", \" so\", \" tired\"], \"topKLogProbs\": [[-0.020237978547811508, -6.104497909545898, -7.070315361022949, -7.073488235473633, -7.1788787841796875, -7.389683723449707, -7.525317192077637, -7.5885772705078125, -7.763050079345703, -8.224308967590332], [-0.372489869594574, -1.827528953552246, -4.3061370849609375, -4.838803291320801, -4.989402770996094, -5.1117401123046875, -5.187068939208984, -5.250604629516602, -5.662858963012695, -5.694143295288086], [-0.005904730875045061, -5.495633602142334, -7.522153377532959, -8.207719802856445, -8.449649810791016, -8.90707778930664, -9.541242599487305, -9.691887855529785, -10.111766815185547, -10.492300033569336], [-0.011737688444554806, -6.031045913696289, -7.209378242492676, -7.243843078613281, -7.747928619384766, -7.977005958557129, -8.321712493896484, -8.35726547241211, -8.502359390258789, -8.701475143432617], [-0.1878184974193573, -1.8141651153564453, -5.624242782592773, -6.579645156860352, -7.822257995605469, -8.386683464050293, -8.49909496307373, -8.57803726196289, -8.823796272277832, -8.960565567016602], [-0.10250605642795563, -3.402726650238037, -3.946469783782959, -5.159356594085693, -5.411166667938232, -6.1134867668151855, -6.131921291351318, -6.200311183929443, -6.424906253814697, -7.0071539878845215], [-0.1382717341184616, -2.544334888458252, -3.050598621368408, -7.775615215301514, -8.049697875976562, -8.075815200805664, -8.107090950012207, -8.183040618896484, -8.868163108825684, -8.878026962280273], [-0.0321299247443676, -3.6293933391571045, -7.695366382598877, -7.812208652496338, -7.8907952308654785, -8.088072776794434, -8.26425552368164, -9.077349662780762, -9.536680221557617, -9.683252334594727], [-1.843049168586731, -2.4128103256225586, -2.9644947052001953, -3.615971565246582, -4.061753273010254, -4.090791702270508, -4.206696510314941, -4.35806941986084, -4.626364707946777, -4.680699348449707], [-0.5746181607246399, -1.2087008953094482, -3.0003554821014404, -3.3169009685516357, -3.463759183883667, -5.873934745788574, -5.896479606628418, -6.705965995788574, -6.999606132507324, -7.05670166015625], [-0.4282326400279999, -2.4605259895324707, -2.6223549842834473, -3.1082043647766113, -3.1714062690734863, -4.331737995147705, -4.498405933380127, -4.684654712677002, -4.858206272125244, -5.0235915184021], [-0.30328744649887085, -1.6684131622314453, -3.2264671325683594, -4.789220809936523, -5.672477722167969, -5.96893310546875, -6.063676834106445, -6.157503128051758, -6.438438415527344, -6.866046905517578], [-0.8339848518371582, -1.6514954566955566, -2.2883095741271973, -2.579624652862549, -4.058284282684326, -4.280709743499756, -4.445773601531982, -4.500208377838135, -4.837723255157471, -4.86464262008667], [-2.4236323833465576, -3.5779502391815186, -3.747347593307495, -3.8657243251800537, -3.9251821041107178, -4.05374813079834, -4.170233726501465, -4.246179580688477, -4.283749580383301, -4.3532514572143555], [-3.0394649505615234, -3.325023651123047, -3.4425315856933594, -3.482461929321289, -3.5253353118896484, -3.5320472717285156, -3.536334991455078, -3.5988540649414062, -3.6118717193603516, -3.767345428466797], [-0.9748451113700867, -3.445889711380005, -3.5092709064483643, -3.8227875232696533, -4.129208564758301, -4.215622901916504, -4.424491882324219, -4.451557159423828, -4.580273628234863, -4.590479850769043], [-0.042138077318668365, -4.544243812561035, -4.993100166320801, -5.2644147872924805, -6.0866193771362305, -6.586895942687988, -6.6464128494262695, -6.650232315063477, -7.103418350219727, -7.140406608581543], [-1.5009835958480835, -1.9959081411361694, -2.2056002616882324, -2.2982020378112793, -2.3899550437927246, -3.7624144554138184, -4.033015727996826, -4.047290325164795, -4.302596569061279, -4.305589199066162], [-0.8899617195129395, -1.8940033912658691, -2.1531224250793457, -2.6077370643615723, -3.57206392288208, -4.008375644683838, -4.123062610626221, -4.19462251663208, -4.703207492828369, -4.783227443695068], [-0.6426234245300293, -1.8766722679138184, -1.905128002166748, -3.432450771331787, -3.950042247772217, -4.064731121063232, -4.408097743988037, -4.558604717254639, -4.580577373504639, -4.752116680145264], [-0.9895710945129395, -1.2832303047180176, -2.834984302520752, -3.317976474761963, -3.6336464881896973, -3.880049228668213, -4.0731587409973145, -4.0978169441223145, -4.23159646987915, -4.357437610626221], [-1.4843578338623047, -2.5058727264404297, -2.769329071044922, -2.864400863647461, -2.9782562255859375, -3.1617069244384766, -3.1705169677734375, -3.6056766510009766, -3.7407455444335938, -3.7464637756347656], [-0.35668182373046875, -2.5939693450927734, -2.6262855529785156, -3.3084774017333984, -3.5330371856689453, -4.559474945068359, -4.66771125793457, -4.916624069213867, -5.343305587768555, -5.446603775024414], [-0.7548707127571106, -1.6354477405548096, -1.728358507156372, -2.720416307449341, -3.955390214920044, -4.285571575164795, -5.190505504608154, -5.417296886444092, -5.692910671234131, -5.70766019821167], [-2.2759742736816406, -2.447378158569336, -3.347470283508301, -3.3950443267822266, -3.646815299987793, -3.908966064453125, -3.999274253845215, -4.045774459838867, -4.176496505737305, -4.207118988037109], [-0.5946731567382812, -1.5350360870361328, -2.5619163513183594, -2.7717952728271484, -3.591419219970703, -4.287841796875, -4.993228912353516, -5.0897979736328125, -5.203357696533203, -5.756500244140625], [-1.2474231719970703, -1.3613033294677734, -2.1486759185791016, -2.4232349395751953, -2.6617660522460938, -3.4216670989990234, -3.936910629272461, -4.392925262451172, -4.43516731262207, -4.45884895324707], [-0.09278443455696106, -4.05391263961792, -4.28864049911499, -4.342662334442139, -5.320607662200928, -5.534726619720459, -5.752249240875244, -5.930094242095947, -6.490260601043701, -6.499328136444092], [-0.7439873218536377, -1.3276169300079346, -3.703294038772583, -3.781306505203247, -3.9016640186309814, -4.004000663757324, -4.276181221008301, -4.478531837463379, -5.095404624938965, -5.32029914855957], [-0.42122718691825867, -2.8185622692108154, -3.101483106613159, -3.188938856124878, -3.4723422527313232, -3.6380908489227295, -4.098782062530518, -4.378454685211182, -4.669825077056885, -4.683760166168213], [-0.06834664195775986, -2.947554349899292, -5.053656101226807, -6.055807590484619, -6.784708499908447, -7.855190753936768, -8.000382423400879, -8.236613273620605, -8.529458045959473, -8.666272163391113], [-0.889788031578064, -1.9251807928085327, -2.3246865272521973, -2.6348557472229004, -3.4432435035705566, -3.7437500953674316, -3.7707600593566895, -3.8446316719055176, -3.9515233039855957, -4.112818241119385], [-0.0339607410132885, -4.20205545425415, -6.14114236831665, -6.209860324859619, -6.44116735458374, -6.61282205581665, -6.623001575469971, -6.731749057769775, -6.739603519439697, -7.492265224456787], [-0.9458173513412476, -2.315518856048584, -2.3769545555114746, -2.8430075645446777, -2.9142165184020996, -3.017808437347412, -3.0437026023864746, -3.5895285606384277, -4.342545986175537, -4.3558669090271], [-1.278747797012329, -1.6791880130767822, -2.243154764175415, -2.6780436038970947, -2.9457952976226807, -3.3872568607330322, -3.8447487354278564, -3.8778812885284424, -4.074468612670898, -4.290727615356445], [-1.6544435024261475, -2.5747010707855225, -2.6344716548919678, -2.7110230922698975, -2.7909810543060303, -2.9975812435150146, -3.1332204341888428, -3.418637990951538, -3.620513677597046, -3.8255879878997803], [-0.6308883428573608, -2.030764102935791, -2.5993924140930176, -2.7181267738342285, -3.8265442848205566, -4.051680088043213, -4.209730625152588, -4.533054828643799, -4.5756001472473145, -4.666740894317627], [-1.1232950687408447, -2.6717894077301025, -2.8861582279205322, -2.978487253189087, -2.99068284034729, -3.0478804111480713, -3.5038721561431885, -3.5882437229156494, -3.6518137454986572, -3.6590216159820557], [-0.29644775390625, -1.4972686767578125, -4.235706329345703, -6.1917266845703125, -6.29533576965332, -6.35308837890625, -6.444324493408203, -6.609521865844727, -6.973005294799805, -7.24797248840332], [-0.5543937683105469, -2.132505416870117, -2.372150421142578, -2.815032958984375, -3.2442665100097656, -3.6007652282714844, -3.939544677734375, -4.62982177734375, -4.962799072265625, -5.13623046875], [-0.09781533479690552, -3.6927363872528076, -3.9080111980438232, -4.1939592361450195, -5.095574378967285, -5.54971981048584, -6.315329551696777, -6.684817314147949, -6.770205497741699, -6.805417060852051], [-1.2705544233322144, -1.3966587781906128, -2.181002616882324, -2.669062614440918, -2.9217100143432617, -3.26711368560791, -3.413538932800293, -3.883183479309082, -4.156699180603027, -4.2995405197143555], [-0.007819517515599728, -6.477696895599365, -6.587197780609131, -6.76526403427124, -6.938920497894287, -7.21808385848999, -7.243297100067139, -8.155594825744629, -9.206570625305176, -9.448287010192871], [-0.18164236843585968, -2.3973500728607178, -4.385944366455078, -5.066688537597656, -5.128849029541016, -5.214832305908203, -5.355841636657715, -5.812589645385742, -6.640542030334473, -6.641702651977539], [-0.3554803729057312, -1.8430063724517822, -2.637809991836548, -4.232261657714844, -4.252098083496094, -5.280946731567383, -5.581424713134766, -5.7038726806640625, -6.054256439208984, -6.070062637329102], [-1.174675703048706, -1.1769263744354248, -2.3811194896698, -3.574639081954956, -3.6462228298187256, -4.428581237792969, -4.515603065490723, -4.629171371459961, -4.8219404220581055, -4.826866149902344], [-1.2866473197937012, -1.3665785789489746, -1.450305461883545, -2.815547466278076, -3.4129652976989746, -3.5802950859069824, -4.259646892547607, -4.410497188568115, -4.686326503753662, -4.819680690765381], [-1.284511923789978, -1.6003555059432983, -1.8806651830673218, -2.330559253692627, -2.822094440460205, -2.978271961212158, -3.0513157844543457, -3.384586811065674, -4.678221225738525, -4.7065606117248535], [-0.003452533157542348, -6.71318244934082, -7.6864166259765625, -7.866558074951172, -8.395820617675781, -8.797486305236816, -9.47498893737793, -10.08416748046875, -10.08804702758789, -10.378552436828613], [-0.7611655592918396, -1.778707504272461, -2.318513870239258, -3.9089527130126953, -4.0619659423828125, -4.076770782470703, -4.324800491333008, -4.37451171875, -4.376003265380859, -4.376029968261719], [-0.43409377336502075, -2.0129549503326416, -3.2345545291900635, -3.5926191806793213, -3.671215295791626, -3.764213800430298, -4.053459644317627, -4.160521030426025, -4.5704731941223145, -4.845633029937744], [-0.19858379662036896, -1.7301313877105713, -6.944765090942383, -7.969247817993164, -8.085386276245117, -8.19151496887207, -8.449525833129883, -9.160225868225098, -9.588837623596191, -9.984594345092773], [-1.8052799701690674, -2.5017712116241455, -2.586858034133911, -2.8280861377716064, -2.8637516498565674, -2.86714768409729, -3.0553886890411377, -3.674962282180786, -3.7268106937408447, -3.859431505203247], [-0.5131959915161133, -1.924729347229004, -2.9035139083862305, -2.921065330505371, -2.9853811264038086, -3.6293020248413086, -4.092734336853027, -5.427997589111328, -5.525561332702637, -5.584092140197754], [-1.3965237140655518, -1.8620693683624268, -2.1119511127471924, -2.61291241645813, -3.2083914279937744, -3.618133783340454, -3.6918938159942627, -3.8739960193634033, -4.162569046020508, -4.266186714172363], [-1.1275970935821533, -1.3566906452178955, -3.086838960647583, -3.452552080154419, -3.5200703144073486, -3.6519806385040283, -3.8299877643585205, -4.088089942932129, -4.1960649490356445, -4.667270660400391], [-0.4600834846496582, -1.3127923011779785, -3.765890598297119, -4.479422092437744, -4.7028889656066895, -4.795820713043213, -5.030157566070557, -5.468201160430908, -5.512045383453369, -5.532801151275635], [-1.2519283294677734, -1.8291149139404297, -2.6022262573242188, -2.739431381225586, -2.9429473876953125, -3.558868408203125, -3.7740249633789062, -4.217281341552734, -4.22056770324707, -4.551769256591797], [-0.5379815697669983, -2.2664647102355957, -2.621513843536377, -2.739819049835205, -4.065554141998291, -4.312038898468018, -4.4349846839904785, -4.5212788581848145, -4.59879732131958, -4.6161580085754395], [-0.24191200733184814, -2.357131004333496, -2.582028388977051, -5.0107927322387695, -5.192748069763184, -5.412188529968262, -5.6585588455200195, -5.699042320251465, -5.758770942687988, -6.2831010818481445], [-0.050112128257751465, -3.3379197120666504, -5.641367435455322, -6.5138840675354, -6.607151508331299, -6.669376850128174, -7.108365535736084, -7.35122537612915, -8.149378776550293, -8.19774341583252], [-0.2582389712333679, -2.471951723098755, -3.396174669265747, -3.807959794998169, -3.808306932449341, -4.562500953674316, -5.042641639709473, -5.083487510681152, -5.352255821228027, -5.45950984954834], [-0.2726271152496338, -2.563659906387329, -2.6477510929107666, -3.4699480533599854, -3.56538987159729, -4.310754776000977, -6.029754638671875, -6.037014007568359, -6.406652450561523, -6.706075668334961], [-0.0036967527121305466, -5.79642915725708, -8.14664077758789, -9.461244583129883, -9.964672088623047, -10.01359748840332, -10.15234375, -11.19976806640625, -11.389558792114258, -11.412906646728516], [-1.3802396059036255, -1.479141354560852, -1.7676469087600708, -2.4589576721191406, -2.7227916717529297, -3.4538612365722656, -3.771778106689453, -4.313604354858398, -4.533607482910156, -4.815101623535156], [-0.0841309204697609, -3.2644708156585693, -4.228188991546631, -5.183648586273193, -5.896180629730225, -6.039113521575928, -6.171563625335693, -6.318734645843506, -6.36781644821167, -6.743028163909912], [-0.2532283365726471, -2.2264645099639893, -3.574834108352661, -4.541092872619629, -4.606949806213379, -4.977408409118652, -5.337008476257324, -5.488327980041504, -5.846911430358887, -6.011813163757324], [-0.8177525997161865, -1.1370179653167725, -2.4109609127044678, -3.2797353267669678, -4.361211776733398, -4.442050933837891, -4.475515365600586, -4.667137145996094, -4.722789764404297, -4.750612258911133], [-0.0016505202511325479, -7.155468463897705, -8.78048038482666, -9.2362699508667, -9.369120597839355, -9.78431224822998, -9.897849082946777, -10.091814994812012, -10.389817237854004, -10.867840766906738], [-0.2397218495607376, -2.567783832550049, -3.9733777046203613, -4.0111870765686035, -4.055180072784424, -4.189627170562744, -4.272867679595947, -4.729145526885986, -4.93442964553833, -5.013723850250244], [-0.6482017636299133, -1.9497230052947998, -2.744734525680542, -2.9432570934295654, -3.294945478439331, -3.4912421703338623, -3.83294939994812, -3.9057776927948, -4.058547019958496, -4.493414878845215], [-0.9068384170532227, -2.589909553527832, -2.8801259994506836, -2.9731664657592773, -2.984347343444824, -3.2325143814086914, -3.6160764694213867, -3.652400016784668, -3.7895421981811523, -3.795302391052246], [-0.48999089002609253, -2.4962146282196045, -2.6280620098114014, -2.8122050762176514, -3.834912061691284, -4.068287372589111, -4.680363178253174, -4.7075276374816895, -4.838058948516846, -5.1276021003723145], [-0.2099856436252594, -2.396739959716797, -3.7827606201171875, -4.102975845336914, -4.2533721923828125, -4.32032585144043, -5.431510925292969, -5.517753601074219, -5.60052490234375, -5.608646392822266], [-1.2549582719802856, -1.7463961839675903, -1.7806674242019653, -1.8416491746902466, -3.588189125061035, -4.121170997619629, -4.203289985656738, -4.2304792404174805, -4.241677284240723, -4.364739418029785], [-1.1053345203399658, -1.2426445484161377, -2.0468666553497314, -2.3002541065216064, -3.6286728382110596, -3.8687584400177, -4.098183631896973, -4.374518394470215, -4.560420036315918, -4.8011064529418945], [-0.11845239996910095, -3.2166216373443604, -3.4023592472076416, -4.34104585647583, -5.193403720855713, -5.461124897003174, -5.67779016494751, -5.840208530426025, -6.054356098175049, -6.985493183135986], [-1.5437091588974, -1.820473074913025, -1.8392070531845093, -2.431112766265869, -2.857800006866455, -3.21336030960083, -3.728994846343994, -3.8952889442443848, -4.184719562530518, -4.387455463409424], [-1.3339064121246338, -1.81770920753479, -1.9535925388336182, -2.603487253189087, -2.773926019668579, -2.8758318424224854, -2.9865000247955322, -3.459209680557251, -3.6767351627349854, -4.189197540283203], [-1.3527289628982544, -1.379508137702942, -1.534403920173645, -3.617396354675293, -3.7183008193969727, -4.081496238708496, -4.251824378967285, -4.358290672302246, -4.5153303146362305, -4.761061668395996], [-0.8167728185653687, -2.3844056129455566, -3.1805481910705566, -3.185965061187744, -3.335096836090088, -3.715414524078369, -4.072058200836182, -4.121044635772705, -4.12502908706665, -4.142355442047119], [-1.7069964408874512, -1.9735217094421387, -2.8056349754333496, -3.0393080711364746, -3.0493712425231934, -3.0742087364196777, -3.0905661582946777, -3.3695883750915527, -3.4065794944763184, -3.4762511253356934], [-0.7300421595573425, -1.9341418743133545, -1.9869163036346436, -3.066288709640503, -3.2312915325164795, -3.742283582687378, -4.1482014656066895, -4.171370029449463, -4.528067111968994, -4.687931537628174], [-0.9212523698806763, -2.6654233932495117, -2.924098014831543, -3.02217960357666, -3.2975587844848633, -3.307803153991699, -3.327866554260254, -3.7506399154663086, -3.7806501388549805, -3.8221845626831055], [-0.17896340787410736, -3.3544814586639404, -4.439797878265381, -4.7614216804504395, -4.931570529937744, -5.003124713897705, -5.3462090492248535, -5.386402606964111, -5.432949542999268, -5.478273868560791], [-0.18972711265087128, -2.5467216968536377, -4.389462947845459, -4.618724346160889, -4.668748378753662, -5.214345455169678, -5.261071681976318, -5.340207576751709, -5.428225994110107, -5.518626689910889], [-0.7267122864723206, -1.9803457260131836, -2.622727394104004, -2.9903154373168945, -2.998072624206543, -3.2191457748413086, -3.8306455612182617, -3.894312858581543, -3.975949287414551, -4.350564002990723], [-0.020010385662317276, -5.251017093658447, -5.67455530166626, -5.886524677276611, -6.079123020172119, -6.6986165046691895, -6.785297870635986, -7.16776704788208, -7.767279148101807, -8.13803482055664], [-0.7513647675514221, -1.7059850692749023, -2.2032651901245117, -3.0225114822387695, -3.234866142272949, -3.852593421936035, -4.09366512298584, -4.310198783874512, -4.543339729309082, -4.547806739807129], [-0.8763109445571899, -1.5423017740249634, -1.8642431497573853, -2.4766430854797363, -3.7231316566467285, -3.9272236824035645, -4.459427356719971, -4.498356342315674, -4.714981555938721, -4.9154839515686035], [-0.00046659549116156995, -8.833687782287598, -8.882493019104004, -10.132534980773926, -10.212037086486816, -11.393338203430176, -11.408698081970215, -11.511855125427246, -11.546452522277832, -12.114046096801758], [-1.5360831022262573, -2.108304977416992, -2.2858123779296875, -2.431581497192383, -2.8270797729492188, -3.043304443359375, -3.1750946044921875, -3.4131622314453125, -3.6463184356689453, -3.6590728759765625], [-1.2832740545272827, -1.382341742515564, -1.5589431524276733, -2.1615853309631348, -3.652876377105713, -3.682203769683838, -3.9527363777160645, -4.271149158477783, -4.416624546051025, -5.082592487335205], [-1.2193286418914795, -1.3985011577606201, -1.9321696758270264, -2.4761435985565186, -2.815045118331909, -2.8223140239715576, -3.122265577316284, -4.1949872970581055, -4.248230934143066, -4.835675239562988], [-0.8976349234580994, -1.5094265937805176, -2.025852680206299, -2.90116548538208, -3.4823861122131348, -3.9814553260803223, -4.281469821929932, -4.530375003814697, -4.63230562210083, -4.833662509918213], [-0.5017741918563843, -1.735287070274353, -1.9957107305526733, -4.081532001495361, -4.8460869789123535, -4.859695911407471, -5.334743976593018, -5.365137577056885, -5.4647088050842285, -5.676289081573486], [-1.3731788396835327, -2.5181336402893066, -3.075880527496338, -3.1481194496154785, -3.295963764190674, -3.4699769020080566, -3.5433621406555176, -3.543529987335205, -3.7463364601135254, -3.8233399391174316], [-0.109302818775177, -3.57485032081604, -3.641247034072876, -4.436283588409424, -4.822815418243408, -5.259323596954346, -5.804226398468018, -6.221785068511963, -6.34154748916626, -6.366329669952393], [-0.03258315473794937, -3.510650634765625, -7.259967803955078, -7.543283462524414, -7.955667495727539, -8.084836959838867, -8.469537734985352, -10.914033889770508, -12.09141731262207, -12.289237976074219], [-1.7006688117980957, -2.4590649604797363, -2.4627881050109863, -2.4675965309143066, -2.7641453742980957, -3.0670418739318848, -3.1854519844055176, -3.293827533721924, -3.3986802101135254, -3.563965320587158], [-0.291087806224823, -1.8130090236663818, -4.385949611663818, -4.612578868865967, -4.788232326507568, -5.132178783416748, -5.208011150360107, -5.2184977531433105, -5.474307537078857, -5.593989849090576], [-0.719337522983551, -0.9613209366798401, -2.8284263610839844, -4.439323425292969, -4.978010177612305, -5.291358947753906, -5.8328752517700195, -6.318382263183594, -6.369429588317871, -6.421497344970703], [-1.7802187204360962, -2.1415505409240723, -2.3215088844299316, -2.900228977203369, -2.9053006172180176, -3.037436008453369, -3.088858127593994, -3.136812686920166, -3.419224262237549, -3.5133137702941895], [-0.8869927525520325, -1.9240639209747314, -2.519213914871216, -3.080664873123169, -3.316546678543091, -3.67718243598938, -4.00349760055542, -4.049285411834717, -4.193301677703857, -4.2090067863464355], [-2.1019766330718994, -2.2705576419830322, -2.502298593521118, -2.7492926120758057, -2.838247537612915, -2.9228575229644775, -3.1833651065826416, -3.183605432510376, -3.2752249240875244, -3.5262339115142822], [-1.239345908164978, -1.9251044988632202, -2.3547024726867676, -3.1700310707092285, -3.5397591590881348, -3.562124729156494, -3.980347156524658, -4.103588581085205, -4.113361835479736, -4.117939472198486], [-0.09708235412836075, -3.6577627658843994, -4.277444839477539, -4.357044219970703, -4.504293441772461, -4.634490966796875, -5.393272399902344, -5.601314544677734, -6.300224304199219, -6.547618865966797], [-0.3409811854362488, -2.5129554271698, -2.8377902507781982, -3.1673457622528076, -3.9924628734588623, -4.179958820343018, -4.2811970710754395, -4.36254358291626, -4.719370365142822, -5.219009876251221], [-0.7249959707260132, -2.150012493133545, -2.2308554649353027, -2.3732733726501465, -2.387692928314209, -3.5528311729431152, -3.7483725547790527, -4.235726833343506, -4.596755504608154, -5.085277080535889], [-0.9774425029754639, -1.8530681133270264, -2.411982297897339, -2.6150577068328857, -2.7503421306610107, -3.4215800762176514, -3.88342547416687, -3.886164426803589, -4.383856773376465, -4.672942161560059], [-0.04679347202181816, -3.96494722366333, -4.641136646270752, -4.795372486114502, -6.8315510749816895, -6.982397556304932, -7.268871784210205, -7.3839497566223145, -7.489758014678955, -7.857241153717041], [-0.31035345792770386, -2.953282594680786, -3.2479069232940674, -3.761418581008911, -4.112614631652832, -4.1251020431518555, -4.619383811950684, -4.64388370513916, -4.796278953552246, -4.85745906829834], [-0.26277634501457214, -2.1181726455688477, -2.456051826477051, -5.3911638259887695, -5.914050102233887, -6.005043983459473, -6.199723243713379, -6.254300117492676, -6.544785499572754, -6.592713356018066], [-0.5433505773544312, -2.7175393104553223, -2.842195987701416, -3.76133394241333, -3.8736729621887207, -3.977125644683838, -4.064489841461182, -4.075668811798096, -4.150932788848877, -4.173208713531494], [-1.4341762065887451, -1.5299479961395264, -1.9002444744110107, -2.6758296489715576, -2.7313029766082764, -2.8934810161590576, -2.8967063426971436, -3.099146604537964, -3.6680171489715576, -4.6420488357543945], [-1.220574975013733, -1.2577797174453735, -2.7986960411071777, -2.9994406700134277, -3.221740245819092, -3.6000800132751465, -3.7003417015075684, -3.7338995933532715, -4.461680889129639, -4.544175624847412], [-0.23072941601276398, -3.5346615314483643, -3.7650253772735596, -4.1820807456970215, -4.230477809906006, -4.369567394256592, -4.410802364349365, -4.443854808807373, -4.6673903465271, -4.708890438079834], [-0.1686599999666214, -2.5616137981414795, -2.7843387126922607, -5.496199607849121, -5.517390251159668, -6.10916805267334, -7.133181571960449, -7.469878196716309, -7.567255973815918, -7.571822166442871], [-0.03456159308552742, -3.5091214179992676, -6.727182865142822, -7.896212100982666, -7.901533603668213, -8.239046096801758, -8.441466331481934, -8.540034294128418, -8.729623794555664, -8.793167114257812], [-0.041516028344631195, -4.718027591705322, -4.776236057281494, -5.768055438995361, -5.774538516998291, -5.931726932525635, -6.047996997833252, -6.555778980255127, -6.7127366065979, -7.016510486602783], [-0.16437621414661407, -2.4174718856811523, -3.868739128112793, -4.466994285583496, -4.826800346374512, -6.00504207611084, -6.166232109069824, -6.271721839904785, -6.527167320251465, -6.677702903747559], [-0.9411216974258423, -1.8419796228408813, -2.60158109664917, -2.783629894256592, -2.8400683403015137, -3.357631206512451, -3.8178515434265137, -4.265321254730225, -4.463851451873779, -4.62208890914917], [-1.1664338111877441, -1.4504914283752441, -1.9690041542053223, -2.520834445953369, -2.894716739654541, -3.2474350929260254, -4.2989277839660645, -4.555185794830322, -4.65734338760376, -4.66165018081665], [-0.6082432866096497, -3.0967419147491455, -3.2729732990264893, -3.279273271560669, -3.345717668533325, -3.5865261554718018, -3.672593355178833, -3.872760057449341, -4.08180570602417, -4.164279460906982], [-1.1109569072723389, -2.094578504562378, -2.3004329204559326, -2.5082404613494873, -2.611309766769409, -3.4852187633514404, -3.604609251022339, -3.620774030685425, -3.7734286785125732, -3.90071177482605], [-0.6794657707214355, -2.0996108055114746, -2.626533031463623, -2.848659038543701, -3.068887233734131, -3.1731066703796387, -4.053109645843506, -4.519050121307373, -4.584981441497803, -5.241530895233154], [-1.4911484718322754, -1.6066193580627441, -2.219088077545166, -2.6915860176086426, -3.222395420074463, -3.320692539215088, -3.3744683265686035, -3.4580636024475098, -3.7232441902160645, -3.782402515411377], [-0.8373185396194458, -1.3131829500198364, -2.0079612731933594, -3.3321914672851562, -3.5005569458007812, -3.81939697265625, -4.4671783447265625, -4.592998504638672, -4.895544052124023, -5.025142669677734], [-0.0037339269183576107, -6.233117580413818, -7.138347148895264, -8.172951698303223, -8.672486305236816, -9.589415550231934, -9.762785911560059, -10.199854850769043, -10.444130897521973, -10.599610328674316], [-1.1195952892303467, -1.2713172435760498, -2.1387641429901123, -2.4109504222869873, -2.4467132091522217, -3.8374545574188232, -4.7074689865112305, -4.770272254943848, -4.886166572570801, -5.233708381652832], [-1.1844749450683594, -1.3241844177246094, -2.4885005950927734, -2.9741172790527344, -3.170734405517578, -3.346952438354492, -4.089082717895508, -4.3595733642578125, -4.749689102172852, -4.86265754699707], [-1.0153727531433105, -2.1631999015808105, -2.7604708671569824, -2.8160014152526855, -3.2777857780456543, -3.427405834197998, -3.5600314140319824, -4.066863536834717, -4.097205638885498, -4.10852575302124], [-0.6537575125694275, -1.7951149940490723, -2.761894702911377, -3.5116944313049316, -3.9450478553771973, -3.996330738067627, -4.053154468536377, -4.572994709014893, -4.786819934844971, -4.798311710357666], [-0.7565585970878601, -1.420966386795044, -1.9133732318878174, -3.0086758136749268, -3.823054552078247, -4.280617713928223, -4.372963905334473, -5.051085472106934, -5.195734977722168, -5.829764366149902], [-0.8668451309204102, -1.5346269607543945, -1.8505373001098633, -3.065587043762207, -3.361924171447754, -3.3716936111450195, -3.9786481857299805, -4.682793617248535, -4.948409080505371, -5.08335018157959], [-0.14208939671516418, -3.7619662284851074, -3.879697322845459, -4.157227993011475, -4.400855541229248, -4.402129650115967, -4.683719158172607, -6.124358654022217, -6.445504665374756, -6.534369945526123], [-0.0002898749662563205, -9.769424438476562, -9.807418823242188, -10.319768905639648, -10.944612503051758, -10.993598937988281, -11.133970260620117, -11.22636890411377, -11.607852935791016, -11.779435157775879], [-0.5908623933792114, -2.5130348205566406, -2.765125274658203, -3.317281723022461, -3.5334625244140625, -4.115068435668945, -4.217123031616211, -4.22657585144043, -4.311502456665039, -4.336002349853516], [-0.17303341627120972, -3.0575225353240967, -3.922276258468628, -4.146698474884033, -4.581213474273682, -4.748531818389893, -4.849279880523682, -5.442572116851807, -5.494207859039307, -5.858719348907471], [-0.4503314197063446, -1.7294681072235107, -1.9426944255828857, -5.353464603424072, -5.516043186187744, -5.613363742828369, -5.796005725860596, -6.252970218658447, -6.364943027496338, -6.636854648590088], [-1.476321816444397, -1.561669945716858, -1.6821247339248657, -2.082017421722412, -2.1354479789733887, -2.613887310028076, -4.449899196624756, -4.539252758026123, -4.73621129989624, -5.2259039878845215], [-0.47746407985687256, -1.336346983909607, -3.1393275260925293, -4.509477138519287, -4.707902431488037, -4.786777019500732, -5.018617153167725, -5.558213710784912, -5.659398555755615, -6.044270992279053], [-1.6391061544418335, -2.1922335624694824, -2.774014949798584, -2.883714199066162, -2.9082236289978027, -2.9262022972106934, -3.0423922538757324, -3.3938002586364746, -3.4541468620300293, -3.458777904510498], [-0.02097589708864689, -4.337897300720215, -6.0606584548950195, -6.6927900314331055, -7.1293535232543945, -7.134495735168457, -7.27060604095459, -8.015889167785645, -8.567399024963379, -8.801312446594238], [-1.750502109527588, -2.045637607574463, -2.1107640266418457, -2.6134114265441895, -2.8450303077697754, -2.9110817909240723, -3.455650806427002, -3.7072300910949707, -3.708134174346924, -3.850904941558838], [-1.093828558921814, -1.169552206993103, -1.618536353111267, -2.87418794631958, -3.158970355987549, -3.925438404083252, -4.961788654327393, -5.144516468048096, -5.226093769073486, -5.271728992462158], [-0.6131972670555115, -1.8658866882324219, -2.665773391723633, -2.8503684997558594, -3.678285598754883, -4.053487777709961, -4.063518524169922, -4.2359466552734375, -4.6979217529296875, -4.9517364501953125], [-0.5003752112388611, -2.489387035369873, -2.9834933280944824, -3.7678942680358887, -3.893876552581787, -4.040996074676514, -4.1345038414001465, -4.317748546600342, -4.766955852508545, -5.0517497062683105], [-0.4996608793735504, -1.5757774114608765, -3.262159824371338, -3.665256977081299, -3.8835339546203613, -3.91666841506958, -3.9641880989074707, -4.392263889312744, -4.9200921058654785, -5.414328098297119], [-0.5248868465423584, -1.9127118587493896, -2.4744832515716553, -3.450437307357788, -3.669299840927124, -3.9967362880706787, -4.019338607788086, -4.04057502746582, -4.48675537109375, -4.563028335571289], [-0.10120955854654312, -2.9748401641845703, -3.215791702270508, -6.222570419311523, -6.512105941772461, -7.740865707397461, -8.747190475463867, -8.781150817871094, -8.82335090637207, -9.702880859375], [-4.768370445162873e-07, -15.009653091430664, -16.646602630615234, -16.90728187561035, -17.88405990600586, -18.06052589416504, -18.22677230834961, -18.581607818603516, -18.613759994506836, -18.632862091064453], [-0.9634779691696167, -2.034724235534668, -2.7465476989746094, -2.75234317779541, -3.5132293701171875, -3.558762550354004, -3.938291549682617, -3.9997920989990234, -4.167990684509277, -4.170927047729492], [-2.3194642066955566, -2.3558993339538574, -3.1568589210510254, -3.2147393226623535, -3.383246898651123, -3.615540027618408, -3.8907132148742676, -3.9164681434631348, -4.205816745758057, -4.302371501922607], [-1.5479261875152588, -2.8916876316070557, -3.0237953662872314, -3.124971628189087, -3.2257044315338135, -3.494793176651001, -3.6809093952178955, -3.729083299636841, -4.260154724121094, -4.283388137817383], [-0.509814441204071, -2.6297810077667236, -2.7403366565704346, -2.9497692584991455, -3.5784332752227783, -3.7098934650421143, -3.7277138233184814, -4.083382606506348, -4.478659629821777, -4.5356645584106445], [-1.2678115367889404, -1.8093516826629639, -1.8235232830047607, -2.4452788829803467, -3.087082624435425, -3.365715742111206, -3.450453519821167, -3.870945692062378, -3.918255567550659, -4.0189714431762695], [-0.9515280723571777, -1.9114794731140137, -2.832017421722412, -3.0379271507263184, -3.337533473968506, -3.742887020111084, -3.8735642433166504, -4.123470783233643, -4.337832927703857, -4.367605686187744], [-2.171112060546875, -2.391653060913086, -2.7205562591552734, -3.2578487396240234, -3.4893417358398438, -3.6892318725585938, -3.927274703979492, -3.9953651428222656, -4.093984603881836, -4.14581298828125], [-2.4455368518829346, -3.310229539871216, -3.437636613845825, -3.5179636478424072, -3.666815996170044, -3.7329695224761963, -3.8520424365997314, -3.8666489124298096, -3.876596689224243, -3.926138162612915], [-0.7217168211936951, -1.3622503280639648, -2.69547176361084, -3.031609535217285, -3.521487236022949, -4.093205451965332, -4.318233489990234, -4.790000915527344, -4.795565605163574, -4.898306846618652], [-1.4549778699874878, -1.9336727857589722, -3.012394905090332, -3.132857322692871, -3.35697078704834, -3.3767499923706055, -3.4579572677612305, -3.627322196960449, -3.715487480163574, -3.8149595260620117], [-0.928566038608551, -1.7140579223632812, -2.603656768798828, -2.8133316040039062, -2.865650177001953, -3.3873939514160156, -3.5062122344970703, -3.648761749267578, -4.004087448120117, -4.115785598754883], [-1.6350857019424438, -2.005673885345459, -2.3206686973571777, -2.5934481620788574, -2.6495261192321777, -2.7588934898376465, -2.7885851860046387, -2.84981107711792, -3.4200778007507324, -3.8958468437194824], [-1.1532529592514038, -2.265956401824951, -2.4496169090270996, -2.573607921600342, -3.124345302581787, -3.4526925086975098, -3.6920504570007324, -3.8715567588806152, -4.009006023406982, -4.036871433258057], [-0.8533508777618408, -1.472703218460083, -2.4125759601593018, -2.9383881092071533, -2.9886906147003174, -3.2573254108428955, -3.607814073562622, -4.065762519836426, -4.508149147033691, -4.9412126541137695], [-1.1067936420440674, -1.4004642963409424, -2.7948715686798096, -3.1332638263702393, -3.7391464710235596, -3.763219118118286, -4.024504661560059, -4.059776306152344, -4.326423645019531, -4.336912155151367], [-0.3188457787036896, -2.4581260681152344, -3.908651351928711, -4.330844879150391, -4.500626564025879, -4.6908979415893555, -4.797201156616211, -5.000249862670898, -5.222867965698242, -5.3551025390625], [-0.8249402046203613, -1.939770221710205, -2.3196702003479004, -2.598604679107666, -2.8208889961242676, -3.946281909942627, -4.101963520050049, -4.345675945281982, -4.736853122711182, -5.155606746673584], [-0.7935125827789307, -2.069671869277954, -3.827415704727173, -3.9018747806549072, -3.939485788345337, -4.121503829956055, -4.289074897766113, -4.335172653198242, -4.424995422363281, -4.42669677734375], [-0.32164621353149414, -1.5706028938293457, -3.779066562652588, -3.8708577156066895, -5.218122959136963, -5.429989337921143, -6.545755863189697, -6.568811893463135, -6.768260478973389, -6.956370830535889], [-1.003499150276184, -2.687191963195801, -2.758955955505371, -3.494236946105957, -3.5270328521728516, -3.8594179153442383, -3.875253677368164, -3.8897104263305664, -4.095883369445801, -4.098232269287109], [-0.5608841180801392, -2.2943644523620605, -3.1340365409851074, -3.9575676918029785, -4.051490306854248, -4.403706073760986, -4.464087963104248, -4.687629222869873, -4.732462406158447, -5.0428242683410645], [-0.8024858236312866, -2.829089641571045, -3.211879253387451, -3.391812801361084, -3.670717716217041, -3.783231258392334, -3.8098196983337402, -3.9890332221984863, -4.193538188934326, -4.4100661277771], [-1.7592356204986572, -1.9216139316558838, -2.155390977859497, -2.5360605716705322, -3.9073526859283447, -4.033847808837891, -4.157377243041992, -4.460928916931152, -4.479830741882324, -4.518803596496582], [-0.9153603911399841, -1.3166208267211914, -2.6429128646850586, -2.8648805618286133, -3.445237159729004, -3.515660285949707, -3.517298698425293, -4.238512992858887, -4.666811943054199, -4.833283424377441], [-0.8651371002197266, -1.54229736328125, -2.7010269165039062, -2.711465835571289, -3.864898681640625, -4.455408096313477, -4.948841094970703, -5.005257606506348, -5.079590797424316, -5.1587982177734375], [-0.7731188535690308, -3.500368118286133, -3.5553665161132812, -3.5704174041748047, -3.7730846405029297, -4.074491500854492, -4.336993217468262, -4.725822448730469, -4.759237289428711, -4.848731994628906], [-2.587486743927002, -3.220287799835205, -3.265594959259033, -3.3050618171691895, -3.3536458015441895, -3.3653206825256348, -3.451712131500244, -3.5286126136779785, -3.535789966583252, -3.6473450660705566], [-1.966299057006836, -2.095020294189453, -2.0976333618164062, -2.108213424682617, -2.695648193359375, -3.207305908203125, -3.3674163818359375, -3.511425018310547, -3.5776424407958984, -3.7452011108398438], [-0.9955637454986572, -1.3299105167388916, -2.001983880996704, -3.7754557132720947, -3.804365396499634, -3.977809190750122, -4.028399467468262, -4.767955780029297, -4.9801225662231445, -4.997615814208984], [-0.9364675283432007, -1.1900686025619507, -2.182166576385498, -3.261643886566162, -4.056382656097412, -4.183980464935303, -4.428506374359131, -4.436883449554443, -4.692630290985107, -4.720468044281006], [-1.603662371635437, -2.2296066284179688, -2.953805923461914, -3.288052558898926, -3.507845878601074, -3.5429210662841797, -3.5826330184936523, -3.766613006591797, -3.838237762451172, -4.1531877517700195], [-0.8462379574775696, -1.6531550884246826, -1.6723926067352295, -2.608097791671753, -4.05698299407959, -4.460982322692871, -4.486870765686035, -4.9635515213012695, -5.140688896179199, -5.224724769592285], [-1.065568208694458, -2.0943844318389893, -2.1396725177764893, -2.6119625568389893, -3.013720750808716, -3.027959108352661, -3.326831102371216, -3.5820438861846924, -4.101555824279785, -4.295680046081543], [-0.9680259227752686, -1.9687373638153076, -3.1111209392547607, -3.258896589279175, -3.6214053630828857, -3.657491445541382, -3.832491636276245, -3.884389638900757, -4.020899772644043, -4.1801652908325195], [-1.3465521335601807, -1.4725191593170166, -2.9039080142974854, -3.4816410541534424, -3.4852802753448486, -3.6336138248443604, -3.6583521366119385, -3.7884113788604736, -4.055657386779785, -4.09886360168457], [-0.20639769732952118, -2.641072988510132, -3.4455068111419678, -3.996896505355835, -4.044372081756592, -5.007638454437256, -5.136912822723389, -5.801711559295654, -5.996934413909912, -6.154557704925537], [-0.1959957629442215, -1.9113335609436035, -4.625126361846924, -5.152451038360596, -5.468470096588135, -5.7141594886779785, -6.312652111053467, -7.214084148406982, -7.27617883682251, -7.899289608001709], [-0.8202694654464722, -1.7166012525558472, -2.189805030822754, -3.0368471145629883, -3.625382423400879, -3.78731632232666, -3.8527956008911133, -3.9533281326293945, -4.0142107009887695, -4.185128211975098], [-0.052747633308172226, -3.355314254760742, -5.081544876098633, -6.354413986206055, -6.50160026550293, -6.867157936096191, -6.906248092651367, -7.23007869720459, -7.279658317565918, -7.505866050720215], [-0.20388591289520264, -2.4745750427246094, -4.065881729125977, -4.1504974365234375, -4.277046203613281, -4.454280853271484, -4.961923599243164, -5.195476531982422, -5.232025146484375, -5.636140823364258], [-1.143463134765625, -1.4821357727050781, -2.4257431030273438, -3.2510194778442383, -3.505624771118164, -4.219663619995117, -4.289992332458496, -4.553607940673828, -4.5758161544799805, -4.618678092956543], [-0.013952808454632759, -4.854330539703369, -6.877737522125244, -7.375034809112549, -7.783509731292725, -8.322366714477539, -8.37386417388916, -8.458611488342285, -8.699201583862305, -8.817965507507324], [-2.1377251148223877, -2.157262086868286, -2.3355114459991455, -2.47631573677063, -2.66723370552063, -3.0581886768341064, -3.184077501296997, -3.4214928150177, -3.4249699115753174, -3.5257275104522705], [-0.14188191294670105, -2.6490936279296875, -4.979707717895508, -5.086414337158203, -5.405952453613281, -5.988659858703613, -6.213647842407227, -6.231888771057129, -6.284070014953613, -6.319009780883789], [-1.2951977252960205, -2.3064815998077393, -2.596562623977661, -2.7419216632843018, -2.888667345046997, -3.9139444828033447, -4.295862197875977, -4.311478614807129, -4.387269020080566, -4.395474433898926], [-0.8956736922264099, -2.3381900787353516, -2.6800975799560547, -2.9655284881591797, -3.6077518463134766, -3.6957340240478516, -3.8474884033203125, -3.8850440979003906, -3.986766815185547, -4.223785400390625], [-0.4706301987171173, -1.8411955833435059, -1.8983187675476074, -4.480752468109131, -4.5081915855407715, -4.704736232757568, -5.019006252288818, -5.580361843109131, -5.720889568328857, -5.744185924530029], [-0.23656286299228668, -2.469434976577759, -2.4718878269195557, -4.273067474365234, -5.49159049987793, -5.98747444152832, -6.138788223266602, -6.188631057739258, -6.302213668823242, -6.544244766235352], [-0.42535340785980225, -2.7721896171569824, -2.8670668601989746, -3.2564597129821777, -3.717986583709717, -4.038154125213623, -4.114224910736084, -4.304599285125732, -4.427661418914795, -4.567361354827881], [-1.8299320936203003, -2.492424964904785, -3.027043342590332, -3.230889320373535, -3.3434972763061523, -3.3455820083618164, -3.400029182434082, -3.4040613174438477, -3.535750389099121, -3.8819074630737305], [-1.115234136581421, -2.8380773067474365, -3.104570150375366, -3.1606156826019287, -3.4197957515716553, -3.625608205795288, -3.711219549179077, -3.8638875484466553, -4.068668365478516, -4.0814971923828125], [-1.6362870931625366, -1.7251046895980835, -1.7372373342514038, -2.5480685234069824, -2.6541266441345215, -3.454556941986084, -3.4559969902038574, -3.5205302238464355, -4.1082892417907715, -4.346578121185303]], \"topKTokens\": [[\"\\n\", \",\", \"Words\", \"Summary\", \"\\n\\n\", \"<|endoftext|>\", \" \", \"Features\", \"Random\", \" the\"], [\" upon\", \" there\", \",\", \" the\", \" she\", \" on\", \" he\", \" day\", \" in\", \" night\"], [\" a\", \" an\", \" the\", \" one\", \" time\", \" some\", \" upon\", \" it\", \" two\", \" something\"], [\" time\", \" day\", \" night\", \" morning\", \" week\", \" Sunday\", \" long\", \" Wednesday\", \" dark\", \" little\"], [\",\", \" there\", \" in\", \" a\", \" was\", \" upon\", \".\", \" on\", \" the\", \" lived\"], [\" there\", \" a\", \" in\", \" two\", \" the\", \" she\", \" it\", \" Jack\", \" he\", \" Sam\"], [\" was\", \" were\", \" lived\", \" are\", \" is\", \"'s\", \",\", \" a\", \" wasn\", \" had\"], [\" a\", \" an\", \" two\", \" one\", \" something\", \" some\", \" the\", \" three\", \" his\", \" Jack\"], [\" little\", \" girl\", \" boy\", \" big\", \" brave\", \" small\", \" man\", \" young\", \" kind\", \" happy\"], [\"my\", \" and\", \",\", \" who\", \".\", \"ma\", \" bear\", \" named\", \"o\", \" called\"], [\" and\", \" who\", \",\", \" bear\", \".\", \" bird\", \" pig\", \" duck\", \" rabbit\", \" bunny\"], [\" a\", \" her\", \" daddy\", \" she\", \" their\", \" his\", \" baby\", \" the\", \" an\", \" some\"], [\" baby\", \" daddy\", \" little\", \" dad\", \" boy\", \" daughter\", \" girl\", \" Daddy\", \" child\", \" man\"], [\",\", \" house\", \" bear\", \" tree\", \" dog\", \" garden\", \" truck\", \" park\", \" box\", \" mountain\"], [\" strong\", \" green\", \" tall\", \" expensive\", \" brown\", \" long\", \" happy\", \" red\", \" loud\", \" shiny\"], [\" house\", \" castle\", \" building\", \" car\", \" boat\", \" theater\", \" store\", \" TV\", \" train\", \" hotel\"], [\".\", \" with\", \"!\", \" in\", \" to\", \" that\", \",\", \" for\", \"!\\\"\", \" and\"], [\" It\", \" They\", \" She\", \" The\", \" Every\", \" One\", \" People\", \" In\", \" Everywhere\", \" This\"], [\" were\", \" had\", \" went\", \" wanted\", \" all\", \" flew\", \" needed\", \" both\", \" watched\", \" took\"], [\" lots\", \" a\", \" many\", \" planes\", \" the\", \" people\", \" so\", \" that\", \" big\", \" it\"], [\" people\", \" planes\", \" airplanes\", \" things\", \" big\", \" exciting\", \" different\", \" interesting\", \" aer\", \" new\"], [\" and\", \" walking\", \" going\", \" there\", \" rushing\", \",\", \" who\", \" taking\", \" in\", \" working\"], [\" around\", \" and\", \" by\", \" on\", \" in\", \" down\", \",\", \" through\", \" there\", \" from\"], [\" and\", \".\", \",\", \" the\", \" with\", \" in\", \" on\", \" them\", \" looking\", \"!\"], [\" talking\", \" looking\", \" having\", \" singing\", \" pointing\", \" were\", \" heard\", \" around\", \" taking\", \" wondered\"], [\" at\", \" to\", \".\", \" hello\", \" their\", \",\", \" flags\", \" back\", \" and\", \" for\"], [\" everyone\", \" them\", \" all\", \" people\", \" the\", \" their\", \" each\", \" some\", \" planes\", \" it\"], [\" the\", \" of\", \" their\", \".\", \" them\", \",\", \" and\", \" kinds\", \" they\", \" different\"], [\" people\", \" planes\", \" airplanes\", \" little\", \" different\", \" things\", \" big\", \" nice\", \" shops\", \" buildings\"], [\".\", \" in\", \",\", \" they\", \" and\", \" flying\", \" as\", \" before\", \" that\", \" she\"], [\".\", \"!\", \",\", \" and\", \" quickly\", \" before\", \" when\", \" much\", \" far\", \" excited\"], [\"\\n\", \" \", \" The\", \" They\", \" Then\", \" Suddenly\", \" Everyone\", \" One\", \" It\", \" Some\"], [\"\\n\", \"The\", \"Suddenly\", \" \", \"Then\", \"When\", \"Mom\", \"One\", \"\\\"\", \"At\"], [\"The\", \"One\", \"Suddenly\", \"Then\", \"When\", \"\\\"\", \"Mom\", \"At\", \"L\", \"Finally\"], [\" little\", \" family\", \" plane\", \" mom\", \" people\", \" airport\", \" pilot\", \" planes\", \" next\", \" modern\"], [\" was\", \" got\", \" had\", \" wanted\", \" went\", \" decided\", \" saw\", \" watched\", \" kept\", \" were\"], [\" the\", \" their\", \",\", \" a\", \" lots\", \" if\", \" one\", \" her\", \" each\", \" his\"], [\" little\", \" dad\", \" airport\", \" pilot\", \" mom\", \" man\", \" family\", \" people\", \" children\", \" 3\"], [\" girl\", \" boy\", \" 3\", \" one\", \" three\", \" man\", \" child\", \" plane\", \" bird\", \" daughter\"], [\",\", \" if\", \" to\", \" what\", \" who\", \" where\", \" why\", \" for\", \" a\", \" and\"], [\" her\", \" mom\", \" she\", \" the\", \" daddy\", \" Mom\", \" dad\", \" Daddy\", \" mum\", \" his\"], [\" if\", \",\", \" to\", \" what\", \" why\", \" lots\", \" many\", \" where\", \" and\", \" for\"], [\" of\", \"-\", \".\", \" and\", \" about\", \" to\", \",\", \" more\", \" at\", \" so\"], [\" questions\", \" people\", \" planes\", \" luggage\", \" other\", \" things\", \" friends\", \" different\", \" mom\", \" money\"], [\".\", \" about\", \",\", \" to\", \" and\", \" for\", \" in\", \"!\", \" all\", \" every\"], [\" but\", \" and\", \" so\", \" while\", \" like\", \" \\\"\", \" the\", \" pointing\", \" she\", \" then\"], [\" they\", \" the\", \" she\", \" he\", \" her\", \" one\", \" each\", \" many\", \" that\", \" their\"], [\" little\", \" family\", \" mom\", \" dad\", \" parents\", \" daddy\", \" girl\", \" mother\", \" plane\", \" father\"], [\"my\", \"m\", \" and\", \" said\", \"ma\", \" answered\", \" was\", \" took\", \"ie\", \" made\"], [\" and\", \" said\", \" answered\", \" told\", \" was\", \" took\", \" got\", \" showed\", \" made\", \" decided\"], [\",\", \" that\", \".\", \" to\", \" with\", \" \\\"\", \" and\", \":\", \" in\", \",\\\"\"], [\" \\\"\", \" \\ufffd\", \" \\\"'\", \" '\", \" \", \"\\n\", \" ''\", \" pointing\", \" \\\"\\\"\", \"\\ufffd\"], [\"Yes\", \"I\", \"We\", \"That\", \"What\", \"No\", \"Let\", \"You\", \"It\", \"The\"], [\" we\", \" you\", \" they\", \" there\", \" the\", \" these\", \" those\", \" planes\", \" all\", \" an\"], [\" there\", \" going\", \" all\", \" here\", \" the\", \" having\", \" ready\", \" almost\", \" getting\", \" flying\"], [\" having\", \" going\", \" discussing\", \" need\", \" getting\", \" have\", \" like\", \" talking\", \" finished\", \" just\"], [\" to\", \" on\", \" for\", \" out\", \" somewhere\", \" through\", \" home\", \" up\", \" the\", \" in\"], [\" see\", \" the\", \" take\", \" get\", \" go\", \" work\", \" fly\", \" a\", \" have\", \" look\"], [\" the\", \"?\\\"\", \" a\", \" this\", \" some\", \" and\", \" an\", \",\", \" more\", \" something\"], [\"\\n\", \" \", \" The\", \" They\", \" and\", \" Daddy\", \" So\", \" Her\", \" She\", \" And\"], [\"\\n\", \"The\", \"Mom\", \"Daddy\", \"Her\", \" \", \"So\", \"\\\"\", \" The\", \"Finally\"], [\"The\", \"Mom\", \"Her\", \"\\\"\", \"So\", \"Daddy\", \"Finally\", \"One\", \"When\", \"But\"], [\" little\", \" family\", \" mom\", \" girl\", \" daddy\", \" dad\", \" people\", \" parents\", \" father\", \" mother\"], [\" girl\", \" boy\", \" one\", \" 3\", \" child\", \" daughter\", \" family\", \" little\", \" three\", \" girls\"], [\" smiled\", \" nodded\", \" was\", \" said\", \" thought\", \" looked\", \" shook\", \"'s\", \" laughed\", \" didn\"], [\" and\", \",\", \".\", \" as\", \" at\", \" a\", \" back\", \" with\", \" excited\", \" again\"], [\" said\", \" nodded\", \" replied\", \" asked\", \" answered\", \" waved\", \" ran\", \" kept\", \" started\", \" continued\"], [\" her\", \".\", \",\", \" in\", \" yes\", \" his\", \" with\", \" their\", \" and\", \" its\"], [\" head\", \" big\", \" heads\", \" eyes\", \" biggest\", \" face\", \" mom\", \" heart\", \" daddy\", \" ears\"], [\".\", \",\", \" yes\", \" excited\", \" in\", \" eagerly\", \" with\", \" to\", \" and\", \" as\"], [\" She\", \"\\n\", \" They\", \" \\\"\", \" Then\", \" The\", \" So\", \" Daddy\", \" As\", \" \"], [\" was\", \" said\", \" knew\", \" had\", \" loved\", \" wanted\", \" felt\", \" liked\", \" looked\", \" thought\"], [\" so\", \" excited\", \" happy\", \" very\", \" ready\", \" a\", \" feeling\", \" having\", \" getting\", \" already\"], [\" to\", \" and\", \" for\", \",\", \".\", \"!\", \" because\", \" that\", \" by\", \" as\"], [\" explore\", \" see\", \" go\", \" learn\", \" be\", \" get\", \" tell\", \" have\", \" start\", \" meet\"], [\" the\", \" what\", \" all\", \" new\", \" her\", \" so\", \" a\", \" how\", \" everything\", \" where\"], [\" the\", \" of\", \" these\", \" around\", \" kinds\", \" about\", \" sorts\", \" her\", \" this\", \" those\"], [\" new\", \" modern\", \" different\", \" planes\", \" things\", \" people\", \" big\", \" fun\", \" amazing\", \" other\"], [\" and\", \".\", \",\", \" taking\", \" flying\", \" in\", \"!\", \" at\", \" take\", \" up\"], [\" so\", \" and\", \" but\", \" the\", \" planes\", \" especially\", \" which\", \" even\", \" she\", \" big\"], [\" she\", \" the\", \" even\", \" all\", \" they\", \" soon\", \" as\", \" when\", \" so\", \" asked\"], [\" was\", \" couldn\", \" wanted\", \" knew\", \" saw\", \" looked\", \" asked\", \" waved\", \" said\", \" started\"], [\" to\", \" on\", \" ready\", \" out\", \" off\", \" a\", \" up\", \" in\", \" into\", \" dressed\"], [\" see\", \" the\", \" go\", \" meet\", \" work\", \" look\", \" explore\", \" sit\", \" take\", \" know\"], [\" airport\", \" top\", \" modern\", \" very\", \" big\", \" plane\", \" beach\", \" front\", \" park\", \" edge\"], [\".\", \" with\", \" too\", \" and\", \" to\", \" first\", \",\", \" just\", \" for\", \" right\"], [\"\\n\", \" She\", \" When\", \" As\", \" The\", \" \", \" They\", \" Everywhere\", \" It\", \" Her\"], [\"\\n\", \"When\", \"<|endoftext|>\", \"The\", \"Summary\", \"At\", \" \", \"Once\", \"As\", \"They\"], [\"When\", \"The\", \"At\", \"As\", \"Once\", \"On\", \"They\", \"Suddenly\", \"Finally\", \"Soon\"], [\" family\", \" little\", \" mom\", \" airport\", \" plane\", \" dad\", \" parents\", \" girl\", \" next\", \" daddy\"], [\" girl\", \" boy\", \" plane\", \" one\", \" family\", \" girls\", \" town\", \" people\", \" planes\", \" airport\"], [\" was\", \" and\", \" saw\", \" looked\", \" watched\", \" asked\", \" had\", \" got\", \" waved\", \"'s\"], [\" up\", \" onto\", \" into\", \" on\", \" aboard\", \" the\", \" out\", \" down\", \" in\", \" off\"], [\" the\", \" onto\", \" on\", \" into\", \" a\", \" to\", \" and\", \" her\", \" high\", \" in\"], [\" stairs\", \" steps\", \" ladder\", \" plane\", \" big\", \" tall\", \" modern\", \" planes\", \" gate\", \" wall\"], [\" and\", \",\", \" to\", \".\", \" of\", \" with\", \" all\", \" as\", \" in\", \" until\"], [\" looked\", \" saw\", \" went\", \" onto\", \" sat\", \" found\", \" started\", \" waved\", \" when\", \" stepped\"], [\",\", \" to\", \" \\\"\", \" goodbye\", \" hello\", \" in\", \":\", \" with\", \",\\\"\", \" excited\"], [\" \\\"\", \" \\ufffd\", \" '\", \" \\\"'\", \"\\n\", \" \", \" ''\", \"\\ufffd\", \" and\", \" \\\"\\\"\"], [\"I\", \"Mom\", \"Wow\", \"Look\", \"This\", \"It\", \"Hello\", \"We\", \"Can\", \"Let\"], [\",\", \" done\", \" I\", \" you\", \" hello\", \" at\", \" then\", \" come\", \" we\", \" what\"], [\",\", \"!\", \"!\\\"\", \".\", \" everyone\", \" for\", \",\\\"\", \" and\", \" you\", \" at\"], [\" mom\", \" you\", \" sweet\", \" Mom\", \" daddy\", \" I\", \" my\", \" everyone\", \" Daddy\", \" little\"], [\"!\", \"!\\\"\", \" is\", \".\", \" was\", \",\", \" can\", \" in\", \" at\", \" cl\"], [\" so\", \" very\", \" going\", \" ready\", \" getting\", \" welcome\", \" safe\", \" here\", \" waiting\", \" happy\"], [\" happy\", \" proud\", \" excited\", \" big\", \" nice\", \" friendly\", \" busy\", \" full\", \" much\", \" brave\"], [\" of\", \"!\\\"\", \"!\", \" and\", \".\", \" to\", \" that\", \".\\\"\", \",\", \"\\\".\"], [\" you\", \" how\", \" me\", \" the\", \" your\", \" yourself\", \" us\", \" what\", \" my\", \" her\"], [\"!\\\"\", \" for\", \".\", \".\\\"\", \"!\", \",\", \"\\\".\", \" all\", \"!\\\".\", \" and\"], [\" You\", \" We\", \" Now\", \" Let\", \" I\", \" Come\", \" Keep\", \" That\", \" The\", \" It\"], [\" you\", \" we\", \" your\", \" the\", \" all\", \" everyone\", \" there\", \" they\", \" these\", \" I\"], [\" ready\", \" all\", \" excited\", \" going\", \" a\", \" sure\", \" getting\", \" the\", \" brave\", \" feeling\"], [\" to\", \" for\", \"?\\\"\", \",\", \"?\\\".\", \" yet\", \" and\", \" now\", \"?\", \" enough\"], [\" go\", \" explore\", \" fly\", \" have\", \" take\", \" start\", \" get\", \" learn\", \" see\", \" climb\"], [\" everyone\", \" off\", \" them\", \" me\", \" it\", \" the\", \" us\", \" your\", \" everybody\", \" her\"], [\" mom\", \" family\", \" new\", \" friends\", \" modern\", \" passport\", \" parents\", \" dad\", \" grandma\", \" daddy\"], [\"?\\\"\", \" the\", \" what\", \" your\", \" and\", \"?\\\".\", \" how\", \",\", \" around\", \" when\"], [\"\\n\", \" \", \" The\", \" And\", \" She\", \" Everyone\", \" Her\", \" Daddy\", \" Mom\", \" They\"], [\"\\n\", \"The\", \" \", \"Summary\", \"<|endoftext|>\", \"Mom\", \"So\", \"\\\"\", \"Everyone\", \"Her\"], [\"The\", \"So\", \"\\\"\", \"Her\", \"Everyone\", \"Mom\", \"And\", \"She\", \"All\", \"At\"], [\" little\", \" family\", \" girl\", \" plane\", \" mom\", \" people\", \" modern\", \" dad\", \" pilot\", \" airport\"], [\" smiled\", \" all\", \" laughed\", \" was\", \" nodded\", \" cheered\", \" said\", \" thanked\", \" looked\", \" replied\"], [\" at\", \" up\", \" around\", \" down\", \" out\", \" surprised\", \" over\", \" back\", \" so\", \" a\"], [\" and\", \" who\", \" so\", \" a\", \" but\", \" then\", \" the\", \" smiled\", \" feeling\", \" proud\"], [\" the\", \" they\", \" sure\", \" then\", \" smiled\", \" after\", \" everyone\", \" soon\", \" nodded\", \" finally\"], [\" all\", \" were\", \" saw\", \" said\", \" nodded\", \" smiled\", \" had\", \" couldn\", \" looked\", \" watched\"], [\" said\", \" smiled\", \" cheered\", \" nodded\", \" agreed\", \" cl\", \" hugged\", \" laughed\", \" shouted\", \" had\"], [\" and\", \" each\", \" the\", \".\", \",\", \" their\", \" her\", \" mom\", \" for\", \" one\"], [\" other\", \"other\", \" one\", \" others\", \" of\", \" family\", \" and\", \" so\", \" plane\", \" mom\"], [\".\", \" tightly\", \" and\", \" tight\", \",\", \" with\", \" again\", \" in\", \" when\", \" before\"], [\" and\", \" happy\", \" then\", \" grateful\", \" thankful\", \" feeling\", \" proud\", \" excited\", \" smiled\", \" \\\"\"], [\" then\", \" the\", \" they\", \" said\", \" thanked\", \" from\", \" smiled\", \" were\", \" everyone\", \" had\"], [\" all\", \" were\", \" said\", \" smiled\", \" knew\", \" laughed\", \" had\", \" thanked\", \" enjoyed\", \" started\"], [\" that\", \" to\", \".\", \" it\", \" with\", \"!\", \" they\", \" the\", \" on\", \",\"], [\"\\n\", \" They\", \" The\", \" From\", \" \", \" Everyone\", \" It\", \" Then\", \" And\", \" Together\"], [\" little\", \" family\", \" girl\", \" modern\", \" mom\", \" plane\", \" airport\", \" dad\", \" daughter\", \" three\"], [\" girl\", \" one\", \" boy\", \" 3\", \" plane\", \" little\", \" family\", \" daughter\", \" three\", \" town\"], [\" was\", \" smiled\", \" felt\", \" had\", \" said\", \"'s\", \" knew\", \" hugged\", \" couldn\", \" ran\"], [\" so\", \" very\", \" the\", \" happy\", \" a\", \" over\", \" thrilled\", \" excited\", \" really\", \" now\"], [\" happy\", \" excited\", \" proud\", \" glad\", \" brave\", \" thrilled\", \" pleased\", \" grateful\", \" eager\", \" thankful\"], [\" that\", \" to\", \" and\", \" she\", \",\", \".\", \"!\", \" with\", \" they\", \" because\"], [\" she\", \" and\", \" they\", \" that\", \" it\", \" because\", \" the\", \" but\", \" with\", \" knowing\"], [\" couldn\", \" had\", \" jumped\", \" was\", \" smiled\", \" ran\", \" felt\", \" hugged\", \" knew\", \" started\"], [\" to\", \" everyone\", \" the\", \" her\", \" them\", \" more\", \" a\", \" all\", \" everything\", \" one\"], [\" explore\", \" go\", \" show\", \" be\", \" stay\", \" keep\", \" make\", \" come\", \" take\", \" share\"], [\" everyone\", \" her\", \" them\", \" the\", \" off\", \" all\", \" that\", \" it\", \" how\", \" their\"], [\" family\", \" mom\", \" appreciation\", \" new\", \" parents\", \" modern\", \" friends\", \" dad\", \" daddy\", \" passport\"], [\" modern\", \" family\", \" toy\", \" friends\", \" friend\", \" passport\", \" plane\", \" adventure\", \" world\", \" things\"], [\" to\", \".\", \" plane\", \" planes\", \" and\", \" at\", \",\", \" car\", \" on\", \" with\"], [\"\\n\", \" She\", \" \", \" The\", \" So\", \" Everyone\", \" From\", \" They\", \" And\", \" Her\"], [\"\\n\", \"Summary\", \"<|endoftext|>\", \"The\", \" \", \"So\", \"She\", \"And\", \"They\", \"From\"], [\"\\n\", \",\", \".\", \"\\n\\n\", \" \\\"\", \".\\\"\", \" The\", \" \", \":\", \" the\"], [\" She\", \" The\", \" Her\", \"\\n\", \" So\", \" Everyone\", \"\\ufffd\", \" It\", \" He\", \" But\"], [\" little\", \" end\", \" moral\", \" End\", \" next\", \" girl\", \" old\", \" sun\", \" man\", \" bear\"], [\" was\", \" took\", \" laughed\", \" went\", \" all\", \" smiled\", \" got\", \" had\", \" watched\", \" arrived\"], [\" to\", \" home\", \" on\", \" outside\", \" back\", \" for\", \" into\", \" out\", \" shopping\", \" inside\"], [\" a\", \" to\", \" the\", \",\", \" their\", \" and\", \" lots\", \" an\", \" many\", \" another\"], [\" the\", \" have\", \" play\", \" explore\", \" find\", \" get\", \" buy\", \" see\", \" enjoy\", \" celebrate\"], [\" beach\", \" store\", \" park\", \" playground\", \" kitchen\", \" big\", \" car\", \" toy\", \" pond\", \" pool\"], [\" store\", \" mall\", \" house\", \" park\", \" playground\", \" restaurant\", \" hospital\", \" shop\", \" market\", \" bathroom\"], [\" and\", \",\", \".\", \" with\", \" to\", \" every\", \" for\", \" had\", \" where\", \" was\"], [\" show\", \" go\", \" see\", \" visit\", \" pick\", \" get\", \" take\", \" meet\", \" celebrate\", \" watch\"], [\" and\", \" back\", \" down\", \" with\", \" on\", \" out\", \" along\", \" up\", \".\", \" together\"], [\" the\", \" their\", \" time\", \",\", \" a\", \" and\", \" board\", \".\", \" big\", \" vacation\"], [\" special\", \" big\", \" fun\", \" trips\", \" new\", \" adventures\", \" trip\", \" amazing\", \" exciting\", \" wonderful\"], [\".\", \",\", \" and\", \" with\", \" too\", \"!\", \" to\", \" every\", \" together\", \" again\"], [\"\\n\", \" They\", \" When\", \" The\", \" Everyone\", \" As\", \" Everywhere\", \" \", \" One\", \" After\"], [\" time\", \" day\", \" morning\", \" week\", \" now\", \" year\", \" person\", \" member\", \" night\", \" step\"], [\" they\", \" the\", \" someone\", \" she\", \",\", \" a\", \" one\", \" he\", \" there\", \" it\"], [\" of\", \" member\", \" day\", \",\", \" little\", \" went\", \" person\", \" was\", \" arrived\", \" item\"], [\" the\", \" them\", \" these\", \" her\", \" those\", \" their\", \" a\", \" it\", \" his\", \" this\"], [\" planes\", \" family\", \" people\", \" luggage\", \" little\", \" flights\", \" new\", \" kids\", \" children\", \" passengers\"], [\" girl\", \" boy\", \" kids\", \" children\", \" girls\", \" people\", \" bird\", \" birds\", \" boys\", \" ones\"], [\"'s\", \",\", \" went\", \" looked\", \" would\", \" was\", \"\\ufffd\", \" and\", \" got\", \" goes\"], [\" the\", \" she\", \" her\", \" they\", \" Lucy\", \" Sarah\", \" Amy\", \" Jane\", \" it\", \" a\"], [\" dad\", \" mom\", \" daddy\", \" parents\", \" family\", \" father\", \" mum\", \" mother\", \" eyes\", \" Dad\"], [\" would\", \",\", \" and\", \" was\", \" went\", \" got\", \" came\", \" smiled\", \" waved\", \" looked\"], [\" would\", \" they\", \" the\", \" she\", \" went\", \" was\", \" her\", \" and\", \" asked\", \" asks\"], [\" always\", \" pick\", \" take\", \" go\", \" watch\", \" wave\", \" come\", \" bring\", \" look\", \" give\"], [\" around\", \" up\", \" and\", \" to\", \" with\", \" together\", \" in\", \" back\", \" on\", \" through\"], [\" her\", \" them\", \" the\", \" joy\", \" excitement\", \" their\", \" a\", \" it\", \" tw\", \" all\"], [\" and\", \".\", \",\", \" to\", \"!\", \" for\", \" in\", \" as\", \" until\", \" on\"], [\" joy\", \" excitement\", \" love\", \" anticipation\", \" cheer\", \" help\", \" make\", \" enthusiasm\", \" smile\", \" have\"], [\" for\", \".\", \" to\", \" that\", \",\", \" she\", \" of\", \"!\", \" they\", \" wherever\"], [\" the\", \" their\", \" her\", \" more\", \" new\", \" something\", \" a\", \" all\", \" many\", \" everyone\"], [\" best\", \" next\", \" new\", \" family\", \" future\", \" fun\", \" surprise\", \" special\", \" airport\", \" rest\"], [\".\", \" to\", \" soon\", \" would\", \" and\", \" of\", \"'s\", \" again\", \" one\", \" on\"], [\"\\n\", \" \", \" She\", \" The\", \" They\", \" And\", \" It\", \" Finally\", \" Her\", \" Everyone\"], [\"Summary\", \"\\n\", \"Story\", \"Words\", \"The\", \"Features\", \"Finally\", \" \", \"One\", \"Eventually\"], [\"The\", \"Finally\", \"One\", \"At\", \"Eventually\", \"When\", \"Soon\", \"And\", \"After\", \"They\"], [\",\", \" the\", \" they\", \" she\", \" her\", \" it\", \" one\", \" their\", \" enough\", \" after\"], [\" the\", \" they\", \" it\", \" their\", \" her\", \" after\", \" she\", \" one\", \" when\", \" a\"], [\" little\", \" family\", \" plane\", \" airport\", \" day\", \" march\", \" girl\", \" dad\", \" whole\", \" mom\"], [\" girl\", \" boy\", \" one\", \" plane\", \" bird\", \" brother\", \" duck\", \" march\", \" ant\", \" family\"], [\" had\", \"'s\", \" got\", \" grew\", \" and\", \" was\", \" found\", \" said\", \" became\", \" started\"], [\" to\", \" her\", \" marching\", \" their\", \" the\", \" singing\", \" his\", \" walking\", \" exploring\", \" looking\"], [\" feel\", \" understand\", \" get\", \" march\", \" miss\", \" grow\", \" realize\", \" sing\", \" enjoy\", \" learn\"], [\" tired\", \" ready\", \" a\", \" excited\", \" bored\", \" off\", \" sleepy\", \" hungry\", \" better\", \" very\"], [\".\", \" and\", \",\", \" from\", \" so\", \" again\", \" too\", \"!\", \" in\", \" at\"], [\".\", \",\", \" and\", \"!\", \" so\", \" -\", \" in\", \" but\", \" because\", \" when\"], [\" She\", \" Her\", \" But\", \" The\", \" They\", \" This\", \" As\", \" So\", \"\\n\", \" When\"], [\" asked\", \" was\", \" looked\", \" wanted\", \" knew\", \" said\", \" saw\", \" had\", \" grabbed\", \" went\"], [\" so\", \" filled\", \" very\", \" getting\", \" a\", \" sad\", \" determined\", \" happy\", \" looking\", \" feeling\"], [\" excited\", \" thirsty\", \" happy\", \" proud\", \" tired\", \" eager\", \" full\", \" glad\", \" close\", \" grateful\"]], \"correctTokenRank\": [294, 0, 0, 0, 0, 0, 0, 0, 73, 0, 0, 0, 13, 0, 28, 95, 0, 1, 22, 2, 0, 1, 0, 0, 24, 0, 2, 0, 1, 38, 0, 0, 0, 0, 1, 11, 0, 0, 0, 9, 4, 5, 0, 0, 2, 2, 1, 2, 0, 15, 0, 0, 20, 0, 15, 1, 0, 13, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 2, 0, 3, 2, 1, 0, 16, 0, 1, 0, 0, 0, 0, 1, 1, 0, 42, 0, 0, 0, 0, 16, 0, 0, 49, 1, 0, 7, 2, 0, 1, 0, 0, 4, 23, 0, 0, 0, 13, 7, 1, 0, 0, 0, 0, 1, 8, 42, 0, 1, 0, 6, 1, 0, 4, 0, 2, 13, 5, 2, 0, 0, 0, 0, 0, 4, 0, 10, 0, 2, 1, 3, 2, 1, 0, 2, 2, 20, 10, 3, 2, 1, 0, 97, 39, 4, 23, 4, 73, 3, 0, 19, 0, 6, 0, 0, 4, 0, 1, 2, 4, 1, 0, 152, 4, 25, 0, 42, 0, 0, 3, 0, 0, 1, 4, 0, 0, 0, 0, 12, 0, 2, 79, 5, 0, 0, 1, 0, 4], \"correctTokenLogProb\": [-11.625454902648926, -0.372489869594574, -0.005904730875045061, -0.011737688444554806, -0.1878184974193573, -0.10250605642795563, -0.1382717341184616, -0.0321299247443676, -5.991303443908691, -0.5746181607246399, -0.4282326400279999, -0.30328744649887085, -5.65919828414917, -2.4236323833465576, -4.623875617980957, -6.590052604675293, -0.042138077318668365, -1.9959081411361694, -5.6335530281066895, -1.905128002166748, -0.9895710945129395, -2.5058727264404297, -0.35668182373046875, -0.7548707127571106, -4.735068321228027, -0.5946731567382812, -2.1486759185791016, -0.09278443455696106, -1.3276169300079346, -7.621414661407471, -0.06834664195775986, -0.889788031578064, -0.0339607410132885, -0.9458173513412476, -1.6791880130767822, -3.8572938442230225, -0.6308883428573608, -1.1232950687408447, -0.29644775390625, -5.13623046875, -5.095574378967285, -3.26711368560791, -0.007819517515599728, -0.18164236843585968, -2.637809991836548, -2.3811194896698, -1.3665785789489746, -1.8806651830673218, -0.003452533157542348, -5.287355422973633, -0.43409377336502075, -0.19858379662036896, -4.841526985168457, -0.5131959915161133, -4.908519744873047, -1.3566906452178955, -0.4600834846496582, -4.830392837524414, -2.2664647102355957, -0.24191200733184814, -0.050112128257751465, -0.2582389712333679, -0.2726271152496338, -0.0036967527121305466, -1.3802396059036255, -0.0841309204697609, -2.2264645099639893, -0.8177525997161865, -0.0016505202511325479, -0.2397218495607376, -0.6482017636299133, -0.9068384170532227, -2.4962146282196045, -0.2099856436252594, -1.7463961839675903, -2.0468666553497314, -0.11845239996910095, -2.431112766265869, -1.9535925388336182, -1.379508137702942, -0.8167728185653687, -4.419667720794678, -0.7300421595573425, -2.6654233932495117, -0.17896340787410736, -0.18972711265087128, -0.7267122864723206, -0.020010385662317276, -1.7059850692749023, -1.5423017740249634, -0.00046659549116156995, -6.36149787902832, -1.2832740545272827, -1.2193286418914795, -0.8976349234580994, -0.5017741918563843, -4.201582431793213, -0.109302818775177, -0.03258315473794937, -6.803769588470459, -1.8130090236663818, -0.719337522983551, -3.136812686920166, -2.519213914871216, -2.1019766330718994, -1.9251044988632202, -0.09708235412836075, -0.3409811854362488, -2.387692928314209, -5.753907203674316, -0.04679347202181816, -0.31035345792770386, -0.26277634501457214, -4.80776834487915, -3.099146604537964, -1.2577797174453735, -0.23072941601276398, -0.1686599999666214, -0.03456159308552742, -0.041516028344631195, -2.4174718856811523, -4.463851451873779, -7.481378078460693, -0.6082432866096497, -2.094578504562378, -0.6794657707214355, -3.3744683265686035, -1.3131829500198364, -0.0037339269183576107, -2.4467132091522217, -1.1844749450683594, -2.7604708671569824, -5.0592942237854, -4.280617713928223, -1.8505373001098633, -0.14208939671516418, -0.0002898749662563205, -0.5908623933792114, -0.17303341627120972, -0.4503314197063446, -2.1354479789733887, -0.47746407985687256, -3.7926478385925293, -0.02097589708864689, -2.1107640266418457, -1.169552206993103, -2.8503684997558594, -2.9834933280944824, -1.5757774114608765, -0.5248868465423584, -3.215791702270508, -16.646602630615234, -5.324350357055664, -4.359821796417236, -3.124971628189087, -2.7403366565704346, -1.8093516826629639, -0.9515280723571777, -6.534619331359863, -5.37264347076416, -3.521487236022949, -4.890478134155273, -2.865650177001953, -8.220545768737793, -2.573607921600342, -0.8533508777618408, -5.152226448059082, -0.3188457787036896, -4.101963520050049, -0.7935125827789307, -0.32164621353149414, -3.5270328521728516, -0.5608841180801392, -2.829089641571045, -2.155390977859497, -3.445237159729004, -1.54229736328125, -0.7731188535690308, -7.224249362945557, -2.695648193359375, -6.626819610595703, -0.9364675283432007, -5.718138694763184, -0.8462379574775696, -1.065568208694458, -3.258896589279175, -1.3465521335601807, -0.20639769732952118, -1.9113335609436035, -3.625382423400879, -0.052747633308172226, -0.20388591289520264, -1.143463134765625, -0.013952808454632759, -4.204424858093262, -0.14188191294670105, -2.596562623977661, -7.299529075622559, -4.704736232757568, -0.23656286299228668, -0.42535340785980225, -2.492424964904785, -1.115234136581421, -2.6541266441345215]}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x77ec1eb34310>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_prompt = model.generate(\n",
    "    \"Once upon a time\",\n",
    "    stop_at_eos=False,  # avoids a bug on MPS\n",
    "    temperature=1,\n",
    "    verbose=True,\n",
    "    max_new_tokens=200,\n",
    ")\n",
    "logits, cache = model.run_with_cache(example_prompt)\n",
    "cv.logits.token_log_probs(\n",
    "    model.to_tokens(example_prompt),\n",
    "    model(example_prompt)[0].log_softmax(dim=-1),\n",
    "    model.to_string,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "er3H1TDoOVHw"
   },
   "source": [
    "# Training an SAE\n",
    "\n",
    "Now we're ready to train out SAE. We'll make a runner config, instantiate the runner and the rest is taken care of for us!\n",
    "\n",
    "During training, you use weights and biases to check key metrics which indicate how well we are able to optimize the variables we care about.\n",
    "\n",
    "To get a better sense of which variables to look at, you can read my (Joseph's) post [here](https://www.lesswrong.com/posts/f9EgfLSurAiqRJySD/open-source-sparse-autoencoders-for-all-residual-stream) and especially look at my weights and biases report [here](https://links-cdn.wandb.ai/wandb-public-images/links/jbloom/uue9i416.html).\n",
    "\n",
    "A few tips:\n",
    "- Feel free to reorganize your wandb dashboard to put L0, CE_Loss_score, explained variance and other key metrics in one section at the top.\n",
    "- Make a [run comparer](https://docs.wandb.ai/guides/app/features/panels/run-comparer) when tuning hyperparameters.\n",
    "- You can download the resulting sparse autoencoder / sparsity estimate from wandb and upload them to huggingface if you want to share your SAE with other.\n",
    "    - cfg.json (training config)\n",
    "    - sae_weight.safetensors (model weights)\n",
    "    - sparsity.safetensors (sparsity estimate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jCHtPycOOVHw"
   },
   "source": [
    "## MLP Out\n",
    "\n",
    "I've tuned the hyperparameters below for a decent SAE which achieves 86% CE Loss recovered and an L0 of ~85, and runs in about 2 hours on an M3 Max. You can get an SAE that looks better faster if you only consider L0 and CE loss but it will likely have more dense features and more dead features. Here's a link to my output with two runs with two different L1's: https://wandb.ai/jbloom/sae_lens_tutorial ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "oAsZCAdJOVHw"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run name: 16384-L1-5-LR-5e-05-Tokens-1.229e+08\n",
      "n_tokens_per_buffer (millions): 0.262144\n",
      "Lower bound: n_contexts_per_buffer (millions): 0.001024\n",
      "Total training steps: 30000\n",
      "Total wandb updates: 1000\n",
      "n_tokens_per_feature_sampling_window (millions): 1048.576\n",
      "n_tokens_per_dead_feature_window (millions): 1048.576\n",
      "We will reset the sparsity calculation 30 times.\n",
      "Number tokens in sparsity calculation window: 4.10e+06\n",
      "Loaded pretrained model tiny-stories-1L-21M into HookedTransformer\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading readme: 100%|██████████| 415/415 [00:00<00:00, 4.30MB/s]\n",
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcurt-tigges\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.17.1"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/curttigges/projects/SAELens/tutorials/wandb/run-20240610_114538-yr1gvjdc</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc' target=\"_blank\">16384-L1-5-LR-5e-05-Tokens-1.229e+08</a></strong> to <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Estimating norm scaling factor: 100%|██████████| 1000/1000 [00:33<00:00, 29.79it/s]\n",
      "30000| MSE Loss 187.703 | L1 156.883:   1%|          | 1228800/122880000 [42:59<70:56:28, 476.34it/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>details/current_l1_coefficient</td><td>▁▅██████████████████████████████████████</td></tr><tr><td>details/current_learning_rate</td><td>████████████████████████████████▇▇▅▅▄▃▂▁</td></tr><tr><td>details/n_training_tokens</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███</td></tr><tr><td>losses/ghost_grad_loss</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/l1_loss</td><td>█▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/mse_loss</td><td>▄▅█▇▆▆▅▄▄▄▄▃▃▃▃▃▃▃▂▃▂▂▂▂▂▂▂▂▁▁▂▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/overall_loss</td><td>▁▅█▇▇▆▆▆▆▆▅▅▅▅▅▅▅▅▅▅▅▄▅▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄</td></tr><tr><td>metrics/CE_loss_score</td><td>█▂▁▂▃▃▄▄▄▄▄▄▄▄▅▅▅▅▅▅▅▅▅▆▅▆▅▅▅▅▆▅▆▆▆▆▆▆▆▅</td></tr><tr><td>metrics/ce_loss_with_ablation</td><td>█▆▇▃▄▄▃▆▆▅▃▅▅▅▅▃▄▅▅▆▄▅▃▅▃▄▅▁▅▅▆▅▆▄▆▄▆▄▅▄</td></tr><tr><td>metrics/ce_loss_with_sae</td><td>▂▇█▆▅▅▅▅▄▄▅▆▄▄▃▄▄▄▁▃▃▃▁▄▄▂▃▂▄▂▁▃▄▃▄▂▄▂▃▄</td></tr><tr><td>metrics/ce_loss_without_sae</td><td>▇▇▇▅▄▄▆▆▅▄▇█▆▆▄▅▆▇▁▄▄▅▁▇▆▄▅▄▆▄▂▅▇▆▇▄▇▄▅█</td></tr><tr><td>metrics/explained_variance</td><td>▅▄▁▂▃▄▄▅▅▅▅▆▆▆▆▆▆▆▆▇▆▇▇▇▇▇▇▇▇█▇▇▇███████</td></tr><tr><td>metrics/explained_variance_std</td><td>▁▅███▇▇▇▇▇▇▆▆▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▇▆▆▆▆▆▆▆▅▆</td></tr><tr><td>metrics/l0</td><td>█▃▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>metrics/l2_norm</td><td>▆▅▁▃▃▂▂▃▁▅▃▃▅▆▂█▄▃▄▆▇▃▆▃▃▁▂▅▅▆▇▅▄▅▃██▃▃▃</td></tr><tr><td>metrics/l2_norm_in</td><td>▅▆▅▆▆▄▅▅▃▆▄▄▆▇▄█▄▃▅▆▇▃▆▃▃▁▂▅▅▆▇▄▅▆▃▇█▄▄▃</td></tr><tr><td>metrics/l2_ratio</td><td>▇▅▁▃▃▃▃▄▃▆▅▅▆▆▄▇▆▆▅▆▇▅▇▆▅▅▅▆▆▇▇▇▆▆▆██▆▆▆</td></tr><tr><td>metrics/mean_log10_feature_sparsity</td><td>█▄▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>sparsity/below_1e-5</td><td>▁▁▁▅▆███▆▆▆▆▆▆▆▆▅▅▅▅▅▅▅▅▅▅▅▅▅▅</td></tr><tr><td>sparsity/below_1e-6</td><td>▁▁▁▁▁▃▆▆▆█▆█████▆▆▆▆▆▆▆▆▆▆▆▆▆▆</td></tr><tr><td>sparsity/dead_features</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▅▅▅▁▁▅█▁▁▁▅▅█▅█▅▅████▅▅▅████</td></tr><tr><td>sparsity/mean_passes_since_fired</td><td>▂▁▁▁▁▁▁▁▂▂▁▁▂▂▂▂▂▂▄▂▂▂▃▃▄▄▄▃▃▃▄▅▆▄▄▅▆▇▇█</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>details/current_l1_coefficient</td><td>5</td></tr><tr><td>details/current_learning_rate</td><td>0.0</td></tr><tr><td>details/n_training_tokens</td><td>122880000</td></tr><tr><td>losses/ghost_grad_loss</td><td>0.0</td></tr><tr><td>losses/l1_loss</td><td>31.37665</td></tr><tr><td>losses/mse_loss</td><td>187.70346</td></tr><tr><td>losses/overall_loss</td><td>344.5867</td></tr><tr><td>metrics/CE_loss_score</td><td>0.90369</td></tr><tr><td>metrics/ce_loss_with_ablation</td><td>8.30545</td></tr><tr><td>metrics/ce_loss_with_sae</td><td>2.62867</td></tr><tr><td>metrics/ce_loss_without_sae</td><td>2.02306</td></tr><tr><td>metrics/explained_variance</td><td>0.66377</td></tr><tr><td>metrics/explained_variance_std</td><td>0.13242</td></tr><tr><td>metrics/l0</td><td>192.95703</td></tr><tr><td>metrics/l2_norm</td><td>24.64933</td></tr><tr><td>metrics/l2_norm_in</td><td>31.38967</td></tr><tr><td>metrics/l2_ratio</td><td>0.77361</td></tr><tr><td>metrics/mean_log10_feature_sparsity</td><td>-2.66941</td></tr><tr><td>sparsity/below_1e-5</td><td>2</td></tr><tr><td>sparsity/below_1e-6</td><td>2</td></tr><tr><td>sparsity/dead_features</td><td>2</td></tr><tr><td>sparsity/mean_passes_since_fired</td><td>0.86823</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">16384-L1-5-LR-5e-05-Tokens-1.229e+08</strong> at: <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc</a><br/> View project at: <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial</a><br/>Synced 6 W&B file(s), 0 media file(s), 3 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20240610_114538-yr1gvjdc/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "total_training_steps = 30_000  # probably we should do more\n",
    "batch_size = 4096\n",
    "total_training_tokens = total_training_steps * batch_size\n",
    "\n",
    "lr_warm_up_steps = 0\n",
    "lr_decay_steps = total_training_steps // 5  # 20% of training\n",
    "l1_warm_up_steps = total_training_steps // 20  # 5% of training\n",
    "\n",
    "cfg = LanguageModelSAERunnerConfig(\n",
    "    # Data Generating Function (Model + Training Distibuion)\n",
    "    model_name=\"tiny-stories-1L-21M\",  # our model (more options here: https://neelnanda-io.github.io/TransformerLens/generated/model_properties_table.html)\n",
    "    hook_name=\"blocks.0.hook_mlp_out\",  # A valid hook point (see more details here: https://neelnanda-io.github.io/TransformerLens/generated/demos/Main_Demo.html#Hook-Points)\n",
    "    hook_layer=0,  # Only one layer in the model.\n",
    "    d_in=1024,  # the width of the mlp output.\n",
    "    dataset_path=\"apollo-research/roneneldan-TinyStories-tokenizer-gpt2\",  # this is a tokenized language dataset on Huggingface for the Tiny Stories corpus.\n",
    "    is_dataset_tokenized=True,\n",
    "    streaming=True,  # we could pre-download the token dataset if it was small.\n",
    "    # SAE Parameters\n",
    "    mse_loss_normalization=None,  # We won't normalize the mse loss,\n",
    "    expansion_factor=16,  # the width of the SAE. Larger will result in better stats but slower training.\n",
    "    b_dec_init_method=\"zeros\",  # The geometric median can be used to initialize the decoder weights.\n",
    "    apply_b_dec_to_input=False,  # We won't apply the decoder weights to the input.\n",
    "    normalize_sae_decoder=False,\n",
    "    scale_sparsity_penalty_by_decoder_norm=True,\n",
    "    decoder_heuristic_init=True,\n",
    "    init_encoder_as_decoder_transpose=True,\n",
    "    normalize_activations=\"expected_average_only_in\",\n",
    "    # Training Parameters\n",
    "    lr=5e-5,  # lower the better, we'll go fairly high to speed up the tutorial.\n",
    "    adam_beta1=0.9,  # adam params (default, but once upon a time we experimented with these.)\n",
    "    adam_beta2=0.999,\n",
    "    lr_scheduler_name=\"constant\",  # constant learning rate with warmup. Could be better schedules out there.\n",
    "    lr_warm_up_steps=lr_warm_up_steps,  # this can help avoid too many dead features initially.\n",
    "    lr_decay_steps=lr_decay_steps,  # this will help us avoid overfitting.\n",
    "    l1_coefficient=5,  # will control how sparse the feature activations are\n",
    "    l1_warm_up_steps=l1_warm_up_steps,  # this can help avoid too many dead features initially.\n",
    "    lp_norm=1.0,  # the L1 penalty (and not a Lp for p < 1)\n",
    "    train_batch_size_tokens=batch_size,\n",
    "    context_size=512,  # will control the lenght of the prompts we feed to the model. Larger is better but slower. so for the tutorial we'll use a short one.\n",
    "    # Activation Store Parameters\n",
    "    n_batches_in_buffer=64,  # controls how many activations we store / shuffle.\n",
    "    training_tokens=total_training_tokens,  # 100 million tokens is quite a few, but we want to see good stats. Get a coffee, come back.\n",
    "    store_batch_size_prompts=16,\n",
    "    # Resampling protocol\n",
    "    use_ghost_grads=False,  # we don't use ghost grads anymore.\n",
    "    feature_sampling_window=1000,  # this controls our reporting of feature sparsity stats\n",
    "    dead_feature_window=1000,  # would effect resampling or ghost grads if we were using it.\n",
    "    dead_feature_threshold=1e-4,  # would effect resampling or ghost grads if we were using it.\n",
    "    # WANDB\n",
    "    log_to_wandb=True,  # always use wandb unless you are just testing code.\n",
    "    wandb_project=\"sae_lens_tutorial\",\n",
    "    wandb_log_frequency=30,\n",
    "    eval_every_n_wandb_logs=20,\n",
    "    # Misc\n",
    "    device=device,\n",
    "    seed=42,\n",
    "    n_checkpoints=0,\n",
    "    checkpoint_path=\"checkpoints\",\n",
    "    dtype=\"float32\",\n",
    ")\n",
    "# look at the next cell to see some instruction for what to do while this is running.\n",
    "sparse_autoencoder = SAETrainingRunner(cfg).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "khR_QkAJOVHw"
   },
   "source": [
    "# TO DO: Understanding TinyStories-1L with our SAE\n",
    "\n",
    "I haven't had time yet to complete this section, but I'd love to see a PR where someones uses an SAE they trained in this tutorial to understand this model better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b4sUumxZOVHw"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Let's start by getting the top 10 logits for each feature\n",
    "projection_onto_unembed = sparse_autoencoder.W_dec @ model.W_U\n",
    "\n",
    "\n",
    "# get the top 10 logits.\n",
    "vals, inds = torch.topk(projection_onto_unembed, 10, dim=1)\n",
    "\n",
    "# get 10 random features\n",
    "random_indices = torch.randint(0, projection_onto_unembed.shape[0], (10,))\n",
    "\n",
    "# Show the top 10 logits promoted by those features\n",
    "top_10_logits_df = pd.DataFrame(\n",
    "    [model.to_str_tokens(i) for i in inds[random_indices]],\n",
    "    index=random_indices.tolist(),\n",
    ").T\n",
    "top_10_logits_df"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}